mirror of
https://github.com/whit3rabbit/anyllm-proxy.git
synced 2026-09-21 16:00:49 +00:00
Merge feat/batch-engine-phase1 into main
Resolves conflicts: take HEAD (security audit) for mcp.rs imports, register_server_blocking error handling, and maybe_execute_tools loop. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Generated
+24
@@ -35,6 +35,28 @@ version = "1.0.102"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c"
|
||||
|
||||
[[package]]
|
||||
name = "anyllm_batch_engine"
|
||||
version = "0.2.0"
|
||||
dependencies = [
|
||||
"anyllm_translate",
|
||||
"async-trait",
|
||||
"hex",
|
||||
"hmac",
|
||||
"pretty_assertions",
|
||||
"reqwest",
|
||||
"rusqlite",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"sha2",
|
||||
"thiserror 2.0.18",
|
||||
"tokio",
|
||||
"tokio-util",
|
||||
"tracing",
|
||||
"url",
|
||||
"uuid",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "anyllm_client"
|
||||
version = "0.2.0"
|
||||
@@ -57,6 +79,7 @@ dependencies = [
|
||||
name = "anyllm_proxy"
|
||||
version = "0.2.0"
|
||||
dependencies = [
|
||||
"anyllm_batch_engine",
|
||||
"anyllm_client",
|
||||
"anyllm_translate",
|
||||
"aws-credential-types",
|
||||
@@ -2896,6 +2919,7 @@ dependencies = [
|
||||
"bytes",
|
||||
"futures-core",
|
||||
"futures-sink",
|
||||
"futures-util",
|
||||
"pin-project-lite",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
+1
-1
@@ -1,5 +1,5 @@
|
||||
[workspace]
|
||||
members = ["crates/translator", "crates/client", "crates/proxy"]
|
||||
members = ["crates/translator", "crates/client", "crates/batch_engine", "crates/proxy"]
|
||||
resolver = "2"
|
||||
|
||||
[workspace.package]
|
||||
|
||||
@@ -0,0 +1,28 @@
|
||||
[package]
|
||||
name = "anyllm_batch_engine"
|
||||
description = "Batch orchestration engine with job queue, workers, and event-driven notifications"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
anyllm_translate = { path = "../translator", version = "0.2.0" }
|
||||
rusqlite = { version = "0.32", features = ["bundled"] }
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json = "1"
|
||||
tokio = { version = "1", features = ["rt", "sync", "time", "macros"] }
|
||||
async-trait = "0.1"
|
||||
thiserror = "2"
|
||||
uuid = { version = "1", features = ["v4"] }
|
||||
tracing = "0.1"
|
||||
reqwest = { version = "0.12", default-features = false, features = ["json", "native-tls"] }
|
||||
url = "2"
|
||||
hmac = "0.12"
|
||||
sha2 = "0.10"
|
||||
tokio-util = { version = "0.7", features = ["rt"] }
|
||||
hex = "0.4"
|
||||
|
||||
[dev-dependencies]
|
||||
tokio = { version = "1", features = ["full"] }
|
||||
pretty_assertions = "1"
|
||||
@@ -0,0 +1,249 @@
|
||||
// crates/batch_engine/src/db.rs
|
||||
//! SQLite schema initialization for batch_engine tables.
|
||||
|
||||
use rusqlite::Connection;
|
||||
|
||||
/// ISO 8601 timestamp for "now" in UTC.
|
||||
pub fn now_iso8601() -> String {
|
||||
// Replicates the pattern used in proxy's admin/db.rs.
|
||||
// Using SystemTime to avoid chrono dependency.
|
||||
let now = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_secs();
|
||||
// Convert epoch seconds to ISO 8601. Simplified: just store epoch
|
||||
// and use SQLite's datetime() for display. But for compatibility
|
||||
// with existing code, produce a formatted string.
|
||||
let secs = now;
|
||||
let days = secs / 86400;
|
||||
let day_secs = secs % 86400;
|
||||
let h = day_secs / 3600;
|
||||
let m = (day_secs % 3600) / 60;
|
||||
let s = day_secs % 60;
|
||||
|
||||
// Civil date from days since epoch (Howard Hinnant algorithm).
|
||||
let z = days as i64 + 719468;
|
||||
let era = if z >= 0 { z } else { z - 146096 } / 146097;
|
||||
let doe = (z - era * 146097) as u64;
|
||||
let yoe = (doe - doe / 1460 + doe / 36524 - doe / 146096) / 365;
|
||||
let y = yoe as i64 + era * 400;
|
||||
let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
|
||||
let mp = (5 * doy + 2) / 153;
|
||||
let d = doy - (153 * mp + 2) / 5 + 1;
|
||||
let m_val = if mp < 10 { mp + 3 } else { mp - 9 };
|
||||
let y_val = if m_val <= 2 { y + 1 } else { y };
|
||||
|
||||
format!(
|
||||
"{:04}-{:02}-{:02}T{:02}:{:02}:{:02}Z",
|
||||
y_val, m_val, d, h, m, s
|
||||
)
|
||||
}
|
||||
|
||||
/// Initialize all batch_engine tables.
|
||||
pub fn init_batch_engine_tables(conn: &Connection) -> rusqlite::Result<()> {
|
||||
conn.execute_batch(
|
||||
"
|
||||
CREATE TABLE IF NOT EXISTS batch_job (
|
||||
batch_id TEXT PRIMARY KEY,
|
||||
status TEXT NOT NULL DEFAULT 'queued',
|
||||
execution_mode TEXT NOT NULL,
|
||||
provider TEXT,
|
||||
provider_batch_id TEXT,
|
||||
priority INTEGER NOT NULL DEFAULT 0,
|
||||
key_id INTEGER,
|
||||
input_file_id TEXT NOT NULL,
|
||||
webhook_url TEXT,
|
||||
metadata TEXT,
|
||||
total INTEGER NOT NULL DEFAULT 0,
|
||||
processing INTEGER NOT NULL DEFAULT 0,
|
||||
succeeded INTEGER NOT NULL DEFAULT 0,
|
||||
failed INTEGER NOT NULL DEFAULT 0,
|
||||
cancelled INTEGER NOT NULL DEFAULT 0,
|
||||
expired INTEGER NOT NULL DEFAULT 0,
|
||||
created_at TEXT NOT NULL,
|
||||
started_at TEXT,
|
||||
completed_at TEXT,
|
||||
expires_at TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_batch_job_dequeue
|
||||
ON batch_job(status, priority DESC, created_at ASC);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_batch_job_key
|
||||
ON batch_job(key_id) WHERE key_id IS NOT NULL;
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_batch_job_native
|
||||
ON batch_job(status, execution_mode)
|
||||
WHERE execution_mode = 'native' AND status = 'processing';
|
||||
|
||||
CREATE TABLE IF NOT EXISTS batch_item (
|
||||
item_id TEXT PRIMARY KEY,
|
||||
batch_id TEXT NOT NULL REFERENCES batch_job(batch_id),
|
||||
custom_id TEXT NOT NULL,
|
||||
status TEXT NOT NULL DEFAULT 'pending',
|
||||
model TEXT NOT NULL,
|
||||
request_body TEXT NOT NULL,
|
||||
source_format TEXT NOT NULL,
|
||||
result_status INTEGER,
|
||||
result_body TEXT,
|
||||
attempts INTEGER NOT NULL DEFAULT 0,
|
||||
max_retries INTEGER NOT NULL DEFAULT 3,
|
||||
last_error TEXT,
|
||||
idempotency_key TEXT,
|
||||
next_retry_at TEXT,
|
||||
lease_id TEXT,
|
||||
lease_expires_at TEXT,
|
||||
created_at TEXT NOT NULL,
|
||||
completed_at TEXT,
|
||||
UNIQUE(batch_id, custom_id)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_batch_item_claim
|
||||
ON batch_item(status, next_retry_at, created_at)
|
||||
WHERE status IN ('pending', 'failed');
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_batch_item_batch
|
||||
ON batch_item(batch_id, status);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_batch_item_lease
|
||||
ON batch_item(lease_expires_at)
|
||||
WHERE lease_id IS NOT NULL;
|
||||
|
||||
CREATE TABLE IF NOT EXISTS batch_dead_letter (
|
||||
item_id TEXT PRIMARY KEY,
|
||||
batch_id TEXT NOT NULL,
|
||||
custom_id TEXT NOT NULL,
|
||||
request_body TEXT NOT NULL,
|
||||
last_error TEXT,
|
||||
attempts INTEGER NOT NULL,
|
||||
failed_at TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS batch_file (
|
||||
file_id TEXT PRIMARY KEY,
|
||||
key_id INTEGER,
|
||||
purpose TEXT NOT NULL DEFAULT 'batch',
|
||||
filename TEXT,
|
||||
byte_size INTEGER NOT NULL,
|
||||
line_count INTEGER NOT NULL,
|
||||
content BLOB NOT NULL,
|
||||
created_at TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS anthropic_batch_map (
|
||||
our_batch_id TEXT PRIMARY KEY,
|
||||
engine_batch_id TEXT NOT NULL,
|
||||
model TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS webhook_delivery (
|
||||
delivery_id TEXT PRIMARY KEY,
|
||||
event_id TEXT NOT NULL,
|
||||
batch_id TEXT NOT NULL,
|
||||
url TEXT NOT NULL,
|
||||
payload TEXT NOT NULL,
|
||||
signing_secret TEXT,
|
||||
status TEXT NOT NULL DEFAULT 'pending',
|
||||
attempts INTEGER NOT NULL DEFAULT 0,
|
||||
max_retries INTEGER NOT NULL DEFAULT 3,
|
||||
next_retry_at TEXT,
|
||||
lease_id TEXT,
|
||||
lease_expires_at TEXT,
|
||||
created_at TEXT NOT NULL,
|
||||
delivered_at TEXT
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_webhook_claim
|
||||
ON webhook_delivery(status, next_retry_at)
|
||||
WHERE status IN ('pending', 'processing');
|
||||
|
||||
CREATE TABLE IF NOT EXISTS batch_event_log (
|
||||
event_id TEXT PRIMARY KEY,
|
||||
batch_id TEXT NOT NULL,
|
||||
sequence INTEGER NOT NULL,
|
||||
event_type TEXT NOT NULL,
|
||||
payload TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL,
|
||||
UNIQUE(batch_id, sequence)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_event_log_batch
|
||||
ON batch_event_log(batch_id, sequence);
|
||||
",
|
||||
)
|
||||
}
|
||||
|
||||
/// Migrate old batch tables (from proxy's admin/db.rs schema) if they exist.
|
||||
/// Renames them to _v1 suffix. Safe to call multiple times.
|
||||
pub fn migrate_old_tables(conn: &Connection) -> rusqlite::Result<()> {
|
||||
// Check if old-schema batch_job exists (has `backend_name` column).
|
||||
let has_old_batch_job: bool = conn
|
||||
.prepare("SELECT 1 FROM pragma_table_info('batch_job') WHERE name = 'backend_name'")
|
||||
.and_then(|mut s| s.exists([]))
|
||||
.unwrap_or(false);
|
||||
|
||||
if has_old_batch_job {
|
||||
conn.execute_batch(
|
||||
"ALTER TABLE batch_job RENAME TO batch_job_v1;
|
||||
ALTER TABLE batch_file RENAME TO batch_file_v1;",
|
||||
)?;
|
||||
tracing::info!("migrated old batch_job and batch_file tables to _v1");
|
||||
}
|
||||
|
||||
// Migrate old anthropic_batch_map if it has openai_batch_id column.
|
||||
let has_old_abm: bool = conn
|
||||
.prepare(
|
||||
"SELECT 1 FROM pragma_table_info('anthropic_batch_map') WHERE name = 'openai_batch_id'",
|
||||
)
|
||||
.and_then(|mut s| s.exists([]))
|
||||
.unwrap_or(false);
|
||||
|
||||
if has_old_abm {
|
||||
conn.execute_batch("ALTER TABLE anthropic_batch_map RENAME TO anthropic_batch_map_v1;")?;
|
||||
tracing::info!("migrated old anthropic_batch_map to _v1");
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn init_tables_succeeds() {
|
||||
let conn = Connection::open_in_memory().unwrap();
|
||||
init_batch_engine_tables(&conn).unwrap();
|
||||
|
||||
// Verify tables exist by querying them.
|
||||
let tables = [
|
||||
"batch_job",
|
||||
"batch_item",
|
||||
"batch_file",
|
||||
"webhook_delivery",
|
||||
"batch_event_log",
|
||||
];
|
||||
for table in tables {
|
||||
let count: i64 = conn
|
||||
.query_row(&format!("SELECT count(*) FROM {table}"), [], |r| r.get(0))
|
||||
.unwrap();
|
||||
assert_eq!(count, 0, "expected empty table {table}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn init_tables_idempotent() {
|
||||
let conn = Connection::open_in_memory().unwrap();
|
||||
init_batch_engine_tables(&conn).unwrap();
|
||||
init_batch_engine_tables(&conn).unwrap(); // no error
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn now_iso8601_format() {
|
||||
let ts = now_iso8601();
|
||||
assert!(ts.ends_with('Z'));
|
||||
assert!(ts.contains('T'));
|
||||
assert_eq!(ts.len(), 20); // "YYYY-MM-DDTHH:MM:SSZ"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,305 @@
|
||||
// crates/batch_engine/src/engine.rs
|
||||
//! BatchEngine: the main entry point for batch operations.
|
||||
//! Thin facade over JobQueue, FileStore, and WebhookQueue.
|
||||
|
||||
use crate::db::now_iso8601;
|
||||
use crate::error::EngineError;
|
||||
use crate::file_store::FileStore;
|
||||
use crate::job::*;
|
||||
use crate::queue::JobQueue;
|
||||
use crate::webhook::{WebhookDelivery, WebhookQueue};
|
||||
use std::sync::Arc;
|
||||
|
||||
/// The main batch engine. Holds references to queue, file store, and webhook queue.
|
||||
pub struct BatchEngine<Q: JobQueue, W: WebhookQueue> {
|
||||
pub queue: Arc<Q>,
|
||||
pub file_store: FileStore,
|
||||
pub webhook_queue: Arc<W>,
|
||||
pub global_webhook_urls: Vec<String>,
|
||||
pub webhook_signing_secret: Option<String>,
|
||||
}
|
||||
|
||||
impl<Q: JobQueue, W: WebhookQueue> BatchEngine<Q, W> {
|
||||
/// Submit a new batch job.
|
||||
pub async fn submit(&self, submission: BatchSubmission) -> Result<BatchJob, EngineError> {
|
||||
// Verify input file exists.
|
||||
self.file_store
|
||||
.get_meta(&submission.input_file_id)
|
||||
.await
|
||||
.map_err(|e| EngineError::Backend(e.to_string()))?
|
||||
.ok_or_else(|| EngineError::FileNotFound(submission.input_file_id.clone()))?;
|
||||
|
||||
let now = now_iso8601();
|
||||
let batch_id = BatchId::new();
|
||||
let total = submission.items.len() as u32;
|
||||
|
||||
let job = BatchJob {
|
||||
id: batch_id.clone(),
|
||||
status: BatchStatus::Queued,
|
||||
execution_mode: submission.execution_mode.clone(),
|
||||
priority: submission.priority,
|
||||
key_id: submission.key_id,
|
||||
input_file_id: submission.input_file_id,
|
||||
webhook_url: submission.webhook_url.clone(),
|
||||
metadata: submission.metadata,
|
||||
request_counts: RequestCounts {
|
||||
total,
|
||||
..Default::default()
|
||||
},
|
||||
created_at: now.clone(),
|
||||
started_at: None,
|
||||
completed_at: None,
|
||||
expires_at: now.clone(), // TODO: add 24h
|
||||
};
|
||||
|
||||
let items: Vec<BatchItem> = submission
|
||||
.items
|
||||
.into_iter()
|
||||
.map(|si| BatchItem {
|
||||
id: ItemId::new(),
|
||||
batch_id: batch_id.clone(),
|
||||
custom_id: si.custom_id,
|
||||
status: ItemStatus::Pending,
|
||||
request: BatchItemRequest {
|
||||
model: si.model,
|
||||
body: si.body,
|
||||
source_format: si.source_format,
|
||||
},
|
||||
result: None,
|
||||
attempts: 0,
|
||||
max_retries: 3,
|
||||
last_error: None,
|
||||
next_retry_at: None,
|
||||
lease_id: None,
|
||||
lease_expires_at: None,
|
||||
idempotency_key: None,
|
||||
created_at: now.clone(),
|
||||
completed_at: None,
|
||||
})
|
||||
.collect();
|
||||
|
||||
self.queue
|
||||
.enqueue(&job, &items)
|
||||
.await
|
||||
.map_err(EngineError::Queue)?;
|
||||
|
||||
// Fire webhook for batch.queued.
|
||||
self.fire_webhook(
|
||||
&batch_id,
|
||||
"batch.queued",
|
||||
serde_json::json!({
|
||||
"batch_id": batch_id.0,
|
||||
"total_items": total,
|
||||
"execution_mode": job.execution_mode.as_str(),
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
Ok(job)
|
||||
}
|
||||
|
||||
/// Get a batch job by ID.
|
||||
pub async fn get(&self, id: &BatchId) -> Result<Option<BatchJob>, EngineError> {
|
||||
self.queue.get(id).await.map_err(EngineError::Queue)
|
||||
}
|
||||
|
||||
/// List batch jobs.
|
||||
pub async fn list(
|
||||
&self,
|
||||
key_id: Option<i64>,
|
||||
cursor: Option<&str>,
|
||||
limit: u32,
|
||||
) -> Result<Vec<BatchJob>, EngineError> {
|
||||
self.queue
|
||||
.list(key_id, cursor, limit)
|
||||
.await
|
||||
.map_err(EngineError::Queue)
|
||||
}
|
||||
|
||||
/// Cancel a batch job.
|
||||
pub async fn cancel(&self, id: &BatchId) -> Result<BatchStatus, EngineError> {
|
||||
let status = self.queue.cancel(id).await.map_err(EngineError::Queue)?;
|
||||
|
||||
if status == BatchStatus::Cancelled {
|
||||
self.fire_webhook(
|
||||
id,
|
||||
"batch.cancelled",
|
||||
serde_json::json!({ "batch_id": id.0 }),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
Ok(status)
|
||||
}
|
||||
|
||||
/// Get items for a batch (used for result retrieval).
|
||||
pub async fn get_items(&self, id: &BatchId) -> Result<Vec<BatchItem>, EngineError> {
|
||||
self.queue.get_items(id).await.map_err(EngineError::Queue)
|
||||
}
|
||||
|
||||
/// Fire a webhook to all configured URLs.
|
||||
async fn fire_webhook(&self, batch_id: &BatchId, event_type: &str, payload: serde_json::Value) {
|
||||
let event_id = format!("evt_{}", uuid::Uuid::new_v4());
|
||||
|
||||
// Collect URLs: global + per-batch.
|
||||
let mut urls: Vec<(String, Option<String>)> = self
|
||||
.global_webhook_urls
|
||||
.iter()
|
||||
.map(|u| (u.clone(), self.webhook_signing_secret.clone()))
|
||||
.collect();
|
||||
|
||||
// Per-batch webhook gets terminal events only.
|
||||
if matches!(
|
||||
event_type,
|
||||
"batch.completed" | "batch.failed" | "batch.cancelled"
|
||||
) {
|
||||
if let Ok(Some(job)) = self.queue.get(batch_id).await {
|
||||
if let Some(url) = job.webhook_url {
|
||||
urls.push((url, self.webhook_signing_secret.clone()));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let full_payload = serde_json::json!({
|
||||
"event_id": event_id,
|
||||
"event_type": event_type,
|
||||
"data": payload,
|
||||
});
|
||||
|
||||
for (url, secret) in urls {
|
||||
let delivery = WebhookDelivery {
|
||||
delivery_id: format!("whd_{}", uuid::Uuid::new_v4()),
|
||||
event_id: event_id.clone(),
|
||||
batch_id: batch_id.0.clone(),
|
||||
url,
|
||||
payload: full_payload.clone(),
|
||||
signing_secret: secret,
|
||||
attempts: 0,
|
||||
max_retries: 3,
|
||||
next_retry_at: None,
|
||||
};
|
||||
if let Err(e) = self.webhook_queue.enqueue(delivery).await {
|
||||
tracing::error!(error = %e, "failed to enqueue webhook delivery");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::db::init_batch_engine_tables;
|
||||
use crate::file_store::FileStore;
|
||||
use crate::queue::sqlite::SqliteQueue;
|
||||
use crate::webhook::sqlite::SqliteWebhookQueue;
|
||||
use rusqlite::Connection;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
async fn test_engine() -> BatchEngine<SqliteQueue, SqliteWebhookQueue> {
|
||||
let conn = Connection::open_in_memory().unwrap();
|
||||
init_batch_engine_tables(&conn).unwrap();
|
||||
let db = Arc::new(Mutex::new(conn));
|
||||
|
||||
BatchEngine {
|
||||
queue: Arc::new(SqliteQueue::new(db.clone())),
|
||||
file_store: FileStore::new(db.clone()),
|
||||
webhook_queue: Arc::new(SqliteWebhookQueue::new(db)),
|
||||
global_webhook_urls: vec![],
|
||||
webhook_signing_secret: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn submit_and_get() {
|
||||
let engine = test_engine().await;
|
||||
|
||||
// Upload a file first.
|
||||
engine
|
||||
.file_store
|
||||
.insert("file-sub1", None, None, b"test", 2)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let job = engine
|
||||
.submit(BatchSubmission {
|
||||
items: vec![
|
||||
SubmissionItem {
|
||||
custom_id: "req-1".into(),
|
||||
model: "gpt-4o".into(),
|
||||
body: serde_json::json!({}),
|
||||
source_format: SourceFormat::OpenAI,
|
||||
},
|
||||
SubmissionItem {
|
||||
custom_id: "req-2".into(),
|
||||
model: "gpt-4o".into(),
|
||||
body: serde_json::json!({}),
|
||||
source_format: SourceFormat::OpenAI,
|
||||
},
|
||||
],
|
||||
execution_mode: ExecutionMode::ProxyNative,
|
||||
input_file_id: "file-sub1".into(),
|
||||
key_id: Some(42),
|
||||
webhook_url: None,
|
||||
metadata: None,
|
||||
priority: 0,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(job.status, BatchStatus::Queued);
|
||||
assert_eq!(job.request_counts.total, 2);
|
||||
assert_eq!(job.key_id, Some(42));
|
||||
|
||||
let fetched = engine.get(&job.id).await.unwrap().unwrap();
|
||||
assert_eq!(fetched.id, job.id);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn submit_missing_file() {
|
||||
let engine = test_engine().await;
|
||||
let result = engine
|
||||
.submit(BatchSubmission {
|
||||
items: vec![],
|
||||
execution_mode: ExecutionMode::ProxyNative,
|
||||
input_file_id: "file-nope".into(),
|
||||
key_id: None,
|
||||
webhook_url: None,
|
||||
metadata: None,
|
||||
priority: 0,
|
||||
})
|
||||
.await;
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cancel_job() {
|
||||
let engine = test_engine().await;
|
||||
engine
|
||||
.file_store
|
||||
.insert("file-cancel", None, None, b"test", 1)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let job = engine
|
||||
.submit(BatchSubmission {
|
||||
items: vec![SubmissionItem {
|
||||
custom_id: "r1".into(),
|
||||
model: "gpt-4o".into(),
|
||||
body: serde_json::json!({}),
|
||||
source_format: SourceFormat::OpenAI,
|
||||
}],
|
||||
execution_mode: ExecutionMode::ProxyNative,
|
||||
input_file_id: "file-cancel".into(),
|
||||
key_id: None,
|
||||
webhook_url: None,
|
||||
metadata: None,
|
||||
priority: 0,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let status = engine.cancel(&job.id).await.unwrap();
|
||||
assert_eq!(status, BatchStatus::Cancelled);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,40 @@
|
||||
// crates/batch_engine/src/error.rs
|
||||
//! Engine and queue error types.
|
||||
|
||||
use thiserror::Error;
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub enum EngineError {
|
||||
#[error("batch not found: {0}")]
|
||||
NotFound(String),
|
||||
|
||||
#[error("file not found: {0}")]
|
||||
FileNotFound(String),
|
||||
|
||||
#[error("validation error: {0}")]
|
||||
Validation(String),
|
||||
|
||||
#[error("queue error: {0}")]
|
||||
Queue(#[from] QueueError),
|
||||
|
||||
#[error("backend error: {0}")]
|
||||
Backend(String),
|
||||
}
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub enum QueueError {
|
||||
#[error("not found")]
|
||||
NotFound,
|
||||
|
||||
#[error("already claimed")]
|
||||
AlreadyClaimed,
|
||||
|
||||
#[error("storage error: {0}")]
|
||||
Storage(String),
|
||||
}
|
||||
|
||||
impl From<rusqlite::Error> for QueueError {
|
||||
fn from(e: rusqlite::Error) -> Self {
|
||||
QueueError::Storage(e.to_string())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,149 @@
|
||||
// crates/batch_engine/src/file_store.rs
|
||||
//! Batch file storage (upload, metadata, content retrieval).
|
||||
|
||||
use crate::db::now_iso8601;
|
||||
use crate::error::QueueError;
|
||||
use rusqlite::{params, Connection};
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
/// Batch file metadata (without content blob).
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct BatchFileMeta {
|
||||
pub file_id: String,
|
||||
pub byte_size: i64,
|
||||
pub line_count: i64,
|
||||
pub filename: Option<String>,
|
||||
pub created_at: String,
|
||||
}
|
||||
|
||||
/// Manages batch file storage in SQLite.
|
||||
#[derive(Clone)]
|
||||
pub struct FileStore {
|
||||
db: Arc<Mutex<Connection>>,
|
||||
}
|
||||
|
||||
impl FileStore {
|
||||
pub fn new(db: Arc<Mutex<Connection>>) -> Self {
|
||||
Self { db }
|
||||
}
|
||||
|
||||
/// Store a batch file. Returns the file_id.
|
||||
pub async fn insert(
|
||||
&self,
|
||||
file_id: &str,
|
||||
key_id: Option<i64>,
|
||||
filename: Option<&str>,
|
||||
content: &[u8],
|
||||
line_count: i64,
|
||||
) -> Result<(), QueueError> {
|
||||
let db = self.db.clone();
|
||||
let file_id = file_id.to_string();
|
||||
let filename = filename.map(|s| s.to_string());
|
||||
let content = content.to_vec();
|
||||
let byte_size = content.len() as i64;
|
||||
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let conn = db.blocking_lock();
|
||||
conn.execute(
|
||||
"INSERT INTO batch_file (file_id, key_id, purpose, filename, byte_size, line_count, content, created_at)
|
||||
VALUES (?1, ?2, 'batch', ?3, ?4, ?5, ?6, ?7)",
|
||||
params![file_id, key_id, filename, byte_size, line_count, content, now_iso8601()],
|
||||
)?;
|
||||
Ok(())
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
/// Get file metadata (without content).
|
||||
pub async fn get_meta(&self, file_id: &str) -> Result<Option<BatchFileMeta>, QueueError> {
|
||||
let db = self.db.clone();
|
||||
let file_id = file_id.to_string();
|
||||
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let conn = db.blocking_lock();
|
||||
let mut stmt = conn.prepare(
|
||||
"SELECT file_id, byte_size, line_count, filename, created_at FROM batch_file WHERE file_id = ?1",
|
||||
)?;
|
||||
let mut rows = stmt.query(params![file_id])?;
|
||||
if let Some(row) = rows.next()? {
|
||||
Ok(Some(BatchFileMeta {
|
||||
file_id: row.get(0)?,
|
||||
byte_size: row.get(1)?,
|
||||
line_count: row.get(2)?,
|
||||
filename: row.get(3)?,
|
||||
created_at: row.get(4)?,
|
||||
}))
|
||||
} else {
|
||||
Ok(None)
|
||||
}
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
/// Get file content (raw JSONL bytes).
|
||||
pub async fn get_content(&self, file_id: &str) -> Result<Option<Vec<u8>>, QueueError> {
|
||||
let db = self.db.clone();
|
||||
let file_id = file_id.to_string();
|
||||
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let conn = db.blocking_lock();
|
||||
let mut stmt = conn.prepare("SELECT content FROM batch_file WHERE file_id = ?1")?;
|
||||
let mut rows = stmt.query(params![file_id])?;
|
||||
if let Some(row) = rows.next()? {
|
||||
let content: Vec<u8> = row.get(0)?;
|
||||
Ok(Some(content))
|
||||
} else {
|
||||
Ok(None)
|
||||
}
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::db::init_batch_engine_tables;
|
||||
|
||||
async fn test_store() -> FileStore {
|
||||
let conn = Connection::open_in_memory().unwrap();
|
||||
init_batch_engine_tables(&conn).unwrap();
|
||||
FileStore::new(Arc::new(Mutex::new(conn)))
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn insert_and_get_meta() {
|
||||
let store = test_store().await;
|
||||
store
|
||||
.insert("file-abc", None, Some("test.jsonl"), b"line1\nline2", 2)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let meta = store.get_meta("file-abc").await.unwrap().unwrap();
|
||||
assert_eq!(meta.file_id, "file-abc");
|
||||
assert_eq!(meta.byte_size, 11);
|
||||
assert_eq!(meta.line_count, 2);
|
||||
assert_eq!(meta.filename.as_deref(), Some("test.jsonl"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn get_content_roundtrip() {
|
||||
let store = test_store().await;
|
||||
let data = b"test content bytes";
|
||||
store.insert("file-xyz", None, None, data, 1).await.unwrap();
|
||||
|
||||
let content = store.get_content("file-xyz").await.unwrap().unwrap();
|
||||
assert_eq!(content, data);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn get_nonexistent_returns_none() {
|
||||
let store = test_store().await;
|
||||
assert!(store.get_meta("file-nope").await.unwrap().is_none());
|
||||
assert!(store.get_content("file-nope").await.unwrap().is_none());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,307 @@
|
||||
// crates/batch_engine/src/job.rs
|
||||
//! Core batch orchestration types. HTTP-agnostic.
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Unique batch job identifier. Format: "batch_{uuid}".
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
||||
pub struct BatchId(pub String);
|
||||
|
||||
/// Unique item identifier within a batch. Format: "item_{uuid}".
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
||||
pub struct ItemId(pub String);
|
||||
|
||||
impl BatchId {
|
||||
pub fn new() -> Self {
|
||||
Self(format!("batch_{}", uuid::Uuid::new_v4()))
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for BatchId {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl ItemId {
|
||||
pub fn new() -> Self {
|
||||
Self(format!("item_{}", uuid::Uuid::new_v4()))
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for ItemId {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Display for BatchId {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.write_str(&self.0)
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Display for ItemId {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.write_str(&self.0)
|
||||
}
|
||||
}
|
||||
|
||||
/// Batch job lifecycle status.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum BatchStatus {
|
||||
Queued,
|
||||
Processing,
|
||||
Completed,
|
||||
Failed,
|
||||
Cancelling,
|
||||
Cancelled,
|
||||
Expired,
|
||||
}
|
||||
|
||||
impl BatchStatus {
|
||||
pub fn as_str(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Queued => "queued",
|
||||
Self::Processing => "processing",
|
||||
Self::Completed => "completed",
|
||||
Self::Failed => "failed",
|
||||
Self::Cancelling => "cancelling",
|
||||
Self::Cancelled => "cancelled",
|
||||
Self::Expired => "expired",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn from_str_status(s: &str) -> Self {
|
||||
match s {
|
||||
"queued" => Self::Queued,
|
||||
"processing" => Self::Processing,
|
||||
"completed" => Self::Completed,
|
||||
"failed" => Self::Failed,
|
||||
"cancelling" => Self::Cancelling,
|
||||
"cancelled" => Self::Cancelled,
|
||||
"expired" => Self::Expired,
|
||||
_ => Self::Failed,
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether this status is terminal (no further transitions).
|
||||
pub fn is_terminal(&self) -> bool {
|
||||
matches!(
|
||||
self,
|
||||
Self::Completed | Self::Failed | Self::Cancelled | Self::Expired
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// How the batch will be executed.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case", tag = "type")]
|
||||
pub enum ExecutionMode {
|
||||
/// Delegate to provider's native batch API (OpenAI, Azure).
|
||||
Native { provider: String },
|
||||
/// Proxy processes items individually against the backend.
|
||||
ProxyNative,
|
||||
}
|
||||
|
||||
impl ExecutionMode {
|
||||
pub fn as_str(&self) -> &str {
|
||||
match self {
|
||||
Self::Native { .. } => "native",
|
||||
Self::ProxyNative => "proxy_native",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn provider(&self) -> Option<&str> {
|
||||
match self {
|
||||
Self::Native { provider } => Some(provider),
|
||||
Self::ProxyNative => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A batch job as seen by the engine.
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct BatchJob {
|
||||
pub id: BatchId,
|
||||
pub status: BatchStatus,
|
||||
pub execution_mode: ExecutionMode,
|
||||
pub priority: u8,
|
||||
pub key_id: Option<i64>,
|
||||
pub webhook_url: Option<String>,
|
||||
pub metadata: Option<serde_json::Value>,
|
||||
pub request_counts: RequestCounts,
|
||||
pub input_file_id: String,
|
||||
pub created_at: String,
|
||||
pub started_at: Option<String>,
|
||||
pub completed_at: Option<String>,
|
||||
pub expires_at: String,
|
||||
}
|
||||
|
||||
/// Counts of requests within a batch job.
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
||||
pub struct RequestCounts {
|
||||
pub total: u32,
|
||||
pub processing: u32,
|
||||
pub succeeded: u32,
|
||||
pub failed: u32,
|
||||
pub cancelled: u32,
|
||||
pub expired: u32,
|
||||
}
|
||||
|
||||
/// Single item within a batch.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct BatchItem {
|
||||
pub id: ItemId,
|
||||
pub batch_id: BatchId,
|
||||
pub custom_id: String,
|
||||
pub status: ItemStatus,
|
||||
pub request: BatchItemRequest,
|
||||
pub result: Option<BatchItemResult>,
|
||||
pub attempts: u8,
|
||||
pub max_retries: u8,
|
||||
pub last_error: Option<String>,
|
||||
pub next_retry_at: Option<String>,
|
||||
pub lease_id: Option<String>,
|
||||
pub lease_expires_at: Option<String>,
|
||||
pub idempotency_key: Option<String>,
|
||||
pub created_at: String,
|
||||
pub completed_at: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum ItemStatus {
|
||||
Pending,
|
||||
Processing,
|
||||
Succeeded,
|
||||
Failed,
|
||||
Cancelled,
|
||||
}
|
||||
|
||||
impl ItemStatus {
|
||||
pub fn as_str(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Pending => "pending",
|
||||
Self::Processing => "processing",
|
||||
Self::Succeeded => "succeeded",
|
||||
Self::Failed => "failed",
|
||||
Self::Cancelled => "cancelled",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn from_str_status(s: &str) -> Self {
|
||||
match s {
|
||||
"pending" => Self::Pending,
|
||||
"processing" => Self::Processing,
|
||||
"succeeded" => Self::Succeeded,
|
||||
"failed" => Self::Failed,
|
||||
"cancelled" => Self::Cancelled,
|
||||
_ => Self::Failed,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn is_terminal(&self) -> bool {
|
||||
matches!(self, Self::Succeeded | Self::Failed | Self::Cancelled)
|
||||
}
|
||||
}
|
||||
|
||||
/// The LLM request payload for a batch item.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct BatchItemRequest {
|
||||
pub model: String,
|
||||
pub body: serde_json::Value,
|
||||
pub source_format: SourceFormat,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum SourceFormat {
|
||||
Anthropic,
|
||||
OpenAI,
|
||||
}
|
||||
|
||||
/// Result of executing a single batch item.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct BatchItemResult {
|
||||
pub status_code: u16,
|
||||
pub body: serde_json::Value,
|
||||
}
|
||||
|
||||
/// Submission request to the engine (from proxy handlers).
|
||||
pub struct BatchSubmission {
|
||||
pub items: Vec<SubmissionItem>,
|
||||
pub execution_mode: ExecutionMode,
|
||||
pub input_file_id: String,
|
||||
pub key_id: Option<i64>,
|
||||
pub webhook_url: Option<String>,
|
||||
pub metadata: Option<serde_json::Value>,
|
||||
pub priority: u8,
|
||||
}
|
||||
|
||||
/// A single item in a batch submission.
|
||||
pub struct SubmissionItem {
|
||||
pub custom_id: String,
|
||||
pub model: String,
|
||||
pub body: serde_json::Value,
|
||||
pub source_format: SourceFormat,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn batch_id_format() {
|
||||
let id = BatchId::new();
|
||||
assert!(id.0.starts_with("batch_"));
|
||||
assert_eq!(id.0.len(), 6 + 36); // "batch_" + uuid
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn item_id_format() {
|
||||
let id = ItemId::new();
|
||||
assert!(id.0.starts_with("item_"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn batch_status_roundtrip() {
|
||||
for status in [
|
||||
BatchStatus::Queued,
|
||||
BatchStatus::Processing,
|
||||
BatchStatus::Completed,
|
||||
BatchStatus::Failed,
|
||||
BatchStatus::Cancelling,
|
||||
BatchStatus::Cancelled,
|
||||
BatchStatus::Expired,
|
||||
] {
|
||||
let s = status.as_str();
|
||||
let parsed = BatchStatus::from_str_status(s);
|
||||
assert_eq!(status, parsed);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn terminal_statuses() {
|
||||
assert!(!BatchStatus::Queued.is_terminal());
|
||||
assert!(!BatchStatus::Processing.is_terminal());
|
||||
assert!(BatchStatus::Completed.is_terminal());
|
||||
assert!(BatchStatus::Failed.is_terminal());
|
||||
assert!(BatchStatus::Cancelled.is_terminal());
|
||||
assert!(BatchStatus::Expired.is_terminal());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn execution_mode_str() {
|
||||
let native = ExecutionMode::Native {
|
||||
provider: "openai".into(),
|
||||
};
|
||||
assert_eq!(native.as_str(), "native");
|
||||
assert_eq!(native.provider(), Some("openai"));
|
||||
|
||||
let proxy = ExecutionMode::ProxyNative;
|
||||
assert_eq!(proxy.as_str(), "proxy_native");
|
||||
assert_eq!(proxy.provider(), None);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,18 @@
|
||||
// crates/batch_engine/src/lib.rs
|
||||
//! Batch orchestration engine: job queue, file storage, webhook delivery.
|
||||
//!
|
||||
//! HTTP-agnostic. The proxy crate wires this into axum routes.
|
||||
|
||||
pub mod db;
|
||||
pub mod engine;
|
||||
pub mod error;
|
||||
pub mod file_store;
|
||||
pub mod job;
|
||||
pub mod queue;
|
||||
pub mod validation;
|
||||
pub mod webhook;
|
||||
|
||||
pub use engine::BatchEngine;
|
||||
pub use error::{EngineError, QueueError};
|
||||
pub use job::*;
|
||||
pub use validation::{validate_jsonl, ValidatedJsonl};
|
||||
@@ -0,0 +1,62 @@
|
||||
// crates/batch_engine/src/queue/mod.rs
|
||||
//! JobQueue trait and implementations.
|
||||
|
||||
pub mod sqlite;
|
||||
|
||||
use crate::error::QueueError;
|
||||
use crate::job::*;
|
||||
use async_trait::async_trait;
|
||||
use std::time::Duration;
|
||||
|
||||
/// A job that has been claimed by a worker. Holds lease metadata.
|
||||
#[derive(Debug)]
|
||||
pub struct LeasedItem {
|
||||
pub item: BatchItem,
|
||||
pub batch_id: BatchId,
|
||||
pub lease_id: String,
|
||||
pub lease_expires_at: String,
|
||||
}
|
||||
|
||||
/// Core queue abstraction. All methods are async + Send.
|
||||
#[async_trait]
|
||||
pub trait JobQueue: Send + Sync + 'static {
|
||||
// -- Job lifecycle --
|
||||
async fn enqueue(&self, job: &BatchJob, items: &[BatchItem]) -> Result<(), QueueError>;
|
||||
async fn get(&self, id: &BatchId) -> Result<Option<BatchJob>, QueueError>;
|
||||
async fn list(
|
||||
&self,
|
||||
key_id: Option<i64>,
|
||||
cursor: Option<&str>,
|
||||
limit: u32,
|
||||
) -> Result<Vec<BatchJob>, QueueError>;
|
||||
async fn cancel(&self, id: &BatchId) -> Result<BatchStatus, QueueError>;
|
||||
|
||||
// -- Item-level operations (proxy-native path, Phase 2) --
|
||||
async fn claim_next_item(&self) -> Result<Option<LeasedItem>, QueueError>;
|
||||
async fn complete_item(&self, id: &ItemId, result: BatchItemResult) -> Result<(), QueueError>;
|
||||
async fn fail_item(&self, id: &ItemId, error: &str) -> Result<(), QueueError>;
|
||||
async fn schedule_retry(
|
||||
&self,
|
||||
id: &ItemId,
|
||||
delay: Duration,
|
||||
error: &str,
|
||||
) -> Result<(), QueueError>;
|
||||
async fn dead_letter(&self, id: &ItemId) -> Result<(), QueueError>;
|
||||
|
||||
// -- Batch completion --
|
||||
async fn is_batch_complete(&self, id: &BatchId) -> Result<bool, QueueError>;
|
||||
async fn complete_batch(&self, id: &BatchId) -> Result<(), QueueError>;
|
||||
|
||||
// -- Native batch support --
|
||||
async fn get_native_jobs_in_progress(&self) -> Result<Vec<BatchJob>, QueueError>;
|
||||
|
||||
// -- Lease management --
|
||||
async fn reclaim_expired_leases(&self) -> Result<u32, QueueError>;
|
||||
|
||||
// -- Progress --
|
||||
async fn update_progress(&self, id: &BatchId, counts: &RequestCounts)
|
||||
-> Result<(), QueueError>;
|
||||
|
||||
// -- Items query --
|
||||
async fn get_items(&self, batch_id: &BatchId) -> Result<Vec<BatchItem>, QueueError>;
|
||||
}
|
||||
@@ -0,0 +1,790 @@
|
||||
// crates/batch_engine/src/queue/sqlite.rs
|
||||
//! SQLite-backed JobQueue implementation.
|
||||
|
||||
use super::{JobQueue, LeasedItem};
|
||||
use crate::db::now_iso8601;
|
||||
use crate::error::QueueError;
|
||||
use crate::job::*;
|
||||
use async_trait::async_trait;
|
||||
use rusqlite::{params, Connection};
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
/// SQLite-backed job queue. Suitable for single-instance deployments.
|
||||
#[derive(Clone)]
|
||||
pub struct SqliteQueue {
|
||||
db: Arc<Mutex<Connection>>,
|
||||
}
|
||||
|
||||
impl SqliteQueue {
|
||||
pub fn new(db: Arc<Mutex<Connection>>) -> Self {
|
||||
Self { db }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl JobQueue for SqliteQueue {
|
||||
async fn enqueue(&self, job: &BatchJob, items: &[BatchItem]) -> Result<(), QueueError> {
|
||||
let db = self.db.clone();
|
||||
let job = job.clone();
|
||||
let items: Vec<BatchItem> = items.to_vec();
|
||||
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let conn = db.blocking_lock();
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
|
||||
tx.execute(
|
||||
"INSERT INTO batch_job (batch_id, status, execution_mode, provider, priority,
|
||||
key_id, input_file_id, webhook_url, metadata, total, created_at, expires_at)
|
||||
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12)",
|
||||
params![
|
||||
job.id.0,
|
||||
job.status.as_str(),
|
||||
job.execution_mode.as_str(),
|
||||
job.execution_mode.provider(),
|
||||
job.priority,
|
||||
job.key_id,
|
||||
job.input_file_id,
|
||||
job.webhook_url,
|
||||
job.metadata
|
||||
.as_ref()
|
||||
.map(|m| serde_json::to_string(m).unwrap_or_default()),
|
||||
job.request_counts.total,
|
||||
job.created_at,
|
||||
job.expires_at,
|
||||
],
|
||||
)?;
|
||||
|
||||
for item in &items {
|
||||
let body_str = serde_json::to_string(&item.request.body)
|
||||
.map_err(|e| QueueError::Storage(e.to_string()))?;
|
||||
tx.execute(
|
||||
"INSERT INTO batch_item (item_id, batch_id, custom_id, status, model,
|
||||
request_body, source_format, max_retries, created_at)
|
||||
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)",
|
||||
params![
|
||||
item.id.0,
|
||||
item.batch_id.0,
|
||||
item.custom_id,
|
||||
item.status.as_str(),
|
||||
item.request.model,
|
||||
body_str,
|
||||
serde_json::to_string(&item.request.source_format)
|
||||
.unwrap_or_else(|_| "\"openai\"".to_string()),
|
||||
item.max_retries,
|
||||
item.created_at,
|
||||
],
|
||||
)?;
|
||||
}
|
||||
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
async fn get(&self, id: &BatchId) -> Result<Option<BatchJob>, QueueError> {
|
||||
let db = self.db.clone();
|
||||
let id = id.0.clone();
|
||||
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let conn = db.blocking_lock();
|
||||
row_to_job(&conn, &id)
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
async fn list(
|
||||
&self,
|
||||
key_id: Option<i64>,
|
||||
cursor: Option<&str>,
|
||||
limit: u32,
|
||||
) -> Result<Vec<BatchJob>, QueueError> {
|
||||
let db = self.db.clone();
|
||||
let cursor = cursor.map(|s| s.to_string());
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let conn = db.blocking_lock();
|
||||
let mut sql = String::from(
|
||||
"SELECT batch_id, status, execution_mode, provider, priority,
|
||||
key_id, input_file_id, webhook_url, metadata,
|
||||
total, processing, succeeded, failed, cancelled, expired,
|
||||
created_at, started_at, completed_at, expires_at
|
||||
FROM batch_job WHERE 1=1",
|
||||
);
|
||||
let mut param_values: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
|
||||
|
||||
if let Some(kid) = key_id {
|
||||
sql.push_str(" AND key_id = ?");
|
||||
param_values.push(Box::new(kid));
|
||||
}
|
||||
if let Some(ref c) = cursor {
|
||||
sql.push_str(
|
||||
" AND created_at < (SELECT created_at FROM batch_job WHERE batch_id = ?)",
|
||||
);
|
||||
param_values.push(Box::new(c.clone()));
|
||||
}
|
||||
sql.push_str(" ORDER BY created_at DESC LIMIT ?");
|
||||
param_values.push(Box::new(limit));
|
||||
|
||||
let params_refs: Vec<&dyn rusqlite::types::ToSql> =
|
||||
param_values.iter().map(|p| p.as_ref()).collect();
|
||||
|
||||
let mut stmt = conn.prepare(&sql)?;
|
||||
let rows = stmt.query_map(params_refs.as_slice(), |row| Ok(batch_job_from_row(row)))?;
|
||||
rows.collect::<Result<Vec<_>, _>>()
|
||||
.map_err(QueueError::from)
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
async fn cancel(&self, id: &BatchId) -> Result<BatchStatus, QueueError> {
|
||||
let db = self.db.clone();
|
||||
let id = id.0.clone();
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let conn = db.blocking_lock();
|
||||
let mut stmt =
|
||||
conn.prepare("SELECT status FROM batch_job WHERE batch_id = ?1")?;
|
||||
let status_str: Option<String> = stmt
|
||||
.query_row(params![id], |row| row.get(0))
|
||||
.ok();
|
||||
|
||||
let Some(status_str) = status_str else {
|
||||
return Err(QueueError::NotFound);
|
||||
};
|
||||
|
||||
let current = BatchStatus::from_str_status(&status_str);
|
||||
let new_status = match current {
|
||||
BatchStatus::Queued => BatchStatus::Cancelled,
|
||||
BatchStatus::Processing => BatchStatus::Cancelling,
|
||||
other if other.is_terminal() => return Ok(other),
|
||||
_ => BatchStatus::Cancelled,
|
||||
};
|
||||
|
||||
conn.execute(
|
||||
"UPDATE batch_job SET status = ?1, completed_at = CASE WHEN ?1 = 'cancelled' THEN ?2 ELSE completed_at END
|
||||
WHERE batch_id = ?3",
|
||||
params![new_status.as_str(), now_iso8601(), id],
|
||||
)?;
|
||||
|
||||
// If directly cancelled (was queued), cancel all pending items.
|
||||
if new_status == BatchStatus::Cancelled {
|
||||
conn.execute(
|
||||
"UPDATE batch_item SET status = 'cancelled' WHERE batch_id = ?1 AND status = 'pending'",
|
||||
params![id],
|
||||
)?;
|
||||
}
|
||||
|
||||
Ok(new_status)
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
async fn claim_next_item(&self) -> Result<Option<LeasedItem>, QueueError> {
|
||||
let db = self.db.clone();
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let conn = db.blocking_lock();
|
||||
let lease_id = format!("lease_{}", uuid::Uuid::new_v4());
|
||||
let now = now_iso8601();
|
||||
// Lease for 120 seconds.
|
||||
let lease_expires = {
|
||||
let secs = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_secs()
|
||||
+ 120;
|
||||
format_epoch_iso8601(secs)
|
||||
};
|
||||
|
||||
let result = conn.query_row(
|
||||
"UPDATE batch_item
|
||||
SET status = 'processing',
|
||||
lease_id = ?1,
|
||||
lease_expires_at = ?2,
|
||||
attempts = attempts + 1
|
||||
WHERE item_id = (
|
||||
SELECT bi.item_id
|
||||
FROM batch_item bi
|
||||
JOIN batch_job bj ON bi.batch_id = bj.batch_id
|
||||
WHERE bi.status IN ('pending')
|
||||
AND (bi.next_retry_at IS NULL OR bi.next_retry_at <= ?3)
|
||||
AND bj.status IN ('queued', 'processing')
|
||||
AND bj.execution_mode = 'proxy_native'
|
||||
ORDER BY bj.priority DESC, bi.created_at ASC
|
||||
LIMIT 1
|
||||
)
|
||||
RETURNING item_id, batch_id, custom_id, status, model, request_body,
|
||||
source_format, result_status, result_body, attempts,
|
||||
max_retries, last_error, next_retry_at, lease_id,
|
||||
lease_expires_at, idempotency_key, created_at, completed_at",
|
||||
params![lease_id, lease_expires, now],
|
||||
|row| {
|
||||
let item = batch_item_from_row(row);
|
||||
Ok(LeasedItem {
|
||||
batch_id: item.batch_id.clone(),
|
||||
lease_id: item.lease_id.clone().unwrap_or_default(),
|
||||
lease_expires_at: item.lease_expires_at.clone().unwrap_or_default(),
|
||||
item,
|
||||
})
|
||||
},
|
||||
);
|
||||
|
||||
match result {
|
||||
Ok(leased) => {
|
||||
// Transition parent job to processing if still queued.
|
||||
conn.execute(
|
||||
"UPDATE batch_job SET status = 'processing', started_at = ?1
|
||||
WHERE batch_id = ?2 AND status = 'queued'",
|
||||
params![now, leased.batch_id.0],
|
||||
)?;
|
||||
Ok(Some(leased))
|
||||
}
|
||||
Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
|
||||
Err(e) => Err(QueueError::from(e)),
|
||||
}
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
async fn complete_item(&self, id: &ItemId, result: BatchItemResult) -> Result<(), QueueError> {
|
||||
let db = self.db.clone();
|
||||
let id = id.0.clone();
|
||||
let result_body =
|
||||
serde_json::to_string(&result.body).map_err(|e| QueueError::Storage(e.to_string()))?;
|
||||
let status_code = result.status_code;
|
||||
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let conn = db.blocking_lock();
|
||||
conn.execute(
|
||||
"UPDATE batch_item SET status = 'succeeded', result_status = ?1,
|
||||
result_body = ?2, lease_id = NULL, lease_expires_at = NULL,
|
||||
completed_at = ?3
|
||||
WHERE item_id = ?4",
|
||||
params![status_code, result_body, now_iso8601(), id],
|
||||
)?;
|
||||
Ok(())
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
async fn fail_item(&self, id: &ItemId, error: &str) -> Result<(), QueueError> {
|
||||
let db = self.db.clone();
|
||||
let id = id.0.clone();
|
||||
let error = error.to_string();
|
||||
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let conn = db.blocking_lock();
|
||||
conn.execute(
|
||||
"UPDATE batch_item SET status = 'failed', last_error = ?1,
|
||||
lease_id = NULL, lease_expires_at = NULL, completed_at = ?2
|
||||
WHERE item_id = ?3",
|
||||
params![error, now_iso8601(), id],
|
||||
)?;
|
||||
Ok(())
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
async fn schedule_retry(
|
||||
&self,
|
||||
id: &ItemId,
|
||||
delay: Duration,
|
||||
error: &str,
|
||||
) -> Result<(), QueueError> {
|
||||
let db = self.db.clone();
|
||||
let id = id.0.clone();
|
||||
let error = error.to_string();
|
||||
let retry_at = {
|
||||
let secs = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_secs()
|
||||
+ delay.as_secs();
|
||||
format_epoch_iso8601(secs)
|
||||
};
|
||||
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let conn = db.blocking_lock();
|
||||
conn.execute(
|
||||
"UPDATE batch_item SET status = 'pending', last_error = ?1,
|
||||
next_retry_at = ?2, lease_id = NULL, lease_expires_at = NULL
|
||||
WHERE item_id = ?3",
|
||||
params![error, retry_at, id],
|
||||
)?;
|
||||
Ok(())
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
async fn dead_letter(&self, id: &ItemId) -> Result<(), QueueError> {
|
||||
let db = self.db.clone();
|
||||
let id = id.0.clone();
|
||||
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let conn = db.blocking_lock();
|
||||
conn.execute(
|
||||
"INSERT OR IGNORE INTO batch_dead_letter (item_id, batch_id, custom_id, request_body, last_error, attempts, failed_at)
|
||||
SELECT item_id, batch_id, custom_id, request_body, last_error, attempts, ?1
|
||||
FROM batch_item WHERE item_id = ?2",
|
||||
params![now_iso8601(), id],
|
||||
)?;
|
||||
Ok(())
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
async fn is_batch_complete(&self, id: &BatchId) -> Result<bool, QueueError> {
|
||||
let db = self.db.clone();
|
||||
let id = id.0.clone();
|
||||
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let conn = db.blocking_lock();
|
||||
let count: i64 = conn.query_row(
|
||||
"SELECT COUNT(*) FROM batch_item
|
||||
WHERE batch_id = ?1 AND status NOT IN ('succeeded', 'failed', 'cancelled')",
|
||||
params![id],
|
||||
|row| row.get(0),
|
||||
)?;
|
||||
Ok(count == 0)
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
async fn complete_batch(&self, id: &BatchId) -> Result<(), QueueError> {
|
||||
let db = self.db.clone();
|
||||
let id = id.0.clone();
|
||||
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let conn = db.blocking_lock();
|
||||
// Count final item states.
|
||||
let (succeeded, failed, cancelled): (i64, i64, i64) = conn.query_row(
|
||||
"SELECT
|
||||
COUNT(CASE WHEN status = 'succeeded' THEN 1 END),
|
||||
COUNT(CASE WHEN status = 'failed' THEN 1 END),
|
||||
COUNT(CASE WHEN status = 'cancelled' THEN 1 END)
|
||||
FROM batch_item WHERE batch_id = ?1",
|
||||
params![id],
|
||||
|row| Ok((row.get(0)?, row.get(1)?, row.get(2)?)),
|
||||
)?;
|
||||
|
||||
conn.execute(
|
||||
"UPDATE batch_job SET status = 'completed',
|
||||
succeeded = ?1, failed = ?2, cancelled = ?3,
|
||||
processing = 0, completed_at = ?4
|
||||
WHERE batch_id = ?5",
|
||||
params![succeeded, failed, cancelled, now_iso8601(), id],
|
||||
)?;
|
||||
Ok(())
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
async fn get_native_jobs_in_progress(&self) -> Result<Vec<BatchJob>, QueueError> {
|
||||
let db = self.db.clone();
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let conn = db.blocking_lock();
|
||||
let mut stmt = conn.prepare(
|
||||
"SELECT batch_id, status, execution_mode, provider, priority,
|
||||
key_id, input_file_id, webhook_url, metadata,
|
||||
total, processing, succeeded, failed, cancelled, expired,
|
||||
created_at, started_at, completed_at, expires_at
|
||||
FROM batch_job
|
||||
WHERE execution_mode = 'native' AND status = 'processing'",
|
||||
)?;
|
||||
let rows = stmt.query_map([], |row| Ok(batch_job_from_row(row)))?;
|
||||
rows.collect::<Result<Vec<_>, _>>()
|
||||
.map_err(QueueError::from)
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
async fn reclaim_expired_leases(&self) -> Result<u32, QueueError> {
|
||||
let db = self.db.clone();
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let conn = db.blocking_lock();
|
||||
let now = now_iso8601();
|
||||
let count = conn.execute(
|
||||
"UPDATE batch_item SET status = 'pending', lease_id = NULL, lease_expires_at = NULL
|
||||
WHERE lease_id IS NOT NULL AND lease_expires_at < ?1 AND status = 'processing'",
|
||||
params![now],
|
||||
)?;
|
||||
if count > 0 {
|
||||
tracing::warn!(count, "reclaimed expired item leases");
|
||||
}
|
||||
Ok(count as u32)
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
async fn update_progress(
|
||||
&self,
|
||||
id: &BatchId,
|
||||
counts: &RequestCounts,
|
||||
) -> Result<(), QueueError> {
|
||||
let db = self.db.clone();
|
||||
let id = id.0.clone();
|
||||
let counts = counts.clone();
|
||||
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let conn = db.blocking_lock();
|
||||
conn.execute(
|
||||
"UPDATE batch_job SET
|
||||
processing = ?1, succeeded = ?2, failed = ?3,
|
||||
cancelled = ?4, expired = ?5
|
||||
WHERE batch_id = ?6",
|
||||
params![
|
||||
counts.processing,
|
||||
counts.succeeded,
|
||||
counts.failed,
|
||||
counts.cancelled,
|
||||
counts.expired,
|
||||
id,
|
||||
],
|
||||
)?;
|
||||
Ok(())
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
async fn get_items(&self, batch_id: &BatchId) -> Result<Vec<BatchItem>, QueueError> {
|
||||
let db = self.db.clone();
|
||||
let batch_id = batch_id.0.clone();
|
||||
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let conn = db.blocking_lock();
|
||||
let mut stmt = conn.prepare(
|
||||
"SELECT item_id, batch_id, custom_id, status, model, request_body,
|
||||
source_format, result_status, result_body, attempts,
|
||||
max_retries, last_error, next_retry_at, lease_id,
|
||||
lease_expires_at, idempotency_key, created_at, completed_at
|
||||
FROM batch_item WHERE batch_id = ?1
|
||||
ORDER BY created_at ASC",
|
||||
)?;
|
||||
let rows = stmt.query_map(params![batch_id], |row| Ok(batch_item_from_row(row)))?;
|
||||
rows.collect::<Result<Vec<_>, _>>()
|
||||
.map_err(QueueError::from)
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
}
|
||||
|
||||
// -- Row mappers --
|
||||
|
||||
fn batch_job_from_row(row: &rusqlite::Row) -> BatchJob {
|
||||
let status_str: String = row.get(1).unwrap_or_default();
|
||||
let exec_mode_str: String = row.get(2).unwrap_or_default();
|
||||
let provider: Option<String> = row.get(3).unwrap_or(None);
|
||||
let metadata_str: Option<String> = row.get(8).unwrap_or(None);
|
||||
|
||||
let execution_mode = match exec_mode_str.as_str() {
|
||||
"native" => ExecutionMode::Native {
|
||||
provider: provider.unwrap_or_else(|| "unknown".into()),
|
||||
},
|
||||
_ => ExecutionMode::ProxyNative,
|
||||
};
|
||||
|
||||
BatchJob {
|
||||
id: BatchId(row.get(0).unwrap_or_default()),
|
||||
status: BatchStatus::from_str_status(&status_str),
|
||||
execution_mode,
|
||||
priority: row.get::<_, i64>(4).unwrap_or(0) as u8,
|
||||
key_id: row.get(5).unwrap_or(None),
|
||||
input_file_id: row.get(6).unwrap_or_default(),
|
||||
webhook_url: row.get(7).unwrap_or(None),
|
||||
metadata: metadata_str.and_then(|s| serde_json::from_str(&s).ok()),
|
||||
request_counts: RequestCounts {
|
||||
total: row.get::<_, i64>(9).unwrap_or(0) as u32,
|
||||
processing: row.get::<_, i64>(10).unwrap_or(0) as u32,
|
||||
succeeded: row.get::<_, i64>(11).unwrap_or(0) as u32,
|
||||
failed: row.get::<_, i64>(12).unwrap_or(0) as u32,
|
||||
cancelled: row.get::<_, i64>(13).unwrap_or(0) as u32,
|
||||
expired: row.get::<_, i64>(14).unwrap_or(0) as u32,
|
||||
},
|
||||
created_at: row.get(15).unwrap_or_default(),
|
||||
started_at: row.get(16).unwrap_or(None),
|
||||
completed_at: row.get(17).unwrap_or(None),
|
||||
expires_at: row.get(18).unwrap_or_default(),
|
||||
}
|
||||
}
|
||||
|
||||
fn batch_item_from_row(row: &rusqlite::Row) -> BatchItem {
|
||||
let status_str: String = row.get(3).unwrap_or_default();
|
||||
let model: String = row.get(4).unwrap_or_default();
|
||||
let body_str: String = row.get(5).unwrap_or_default();
|
||||
let source_fmt_str: String = row.get(6).unwrap_or_default();
|
||||
let result_status: Option<i64> = row.get(7).unwrap_or(None);
|
||||
let result_body_str: Option<String> = row.get(8).unwrap_or(None);
|
||||
|
||||
let source_format =
|
||||
serde_json::from_str::<SourceFormat>(&source_fmt_str).unwrap_or(SourceFormat::OpenAI);
|
||||
|
||||
let body = serde_json::from_str(&body_str).unwrap_or(serde_json::Value::Null);
|
||||
|
||||
let result = match (result_status, result_body_str) {
|
||||
(Some(code), Some(body_s)) => {
|
||||
let body_val = serde_json::from_str(&body_s).unwrap_or(serde_json::Value::Null);
|
||||
Some(BatchItemResult {
|
||||
status_code: code as u16,
|
||||
body: body_val,
|
||||
})
|
||||
}
|
||||
_ => None,
|
||||
};
|
||||
|
||||
BatchItem {
|
||||
id: ItemId(row.get(0).unwrap_or_default()),
|
||||
batch_id: BatchId(row.get(1).unwrap_or_default()),
|
||||
custom_id: row.get(2).unwrap_or_default(),
|
||||
status: ItemStatus::from_str_status(&status_str),
|
||||
request: BatchItemRequest {
|
||||
model,
|
||||
body,
|
||||
source_format,
|
||||
},
|
||||
result,
|
||||
attempts: row.get::<_, i64>(9).unwrap_or(0) as u8,
|
||||
max_retries: row.get::<_, i64>(10).unwrap_or(3) as u8,
|
||||
last_error: row.get(11).unwrap_or(None),
|
||||
next_retry_at: row.get(12).unwrap_or(None),
|
||||
lease_id: row.get(13).unwrap_or(None),
|
||||
lease_expires_at: row.get(14).unwrap_or(None),
|
||||
idempotency_key: row.get(15).unwrap_or(None),
|
||||
created_at: row.get(16).unwrap_or_default(),
|
||||
completed_at: row.get(17).unwrap_or(None),
|
||||
}
|
||||
}
|
||||
|
||||
fn row_to_job(conn: &Connection, batch_id: &str) -> Result<Option<BatchJob>, QueueError> {
|
||||
let mut stmt = conn.prepare(
|
||||
"SELECT batch_id, status, execution_mode, provider, priority,
|
||||
key_id, input_file_id, webhook_url, metadata,
|
||||
total, processing, succeeded, failed, cancelled, expired,
|
||||
created_at, started_at, completed_at, expires_at
|
||||
FROM batch_job WHERE batch_id = ?1",
|
||||
)?;
|
||||
let mut rows = stmt.query(params![batch_id])?;
|
||||
if let Some(row) = rows.next()? {
|
||||
Ok(Some(batch_job_from_row(row)))
|
||||
} else {
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert epoch seconds to ISO 8601 string.
|
||||
pub(crate) fn format_epoch_iso8601(secs: u64) -> String {
|
||||
let days = secs / 86400;
|
||||
let day_secs = secs % 86400;
|
||||
let h = day_secs / 3600;
|
||||
let m = (day_secs % 3600) / 60;
|
||||
let s = day_secs % 60;
|
||||
|
||||
let z = days as i64 + 719468;
|
||||
let era = if z >= 0 { z } else { z - 146096 } / 146097;
|
||||
let doe = (z - era * 146097) as u64;
|
||||
let yoe = (doe - doe / 1460 + doe / 36524 - doe / 146096) / 365;
|
||||
let y = yoe as i64 + era * 400;
|
||||
let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
|
||||
let mp = (5 * doy + 2) / 153;
|
||||
let d = doy - (153 * mp + 2) / 5 + 1;
|
||||
let m_val = if mp < 10 { mp + 3 } else { mp - 9 };
|
||||
let y_val = if m_val <= 2 { y + 1 } else { y };
|
||||
|
||||
format!(
|
||||
"{:04}-{:02}-{:02}T{:02}:{:02}:{:02}Z",
|
||||
y_val, m_val, d, h, m, s
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::db::init_batch_engine_tables;
|
||||
|
||||
async fn test_queue() -> SqliteQueue {
|
||||
let conn = Connection::open_in_memory().unwrap();
|
||||
init_batch_engine_tables(&conn).unwrap();
|
||||
SqliteQueue::new(Arc::new(Mutex::new(conn)))
|
||||
}
|
||||
|
||||
fn make_job(id: &str) -> (BatchJob, Vec<BatchItem>) {
|
||||
let batch_id = BatchId(id.into());
|
||||
let now = crate::db::now_iso8601();
|
||||
let job = BatchJob {
|
||||
id: batch_id.clone(),
|
||||
status: BatchStatus::Queued,
|
||||
execution_mode: ExecutionMode::ProxyNative,
|
||||
priority: 0,
|
||||
key_id: None,
|
||||
input_file_id: "file-test".into(),
|
||||
webhook_url: None,
|
||||
metadata: None,
|
||||
request_counts: RequestCounts {
|
||||
total: 2,
|
||||
..Default::default()
|
||||
},
|
||||
created_at: now.clone(),
|
||||
started_at: None,
|
||||
completed_at: None,
|
||||
expires_at: now.clone(),
|
||||
};
|
||||
let items = vec![
|
||||
BatchItem {
|
||||
id: ItemId(format!("{id}_item_1")),
|
||||
batch_id: batch_id.clone(),
|
||||
custom_id: "req-1".into(),
|
||||
status: ItemStatus::Pending,
|
||||
request: BatchItemRequest {
|
||||
model: "gpt-4o".into(),
|
||||
body: serde_json::json!({"messages": []}),
|
||||
source_format: SourceFormat::OpenAI,
|
||||
},
|
||||
result: None,
|
||||
attempts: 0,
|
||||
max_retries: 3,
|
||||
last_error: None,
|
||||
next_retry_at: None,
|
||||
lease_id: None,
|
||||
lease_expires_at: None,
|
||||
idempotency_key: None,
|
||||
created_at: now.clone(),
|
||||
completed_at: None,
|
||||
},
|
||||
BatchItem {
|
||||
id: ItemId(format!("{id}_item_2")),
|
||||
batch_id: batch_id.clone(),
|
||||
custom_id: "req-2".into(),
|
||||
status: ItemStatus::Pending,
|
||||
request: BatchItemRequest {
|
||||
model: "gpt-4o".into(),
|
||||
body: serde_json::json!({"messages": []}),
|
||||
source_format: SourceFormat::OpenAI,
|
||||
},
|
||||
result: None,
|
||||
attempts: 0,
|
||||
max_retries: 3,
|
||||
last_error: None,
|
||||
next_retry_at: None,
|
||||
lease_id: None,
|
||||
lease_expires_at: None,
|
||||
idempotency_key: None,
|
||||
created_at: now.clone(),
|
||||
completed_at: None,
|
||||
},
|
||||
];
|
||||
(job, items)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn enqueue_and_get() {
|
||||
let q = test_queue().await;
|
||||
let (job, items) = make_job("batch_test1");
|
||||
q.enqueue(&job, &items).await.unwrap();
|
||||
|
||||
let fetched = q.get(&BatchId("batch_test1".into())).await.unwrap();
|
||||
assert!(fetched.is_some());
|
||||
let fetched = fetched.unwrap();
|
||||
assert_eq!(fetched.status, BatchStatus::Queued);
|
||||
assert_eq!(fetched.request_counts.total, 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn get_nonexistent() {
|
||||
let q = test_queue().await;
|
||||
let fetched = q.get(&BatchId("nope".into())).await.unwrap();
|
||||
assert!(fetched.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cancel_queued_job() {
|
||||
let q = test_queue().await;
|
||||
let (job, items) = make_job("batch_cancel");
|
||||
q.enqueue(&job, &items).await.unwrap();
|
||||
|
||||
let status = q.cancel(&BatchId("batch_cancel".into())).await.unwrap();
|
||||
assert_eq!(status, BatchStatus::Cancelled);
|
||||
|
||||
let fetched = q
|
||||
.get(&BatchId("batch_cancel".into()))
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(fetched.status, BatchStatus::Cancelled);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn list_with_pagination() {
|
||||
let q = test_queue().await;
|
||||
for i in 0..5 {
|
||||
let (job, items) = make_job(&format!("batch_list_{i}"));
|
||||
q.enqueue(&job, &items).await.unwrap();
|
||||
}
|
||||
|
||||
let all = q.list(None, None, 10).await.unwrap();
|
||||
assert_eq!(all.len(), 5);
|
||||
|
||||
let page = q.list(None, None, 2).await.unwrap();
|
||||
assert_eq!(page.len(), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn get_items() {
|
||||
let q = test_queue().await;
|
||||
let (job, items) = make_job("batch_items");
|
||||
q.enqueue(&job, &items).await.unwrap();
|
||||
|
||||
let fetched = q.get_items(&BatchId("batch_items".into())).await.unwrap();
|
||||
assert_eq!(fetched.len(), 2);
|
||||
assert_eq!(fetched[0].custom_id, "req-1");
|
||||
assert_eq!(fetched[1].custom_id, "req-2");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn complete_item_and_batch() {
|
||||
let q = test_queue().await;
|
||||
let (job, items) = make_job("batch_complete");
|
||||
q.enqueue(&job, &items).await.unwrap();
|
||||
|
||||
// Complete both items.
|
||||
let result = BatchItemResult {
|
||||
status_code: 200,
|
||||
body: serde_json::json!({"id": "resp-1"}),
|
||||
};
|
||||
q.complete_item(&ItemId("batch_complete_item_1".into()), result.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
q.complete_item(&ItemId("batch_complete_item_2".into()), result)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(q
|
||||
.is_batch_complete(&BatchId("batch_complete".into()))
|
||||
.await
|
||||
.unwrap());
|
||||
|
||||
q.complete_batch(&BatchId("batch_complete".into()))
|
||||
.await
|
||||
.unwrap();
|
||||
let job = q
|
||||
.get(&BatchId("batch_complete".into()))
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(job.status, BatchStatus::Completed);
|
||||
assert_eq!(job.request_counts.succeeded, 2);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,187 @@
|
||||
// crates/batch_engine/src/validation.rs
|
||||
//! JSONL batch file validation.
|
||||
|
||||
use std::collections::HashSet;
|
||||
|
||||
/// Maximum number of lines in a JSONL batch file.
|
||||
const MAX_LINE_COUNT: usize = 50_000;
|
||||
|
||||
/// Maximum file size in bytes (100 MB).
|
||||
const MAX_FILE_SIZE: usize = 100 * 1024 * 1024;
|
||||
|
||||
/// Maximum length of a custom_id field.
|
||||
const MAX_CUSTOM_ID_LEN: usize = 64;
|
||||
|
||||
/// Result of JSONL validation: line count on success.
|
||||
#[derive(Debug)]
|
||||
pub struct ValidatedJsonl {
|
||||
pub line_count: usize,
|
||||
}
|
||||
|
||||
/// Validate a JSONL batch file.
|
||||
///
|
||||
/// Each line must be valid JSON with a unique `custom_id` (string, max 64 chars)
|
||||
/// and a `body` object containing a `model` field. Max 50,000 lines, 100 MB.
|
||||
pub fn validate_jsonl(mut reader: impl std::io::BufRead) -> Result<ValidatedJsonl, String> {
|
||||
let mut seen_ids = HashSet::new();
|
||||
let mut line_count = 0usize;
|
||||
let mut raw_line_num = 0usize;
|
||||
let mut bytes_read = 0usize;
|
||||
let mut line_buf = String::new();
|
||||
|
||||
loop {
|
||||
line_buf.clear();
|
||||
let n = reader
|
||||
.read_line(&mut line_buf)
|
||||
.map_err(|e| format!("Read error: {e}"))?;
|
||||
if n == 0 {
|
||||
break;
|
||||
}
|
||||
raw_line_num += 1;
|
||||
bytes_read += n;
|
||||
if bytes_read > MAX_FILE_SIZE {
|
||||
return Err(format!(
|
||||
"File exceeds maximum size of {} bytes",
|
||||
MAX_FILE_SIZE
|
||||
));
|
||||
}
|
||||
|
||||
let line = line_buf.trim();
|
||||
if line.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
line_count += 1;
|
||||
if line_count > MAX_LINE_COUNT {
|
||||
return Err(format!("File exceeds maximum of {MAX_LINE_COUNT} lines"));
|
||||
}
|
||||
|
||||
let parsed: serde_json::Value = serde_json::from_str(line)
|
||||
.map_err(|e| format!("Line {raw_line_num}: invalid JSON: {e}"))?;
|
||||
|
||||
let obj = parsed
|
||||
.as_object()
|
||||
.ok_or_else(|| format!("Line {raw_line_num}: expected JSON object"))?;
|
||||
|
||||
let custom_id = obj
|
||||
.get("custom_id")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| format!("Line {raw_line_num}: missing or non-string 'custom_id'"))?;
|
||||
|
||||
if custom_id.len() > MAX_CUSTOM_ID_LEN {
|
||||
return Err(format!(
|
||||
"Line {raw_line_num}: custom_id exceeds maximum length of {MAX_CUSTOM_ID_LEN} characters"
|
||||
));
|
||||
}
|
||||
|
||||
if !seen_ids.insert(custom_id.to_string()) {
|
||||
return Err(format!(
|
||||
"Line {raw_line_num}: duplicate custom_id '{custom_id}'"
|
||||
));
|
||||
}
|
||||
|
||||
let body = obj
|
||||
.get("body")
|
||||
.and_then(|v| v.as_object())
|
||||
.ok_or_else(|| format!("Line {raw_line_num}: missing or non-object 'body'"))?;
|
||||
|
||||
if !body.contains_key("model") {
|
||||
return Err(format!("Line {raw_line_num}: body missing 'model' field"));
|
||||
}
|
||||
}
|
||||
|
||||
if line_count == 0 {
|
||||
return Err("File is empty".to_string());
|
||||
}
|
||||
|
||||
Ok(ValidatedJsonl { line_count })
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::io::{BufReader, Cursor};
|
||||
|
||||
fn check(data: &str) -> Result<ValidatedJsonl, String> {
|
||||
validate_jsonl(BufReader::new(Cursor::new(data.as_bytes())))
|
||||
}
|
||||
|
||||
fn check_bytes(data: &[u8]) -> Result<ValidatedJsonl, String> {
|
||||
validate_jsonl(BufReader::new(Cursor::new(data)))
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn valid_jsonl() {
|
||||
let data = r#"{"custom_id": "req-1", "body": {"model": "gpt-4o", "messages": []}}
|
||||
{"custom_id": "req-2", "body": {"model": "gpt-4o", "messages": []}}"#;
|
||||
let result = check(data);
|
||||
assert!(result.is_ok());
|
||||
assert_eq!(result.unwrap().line_count, 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_custom_id() {
|
||||
let data = r#"{"body": {"model": "gpt-4o"}}"#;
|
||||
let result = check(data);
|
||||
assert!(result.is_err());
|
||||
assert!(result.unwrap_err().contains("custom_id"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_body_model() {
|
||||
let data = r#"{"custom_id": "req-1", "body": {"messages": []}}"#;
|
||||
let result = check(data);
|
||||
assert!(result.is_err());
|
||||
assert!(result.unwrap_err().contains("model"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn duplicate_custom_id() {
|
||||
let data = r#"{"custom_id": "req-1", "body": {"model": "gpt-4o"}}
|
||||
{"custom_id": "req-1", "body": {"model": "gpt-4o"}}"#;
|
||||
let result = check(data);
|
||||
assert!(result.is_err());
|
||||
assert!(result.unwrap_err().contains("duplicate"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn oversized_custom_id() {
|
||||
let long_id = "a".repeat(65);
|
||||
let data = format!(r#"{{"custom_id": "{long_id}", "body": {{"model": "gpt-4o"}}}}"#);
|
||||
let result = check(&data);
|
||||
assert!(result.is_err());
|
||||
assert!(result.unwrap_err().contains("maximum length"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_file() {
|
||||
let result = check_bytes(b"");
|
||||
assert!(result.is_err());
|
||||
assert!(result.unwrap_err().contains("empty"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_json_line() {
|
||||
let data = b"not json at all";
|
||||
let result = check_bytes(data);
|
||||
assert!(result.is_err());
|
||||
assert!(result.unwrap_err().contains("invalid JSON"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn blank_lines_skipped() {
|
||||
let data = r#"{"custom_id": "req-1", "body": {"model": "gpt-4o"}}
|
||||
|
||||
{"custom_id": "req-2", "body": {"model": "gpt-4o"}}"#;
|
||||
let result = check(data);
|
||||
assert!(result.is_ok());
|
||||
assert_eq!(result.unwrap().line_count, 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn error_reports_absolute_line_number_with_blank_lines() {
|
||||
let data = "\n{\"custom_id\": \"ok\", \"body\": INVALID}";
|
||||
let err = check(data).unwrap_err();
|
||||
assert!(err.contains("Line 2"), "expected 'Line 2' in: {err}");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,145 @@
|
||||
// crates/batch_engine/src/webhook/dispatcher.rs
|
||||
//! Background webhook delivery loop with HMAC signing and retries.
|
||||
|
||||
use super::WebhookQueue;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use tokio::task::JoinHandle;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
/// Configuration for the webhook dispatcher.
|
||||
pub struct WebhookConfig {
|
||||
pub poll_interval: Duration,
|
||||
pub reclaim_interval: Duration,
|
||||
pub max_concurrent: usize,
|
||||
}
|
||||
|
||||
impl Default for WebhookConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
poll_interval: Duration::from_secs(1),
|
||||
reclaim_interval: Duration::from_secs(30),
|
||||
max_concurrent: 8,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Handle to the running webhook dispatcher.
|
||||
pub struct WebhookHandle {
|
||||
shutdown: CancellationToken,
|
||||
join_handle: JoinHandle<()>,
|
||||
}
|
||||
|
||||
impl WebhookHandle {
|
||||
pub async fn shutdown(self) {
|
||||
self.shutdown.cancel();
|
||||
let _ = self.join_handle.await;
|
||||
}
|
||||
}
|
||||
|
||||
/// Start the webhook dispatcher background loop.
|
||||
pub fn start_dispatcher<Q: WebhookQueue>(
|
||||
queue: Arc<Q>,
|
||||
client: reqwest::Client,
|
||||
config: WebhookConfig,
|
||||
) -> WebhookHandle {
|
||||
let shutdown = CancellationToken::new();
|
||||
let token = shutdown.clone();
|
||||
|
||||
let join_handle = tokio::spawn(async move {
|
||||
let semaphore = Arc::new(tokio::sync::Semaphore::new(config.max_concurrent));
|
||||
let mut poll_interval = tokio::time::interval(config.poll_interval);
|
||||
let mut reclaim_interval = tokio::time::interval(config.reclaim_interval);
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = token.cancelled() => break,
|
||||
_ = reclaim_interval.tick() => {
|
||||
if let Ok(count) = queue.reclaim_expired_leases().await {
|
||||
if count > 0 {
|
||||
tracing::warn!(count, "reclaimed expired webhook leases");
|
||||
}
|
||||
}
|
||||
}
|
||||
_ = poll_interval.tick() => {
|
||||
let Ok(permit) = semaphore.clone().try_acquire_owned() else {
|
||||
continue;
|
||||
};
|
||||
|
||||
match queue.claim_next().await {
|
||||
Ok(Some(leased)) => {
|
||||
let queue = queue.clone();
|
||||
let client = client.clone();
|
||||
tokio::spawn(async move {
|
||||
deliver(queue.as_ref(), &client, &leased.delivery).await;
|
||||
drop(permit);
|
||||
});
|
||||
}
|
||||
Ok(None) => {
|
||||
drop(permit);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::error!(error = %e, "webhook queue claim error");
|
||||
drop(permit);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
tracing::info!("webhook dispatcher shut down");
|
||||
});
|
||||
|
||||
WebhookHandle {
|
||||
shutdown,
|
||||
join_handle,
|
||||
}
|
||||
}
|
||||
|
||||
async fn deliver<Q: WebhookQueue>(
|
||||
queue: &Q,
|
||||
client: &reqwest::Client,
|
||||
delivery: &super::WebhookDelivery,
|
||||
) {
|
||||
let mut request = client
|
||||
.post(&delivery.url)
|
||||
.header("Content-Type", "application/json")
|
||||
.header("X-Webhook-Id", &delivery.event_id);
|
||||
|
||||
// HMAC signing.
|
||||
if let Some(ref secret) = delivery.signing_secret {
|
||||
use hmac::{Hmac, Mac};
|
||||
use sha2::Sha256;
|
||||
let payload_bytes = serde_json::to_vec(&delivery.payload).unwrap_or_default();
|
||||
let mut mac =
|
||||
Hmac::<Sha256>::new_from_slice(secret.as_bytes()).expect("HMAC key length ok");
|
||||
mac.update(&payload_bytes);
|
||||
let sig = hex::encode(mac.finalize().into_bytes());
|
||||
request = request.header("X-Webhook-Signature", format!("sha256={sig}"));
|
||||
}
|
||||
|
||||
let response = request.json(&delivery.payload).send().await;
|
||||
|
||||
match response {
|
||||
Ok(r) if r.status().is_success() => {
|
||||
if let Err(e) = queue.ack(&delivery.delivery_id).await {
|
||||
tracing::error!(error = %e, "failed to ack webhook delivery");
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
if delivery.attempts < delivery.max_retries {
|
||||
let delay = Duration::from_secs(1 << delivery.attempts.min(4));
|
||||
if let Err(e) = queue.schedule_retry(&delivery.delivery_id, delay).await {
|
||||
tracing::error!(error = %e, "failed to schedule webhook retry");
|
||||
}
|
||||
} else {
|
||||
tracing::warn!(
|
||||
delivery_id = %delivery.delivery_id,
|
||||
"webhook delivery exhausted retries, moving to dead letter"
|
||||
);
|
||||
if let Err(e) = queue.dead_letter(&delivery.delivery_id).await {
|
||||
tracing::error!(error = %e, "failed to dead-letter webhook");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,43 @@
|
||||
// crates/batch_engine/src/webhook/mod.rs
|
||||
//! Durable webhook delivery queue and dispatcher.
|
||||
|
||||
pub mod dispatcher;
|
||||
pub mod sqlite;
|
||||
|
||||
use crate::error::QueueError;
|
||||
use async_trait::async_trait;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::time::Duration;
|
||||
|
||||
/// A webhook delivery request.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct WebhookDelivery {
|
||||
pub delivery_id: String,
|
||||
pub event_id: String,
|
||||
pub batch_id: String,
|
||||
pub url: String,
|
||||
pub payload: serde_json::Value,
|
||||
#[serde(skip)]
|
||||
pub signing_secret: Option<String>,
|
||||
pub attempts: u8,
|
||||
pub max_retries: u8,
|
||||
pub next_retry_at: Option<String>,
|
||||
}
|
||||
|
||||
/// A claimed webhook delivery with lease info.
|
||||
#[derive(Debug)]
|
||||
pub struct LeasedDelivery {
|
||||
pub delivery: WebhookDelivery,
|
||||
pub lease_id: String,
|
||||
}
|
||||
|
||||
/// Durable webhook delivery queue.
|
||||
#[async_trait]
|
||||
pub trait WebhookQueue: Send + Sync + 'static {
|
||||
async fn enqueue(&self, delivery: WebhookDelivery) -> Result<(), QueueError>;
|
||||
async fn claim_next(&self) -> Result<Option<LeasedDelivery>, QueueError>;
|
||||
async fn ack(&self, delivery_id: &str) -> Result<(), QueueError>;
|
||||
async fn schedule_retry(&self, delivery_id: &str, delay: Duration) -> Result<(), QueueError>;
|
||||
async fn dead_letter(&self, delivery_id: &str) -> Result<(), QueueError>;
|
||||
async fn reclaim_expired_leases(&self) -> Result<u32, QueueError>;
|
||||
}
|
||||
@@ -0,0 +1,263 @@
|
||||
// crates/batch_engine/src/webhook/sqlite.rs
|
||||
//! SQLite-backed webhook delivery queue.
|
||||
|
||||
use super::{LeasedDelivery, WebhookDelivery, WebhookQueue};
|
||||
use crate::db::now_iso8601;
|
||||
use crate::error::QueueError;
|
||||
use async_trait::async_trait;
|
||||
use rusqlite::{params, Connection};
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
/// SQLite-backed webhook delivery queue.
|
||||
#[derive(Clone)]
|
||||
pub struct SqliteWebhookQueue {
|
||||
db: Arc<Mutex<Connection>>,
|
||||
}
|
||||
|
||||
impl SqliteWebhookQueue {
|
||||
pub fn new(db: Arc<Mutex<Connection>>) -> Self {
|
||||
Self { db }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl WebhookQueue for SqliteWebhookQueue {
|
||||
async fn enqueue(&self, delivery: WebhookDelivery) -> Result<(), QueueError> {
|
||||
let db = self.db.clone();
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let conn = db.blocking_lock();
|
||||
let payload_str = serde_json::to_string(&delivery.payload)
|
||||
.map_err(|e| QueueError::Storage(e.to_string()))?;
|
||||
conn.execute(
|
||||
"INSERT INTO webhook_delivery
|
||||
(delivery_id, event_id, batch_id, url, payload, signing_secret,
|
||||
status, attempts, max_retries, created_at)
|
||||
VALUES (?1, ?2, ?3, ?4, ?5, ?6, 'pending', ?7, ?8, ?9)",
|
||||
params![
|
||||
delivery.delivery_id,
|
||||
delivery.event_id,
|
||||
delivery.batch_id,
|
||||
delivery.url,
|
||||
payload_str,
|
||||
delivery.signing_secret,
|
||||
delivery.attempts,
|
||||
delivery.max_retries,
|
||||
now_iso8601(),
|
||||
],
|
||||
)?;
|
||||
Ok(())
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
async fn claim_next(&self) -> Result<Option<LeasedDelivery>, QueueError> {
|
||||
let db = self.db.clone();
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let conn = db.blocking_lock();
|
||||
let lease_id = format!("whl_{}", uuid::Uuid::new_v4());
|
||||
let now = now_iso8601();
|
||||
let lease_expires = {
|
||||
let secs = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_secs()
|
||||
+ 60;
|
||||
super::super::queue::sqlite::format_epoch_iso8601(secs)
|
||||
};
|
||||
|
||||
let result = conn.query_row(
|
||||
"UPDATE webhook_delivery
|
||||
SET status = 'processing', lease_id = ?1, lease_expires_at = ?2,
|
||||
attempts = attempts + 1
|
||||
WHERE delivery_id = (
|
||||
SELECT delivery_id FROM webhook_delivery
|
||||
WHERE status = 'pending'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= ?3)
|
||||
ORDER BY created_at ASC
|
||||
LIMIT 1
|
||||
)
|
||||
RETURNING delivery_id, event_id, batch_id, url, payload,
|
||||
signing_secret, attempts, max_retries, next_retry_at",
|
||||
params![lease_id, lease_expires, now],
|
||||
|row| {
|
||||
let payload_str: String = row.get(4)?;
|
||||
let payload =
|
||||
serde_json::from_str(&payload_str).unwrap_or(serde_json::Value::Null);
|
||||
Ok(LeasedDelivery {
|
||||
delivery: WebhookDelivery {
|
||||
delivery_id: row.get(0)?,
|
||||
event_id: row.get(1)?,
|
||||
batch_id: row.get(2)?,
|
||||
url: row.get(3)?,
|
||||
payload,
|
||||
signing_secret: row.get(5)?,
|
||||
attempts: row.get::<_, i64>(6).unwrap_or(0) as u8,
|
||||
max_retries: row.get::<_, i64>(7).unwrap_or(3) as u8,
|
||||
next_retry_at: row.get(8)?,
|
||||
},
|
||||
lease_id: lease_id.clone(),
|
||||
})
|
||||
},
|
||||
);
|
||||
|
||||
match result {
|
||||
Ok(leased) => Ok(Some(leased)),
|
||||
Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
|
||||
Err(e) => Err(QueueError::from(e)),
|
||||
}
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
async fn ack(&self, delivery_id: &str) -> Result<(), QueueError> {
|
||||
let db = self.db.clone();
|
||||
let id = delivery_id.to_string();
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let conn = db.blocking_lock();
|
||||
conn.execute(
|
||||
"UPDATE webhook_delivery SET status = 'delivered', delivered_at = ?1,
|
||||
lease_id = NULL, lease_expires_at = NULL
|
||||
WHERE delivery_id = ?2",
|
||||
params![now_iso8601(), id],
|
||||
)?;
|
||||
Ok(())
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
async fn schedule_retry(&self, delivery_id: &str, delay: Duration) -> Result<(), QueueError> {
|
||||
let db = self.db.clone();
|
||||
let id = delivery_id.to_string();
|
||||
let retry_at = {
|
||||
let secs = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_secs()
|
||||
+ delay.as_secs();
|
||||
super::super::queue::sqlite::format_epoch_iso8601(secs)
|
||||
};
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let conn = db.blocking_lock();
|
||||
conn.execute(
|
||||
"UPDATE webhook_delivery SET status = 'pending', next_retry_at = ?1,
|
||||
lease_id = NULL, lease_expires_at = NULL
|
||||
WHERE delivery_id = ?2",
|
||||
params![retry_at, id],
|
||||
)?;
|
||||
Ok(())
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
async fn dead_letter(&self, delivery_id: &str) -> Result<(), QueueError> {
|
||||
let db = self.db.clone();
|
||||
let id = delivery_id.to_string();
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let conn = db.blocking_lock();
|
||||
conn.execute(
|
||||
"UPDATE webhook_delivery SET status = 'dead_letter',
|
||||
lease_id = NULL, lease_expires_at = NULL
|
||||
WHERE delivery_id = ?1",
|
||||
params![id],
|
||||
)?;
|
||||
Ok(())
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
async fn reclaim_expired_leases(&self) -> Result<u32, QueueError> {
|
||||
let db = self.db.clone();
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let conn = db.blocking_lock();
|
||||
let now = now_iso8601();
|
||||
let count = conn.execute(
|
||||
"UPDATE webhook_delivery SET status = 'pending', lease_id = NULL, lease_expires_at = NULL
|
||||
WHERE lease_id IS NOT NULL AND lease_expires_at < ?1 AND status = 'processing'",
|
||||
params![now],
|
||||
)?;
|
||||
Ok(count as u32)
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::db::init_batch_engine_tables;
|
||||
|
||||
async fn test_wq() -> SqliteWebhookQueue {
|
||||
let conn = Connection::open_in_memory().unwrap();
|
||||
init_batch_engine_tables(&conn).unwrap();
|
||||
SqliteWebhookQueue::new(Arc::new(Mutex::new(conn)))
|
||||
}
|
||||
|
||||
fn make_delivery(id: &str) -> WebhookDelivery {
|
||||
WebhookDelivery {
|
||||
delivery_id: id.into(),
|
||||
event_id: format!("evt_{id}"),
|
||||
batch_id: "batch_1".into(),
|
||||
url: "https://example.com/webhook".into(),
|
||||
payload: serde_json::json!({"type": "batch.completed"}),
|
||||
signing_secret: None,
|
||||
attempts: 0,
|
||||
max_retries: 3,
|
||||
next_retry_at: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn enqueue_and_claim() {
|
||||
let wq = test_wq().await;
|
||||
wq.enqueue(make_delivery("whd_1")).await.unwrap();
|
||||
|
||||
let claimed = wq.claim_next().await.unwrap();
|
||||
assert!(claimed.is_some());
|
||||
let claimed = claimed.unwrap();
|
||||
assert_eq!(claimed.delivery.delivery_id, "whd_1");
|
||||
|
||||
// Queue is now empty.
|
||||
assert!(wq.claim_next().await.unwrap().is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ack_delivery() {
|
||||
let wq = test_wq().await;
|
||||
wq.enqueue(make_delivery("whd_ack")).await.unwrap();
|
||||
let claimed = wq.claim_next().await.unwrap().unwrap();
|
||||
wq.ack(&claimed.delivery.delivery_id).await.unwrap();
|
||||
|
||||
// Should not be claimable again.
|
||||
assert!(wq.claim_next().await.unwrap().is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn retry_and_dead_letter() {
|
||||
let wq = test_wq().await;
|
||||
wq.enqueue(make_delivery("whd_retry")).await.unwrap();
|
||||
let claimed = wq.claim_next().await.unwrap().unwrap();
|
||||
|
||||
// Schedule retry with 0 delay (immediate).
|
||||
wq.schedule_retry(&claimed.delivery.delivery_id, Duration::from_secs(0))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Should be claimable again.
|
||||
let claimed2 = wq.claim_next().await.unwrap();
|
||||
assert!(claimed2.is_some());
|
||||
|
||||
// Dead letter.
|
||||
wq.dead_letter(&claimed2.unwrap().delivery.delivery_id)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(wq.claim_next().await.unwrap().is_none());
|
||||
}
|
||||
}
|
||||
@@ -9,6 +9,7 @@ repository.workspace = true
|
||||
[dependencies]
|
||||
anyllm_translate = { path = "../translator", version = "0.2.0" }
|
||||
anyllm_client = { path = "../client", version = "0.2.0" }
|
||||
anyllm_batch_engine = { path = "../batch_engine", version = "0.2.0" }
|
||||
axum = { version = "0.8", features = ["ws", "multipart"] }
|
||||
tokio = { version = "1", features = ["full"] }
|
||||
reqwest = { version = "0.12", default-features = false, features = ["json", "stream", "native-tls", "http2", "multipart"] }
|
||||
|
||||
@@ -77,39 +77,6 @@ pub fn init_db(conn: &Connection) -> rusqlite::Result<()> {
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_audit_log_timestamp ON audit_log(timestamp);
|
||||
|
||||
-- Note: batch file JSONL is stored directly in SQLite. For large batch files
|
||||
-- (>10MB), consider external blob storage. Current design prioritizes simplicity
|
||||
-- and single-binary deployment over storage efficiency.
|
||||
CREATE TABLE IF NOT EXISTS batch_file (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
file_id TEXT NOT NULL UNIQUE,
|
||||
key_id INTEGER,
|
||||
purpose TEXT NOT NULL,
|
||||
filename TEXT,
|
||||
byte_size INTEGER NOT NULL,
|
||||
line_count INTEGER NOT NULL,
|
||||
content BLOB NOT NULL,
|
||||
created_at TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS batch_job (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
batch_id TEXT NOT NULL UNIQUE,
|
||||
key_id INTEGER,
|
||||
input_file_id TEXT NOT NULL,
|
||||
backend_batch_id TEXT,
|
||||
backend_name TEXT NOT NULL,
|
||||
status TEXT NOT NULL,
|
||||
request_counts_total INTEGER NOT NULL DEFAULT 0,
|
||||
request_counts_completed INTEGER NOT NULL DEFAULT 0,
|
||||
request_counts_failed INTEGER NOT NULL DEFAULT 0,
|
||||
output_file_id TEXT,
|
||||
error_file_id TEXT,
|
||||
created_at TEXT NOT NULL,
|
||||
completed_at TEXT,
|
||||
expires_at TEXT,
|
||||
metadata TEXT
|
||||
);
|
||||
",
|
||||
)?;
|
||||
|
||||
|
||||
@@ -323,10 +323,7 @@ pub fn admin_router(shared: SharedState, token: Arc<String>) -> Router {
|
||||
"/admin/api/mcp-servers",
|
||||
get(list_mcp_servers).post(add_mcp_server),
|
||||
)
|
||||
.route(
|
||||
"/admin/api/mcp-servers/{name}",
|
||||
delete(remove_mcp_server),
|
||||
)
|
||||
.route("/admin/api/mcp-servers/{name}", delete(remove_mcp_server))
|
||||
.with_state(shared.clone())
|
||||
// Innermost: CSRF check runs after auth succeeds.
|
||||
.layer(middleware::from_fn_with_state(shared.clone(), validate_csrf))
|
||||
@@ -1647,7 +1644,11 @@ async fn list_mcp_servers(State(shared): State<SharedState>) -> axum::response::
|
||||
return (StatusCode::OK, Json(serde_json::json!({"servers": []}))).into_response();
|
||||
};
|
||||
let servers = mgr.list_servers_blocking();
|
||||
(StatusCode::OK, Json(serde_json::json!({"servers": servers}))).into_response()
|
||||
(
|
||||
StatusCode::OK,
|
||||
Json(serde_json::json!({"servers": servers})),
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
|
||||
/// POST /admin/api/mcp-servers - Register an MCP server. Body: { name, url }.
|
||||
|
||||
@@ -122,8 +122,7 @@ impl BackendError {
|
||||
_ => {
|
||||
let status = self.api_error_status();
|
||||
let msg = self.api_error_message();
|
||||
infer_error_kind(status.unwrap_or(0), Some(msg.as_str()))
|
||||
.unwrap_or("unknown")
|
||||
infer_error_kind(status.unwrap_or(0), Some(msg.as_str())).unwrap_or("unknown")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -102,7 +102,7 @@ pub(crate) async fn create_anthropic_batch(
|
||||
let model_clone = model.clone();
|
||||
let result = tokio::task::spawn_blocking(move || {
|
||||
let conn = db.lock().unwrap_or_else(|e| e.into_inner());
|
||||
db::init_batch_tables(&conn)?;
|
||||
db::init_anthropic_batch_map_table(&conn)?;
|
||||
db::insert_anthropic_batch_map(&conn, &our_id, &oai_id)?;
|
||||
// Store model alongside mapping for result translation.
|
||||
conn.execute(
|
||||
@@ -185,7 +185,7 @@ pub(crate) async fn get_anthropic_batch(
|
||||
let batch_id_clone = batch_id.clone();
|
||||
let mapping = tokio::task::spawn_blocking(move || {
|
||||
let conn = db.lock().unwrap_or_else(|e| e.into_inner());
|
||||
db::init_batch_tables(&conn)?;
|
||||
db::init_anthropic_batch_map_table(&conn)?;
|
||||
db::get_anthropic_batch_map(&conn, &batch_id_clone)
|
||||
})
|
||||
.await;
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
// SQLite CRUD for batch_file and batch_job tables.
|
||||
// Anthropic-to-OpenAI batch ID mapping table.
|
||||
// The batch_file and batch_job tables are now managed by anyllm_batch_engine.
|
||||
|
||||
use super::{BatchJob, BatchStatus, RequestCounts};
|
||||
use crate::admin::db::now_iso8601;
|
||||
use rusqlite::{params, Connection};
|
||||
|
||||
/// Mapping from our Anthropic batch ID to the upstream OpenAI batch ID.
|
||||
@@ -12,7 +11,7 @@ pub struct AnthropicBatchMap {
|
||||
pub model: String,
|
||||
}
|
||||
|
||||
/// Create the anthropic_batch_map table if it doesn't exist.
|
||||
/// Create the anthropic_batch_map table if it doesn't exist (old schema for proxy use).
|
||||
pub fn init_anthropic_batch_map_table(conn: &Connection) -> rusqlite::Result<()> {
|
||||
conn.execute_batch(
|
||||
"CREATE TABLE IF NOT EXISTS anthropic_batch_map (
|
||||
@@ -73,385 +72,13 @@ pub fn set_anthropic_batch_output_file(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Create the batch_file and batch_job tables if they do not exist.
|
||||
pub fn init_batch_tables(conn: &Connection) -> rusqlite::Result<()> {
|
||||
conn.execute_batch(
|
||||
"
|
||||
CREATE TABLE IF NOT EXISTS batch_file (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
file_id TEXT NOT NULL UNIQUE,
|
||||
key_id INTEGER,
|
||||
purpose TEXT NOT NULL,
|
||||
filename TEXT,
|
||||
byte_size INTEGER NOT NULL,
|
||||
line_count INTEGER NOT NULL,
|
||||
content BLOB NOT NULL,
|
||||
created_at TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS batch_job (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
batch_id TEXT NOT NULL UNIQUE,
|
||||
key_id INTEGER,
|
||||
input_file_id TEXT NOT NULL,
|
||||
backend_batch_id TEXT,
|
||||
backend_name TEXT NOT NULL,
|
||||
status TEXT NOT NULL,
|
||||
request_counts_total INTEGER NOT NULL DEFAULT 0,
|
||||
request_counts_completed INTEGER NOT NULL DEFAULT 0,
|
||||
request_counts_failed INTEGER NOT NULL DEFAULT 0,
|
||||
output_file_id TEXT,
|
||||
error_file_id TEXT,
|
||||
created_at TEXT NOT NULL,
|
||||
completed_at TEXT,
|
||||
expires_at TEXT,
|
||||
metadata TEXT
|
||||
);
|
||||
",
|
||||
)?;
|
||||
// Also create the Anthropic batch ID mapping table.
|
||||
init_anthropic_batch_map_table(conn)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Insert a new batch file record.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn insert_batch_file(
|
||||
conn: &Connection,
|
||||
file_id: &str,
|
||||
key_id: Option<i64>,
|
||||
purpose: &str,
|
||||
filename: Option<&str>,
|
||||
byte_size: i64,
|
||||
line_count: i64,
|
||||
content: &[u8],
|
||||
) -> rusqlite::Result<()> {
|
||||
conn.execute(
|
||||
"INSERT INTO batch_file (file_id, key_id, purpose, filename, byte_size, line_count, content, created_at)
|
||||
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
|
||||
params![
|
||||
file_id,
|
||||
key_id,
|
||||
purpose,
|
||||
filename,
|
||||
byte_size,
|
||||
line_count,
|
||||
content,
|
||||
now_iso8601(),
|
||||
],
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Check if a batch file exists by file_id. Returns (byte_size, line_count, created_at) if found.
|
||||
pub fn get_batch_file_meta(
|
||||
conn: &Connection,
|
||||
file_id: &str,
|
||||
) -> rusqlite::Result<Option<(i64, i64, String)>> {
|
||||
let mut stmt = conn
|
||||
.prepare("SELECT byte_size, line_count, created_at FROM batch_file WHERE file_id = ?1")?;
|
||||
let mut rows = stmt.query_map(params![file_id], |row| {
|
||||
Ok((row.get(0)?, row.get(1)?, row.get(2)?))
|
||||
})?;
|
||||
rows.next().transpose()
|
||||
}
|
||||
|
||||
/// Insert a new batch job record.
|
||||
pub fn insert_batch_job(
|
||||
conn: &Connection,
|
||||
batch_id: &str,
|
||||
key_id: Option<i64>,
|
||||
input_file_id: &str,
|
||||
backend_name: &str,
|
||||
line_count: i64,
|
||||
metadata: Option<&serde_json::Value>,
|
||||
) -> rusqlite::Result<()> {
|
||||
let meta_str = metadata.map(|m| serde_json::to_string(m).unwrap_or_default());
|
||||
conn.execute(
|
||||
"INSERT INTO batch_job (batch_id, key_id, input_file_id, backend_name, status, request_counts_total, created_at, metadata)
|
||||
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
|
||||
params![
|
||||
batch_id,
|
||||
key_id,
|
||||
input_file_id,
|
||||
backend_name,
|
||||
BatchStatus::Validating.as_str(),
|
||||
line_count,
|
||||
now_iso8601(),
|
||||
meta_str,
|
||||
],
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Fetch a single batch job by batch_id.
|
||||
pub fn get_batch_job(conn: &Connection, batch_id: &str) -> rusqlite::Result<Option<BatchJob>> {
|
||||
let mut stmt = conn.prepare(
|
||||
"SELECT batch_id, input_file_id, backend_name, status,
|
||||
request_counts_total, request_counts_completed, request_counts_failed,
|
||||
output_file_id, error_file_id, created_at, completed_at, expires_at, metadata
|
||||
FROM batch_job WHERE batch_id = ?1",
|
||||
)?;
|
||||
let mut rows = stmt.query_map(params![batch_id], row_to_batch_job)?;
|
||||
rows.next().transpose()
|
||||
}
|
||||
|
||||
/// Update the status (and optional completion fields) of a batch job.
|
||||
pub fn update_batch_job_status(
|
||||
conn: &Connection,
|
||||
batch_id: &str,
|
||||
status: &BatchStatus,
|
||||
completed_count: Option<i64>,
|
||||
failed_count: Option<i64>,
|
||||
output_file_id: Option<&str>,
|
||||
error_file_id: Option<&str>,
|
||||
) -> rusqlite::Result<bool> {
|
||||
let completed_at = if matches!(status, BatchStatus::Completed | BatchStatus::Failed) {
|
||||
Some(now_iso8601())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let changed = conn.execute(
|
||||
"UPDATE batch_job SET status = ?1,
|
||||
request_counts_completed = COALESCE(?2, request_counts_completed),
|
||||
request_counts_failed = COALESCE(?3, request_counts_failed),
|
||||
output_file_id = COALESCE(?4, output_file_id),
|
||||
error_file_id = COALESCE(?5, error_file_id),
|
||||
completed_at = COALESCE(?6, completed_at)
|
||||
WHERE batch_id = ?7",
|
||||
params![
|
||||
status.as_str(),
|
||||
completed_count,
|
||||
failed_count,
|
||||
output_file_id,
|
||||
error_file_id,
|
||||
completed_at,
|
||||
batch_id,
|
||||
],
|
||||
)?;
|
||||
Ok(changed > 0)
|
||||
}
|
||||
|
||||
/// List batch jobs, optionally filtered by key_id, with cursor pagination.
|
||||
pub fn list_batch_jobs(
|
||||
conn: &Connection,
|
||||
key_id: Option<i64>,
|
||||
limit: u32,
|
||||
after: Option<&str>,
|
||||
) -> rusqlite::Result<Vec<BatchJob>> {
|
||||
let mut sql = String::from(
|
||||
"SELECT batch_id, input_file_id, backend_name, status,
|
||||
request_counts_total, request_counts_completed, request_counts_failed,
|
||||
output_file_id, error_file_id, created_at, completed_at, expires_at, metadata
|
||||
FROM batch_job WHERE 1=1",
|
||||
);
|
||||
let mut param_values: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
|
||||
|
||||
if let Some(kid) = key_id {
|
||||
sql.push_str(" AND key_id = ?");
|
||||
param_values.push(Box::new(kid));
|
||||
}
|
||||
if let Some(cursor) = after {
|
||||
sql.push_str(" AND batch_id < ?");
|
||||
param_values.push(Box::new(cursor.to_string()));
|
||||
}
|
||||
|
||||
sql.push_str(" ORDER BY id DESC LIMIT ?");
|
||||
param_values.push(Box::new(limit));
|
||||
|
||||
let params_refs: Vec<&dyn rusqlite::types::ToSql> =
|
||||
param_values.iter().map(|p| p.as_ref()).collect();
|
||||
|
||||
let mut stmt = conn.prepare(&sql)?;
|
||||
let rows = stmt.query_map(params_refs.as_slice(), row_to_batch_job)?;
|
||||
rows.collect()
|
||||
}
|
||||
|
||||
/// Map a SQLite row to a BatchJob.
|
||||
fn row_to_batch_job(row: &rusqlite::Row) -> rusqlite::Result<BatchJob> {
|
||||
let status_str: String = row.get(3)?;
|
||||
let created_at_str: String = row.get(9)?;
|
||||
let metadata_str: Option<String> = row.get(12)?;
|
||||
|
||||
Ok(BatchJob {
|
||||
id: row.get(0)?,
|
||||
object: "batch".to_string(),
|
||||
endpoint: "/v1/chat/completions".to_string(),
|
||||
status: BatchStatus::from_str_status(&status_str),
|
||||
input_file_id: row.get(1)?,
|
||||
completion_window: "24h".to_string(),
|
||||
created_at: iso8601_to_epoch(&created_at_str),
|
||||
request_counts: RequestCounts {
|
||||
total: row.get(4)?,
|
||||
completed: row.get(5)?,
|
||||
failed: row.get(6)?,
|
||||
},
|
||||
metadata: metadata_str.and_then(|s| serde_json::from_str(&s).ok()),
|
||||
output_file_id: row.get(7)?,
|
||||
error_file_id: row.get(8)?,
|
||||
completed_at: row
|
||||
.get::<_, Option<String>>(10)?
|
||||
.map(|s| iso8601_to_epoch(&s)),
|
||||
expires_at: row
|
||||
.get::<_, Option<String>>(11)?
|
||||
.map(|s| iso8601_to_epoch(&s)),
|
||||
})
|
||||
}
|
||||
|
||||
/// Approximate conversion from ISO 8601 string to unix epoch seconds.
|
||||
/// Falls back to 0 on parse failure (non-critical metadata field).
|
||||
fn iso8601_to_epoch(s: &str) -> i64 {
|
||||
// Parse "YYYY-MM-DDTHH:MM:SSZ" manually (no chrono dependency).
|
||||
let parts: Vec<&str> = s.split('T').collect();
|
||||
if parts.len() != 2 {
|
||||
return 0;
|
||||
}
|
||||
let date_parts: Vec<u64> = parts[0].split('-').filter_map(|p| p.parse().ok()).collect();
|
||||
let time_str = parts[1].trim_end_matches('Z');
|
||||
let time_parts: Vec<u64> = time_str.split(':').filter_map(|p| p.parse().ok()).collect();
|
||||
|
||||
if date_parts.len() != 3 || time_parts.len() != 3 {
|
||||
return 0;
|
||||
}
|
||||
|
||||
let (y, m, d) = (date_parts[0], date_parts[1], date_parts[2]);
|
||||
let (hh, mm, ss) = (time_parts[0], time_parts[1], time_parts[2]);
|
||||
|
||||
// Days from epoch using the inverse of the Howard Hinnant algorithm.
|
||||
let y_adj = if m <= 2 { y - 1 } else { y };
|
||||
let era = y_adj / 400;
|
||||
let yoe = y_adj - era * 400;
|
||||
let m_adj = if m > 2 { m - 3 } else { m + 9 };
|
||||
let doy = (153 * m_adj + 2) / 5 + d - 1;
|
||||
let doe = yoe * 365 + yoe / 4 - yoe / 100 + doy;
|
||||
let days = era * 146097 + doe - 719468;
|
||||
|
||||
(days * 86400 + hh * 3600 + mm * 60 + ss) as i64
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::admin::db::init_db;
|
||||
|
||||
fn test_db() -> Connection {
|
||||
let conn = Connection::open_in_memory().unwrap();
|
||||
init_db(&conn).unwrap();
|
||||
init_batch_tables(&conn).unwrap();
|
||||
conn
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn insert_and_get_batch_file() {
|
||||
let conn = test_db();
|
||||
insert_batch_file(
|
||||
&conn,
|
||||
"file-abc123",
|
||||
None,
|
||||
"batch",
|
||||
Some("test.jsonl"),
|
||||
1024,
|
||||
10,
|
||||
b"test content",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let meta = get_batch_file_meta(&conn, "file-abc123").unwrap();
|
||||
assert!(meta.is_some());
|
||||
let (size, count, _created) = meta.unwrap();
|
||||
assert_eq!(size, 1024);
|
||||
assert_eq!(count, 10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn insert_and_get_batch_job() {
|
||||
let conn = test_db();
|
||||
insert_batch_file(&conn, "file-input1", None, "batch", None, 512, 5, b"data").unwrap();
|
||||
|
||||
insert_batch_job(&conn, "batch-job1", None, "file-input1", "openai", 5, None).unwrap();
|
||||
|
||||
let job = get_batch_job(&conn, "batch-job1").unwrap();
|
||||
assert!(job.is_some());
|
||||
let job = job.unwrap();
|
||||
assert_eq!(job.id, "batch-job1");
|
||||
assert_eq!(job.status, BatchStatus::Validating);
|
||||
assert_eq!(job.request_counts.total, 5);
|
||||
assert_eq!(job.input_file_id, "file-input1");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn update_batch_job_status_works() {
|
||||
let conn = test_db();
|
||||
insert_batch_file(&conn, "file-u1", None, "batch", None, 100, 2, b"d").unwrap();
|
||||
insert_batch_job(&conn, "batch-u1", None, "file-u1", "openai", 2, None).unwrap();
|
||||
|
||||
let ok = update_batch_job_status(
|
||||
&conn,
|
||||
"batch-u1",
|
||||
&BatchStatus::Completed,
|
||||
Some(2),
|
||||
Some(0),
|
||||
Some("file-out1"),
|
||||
None,
|
||||
)
|
||||
.unwrap();
|
||||
assert!(ok);
|
||||
|
||||
let job = get_batch_job(&conn, "batch-u1").unwrap().unwrap();
|
||||
assert_eq!(job.status, BatchStatus::Completed);
|
||||
assert_eq!(job.request_counts.completed, 2);
|
||||
assert!(job.output_file_id.is_some());
|
||||
assert!(job.completed_at.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn list_batch_jobs_with_pagination() {
|
||||
let conn = test_db();
|
||||
insert_batch_file(&conn, "file-l1", None, "batch", None, 10, 1, b"d").unwrap();
|
||||
|
||||
for i in 0..5 {
|
||||
insert_batch_job(
|
||||
&conn,
|
||||
&format!("batch-l{i}"),
|
||||
Some(1),
|
||||
"file-l1",
|
||||
"openai",
|
||||
1,
|
||||
None,
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
let all = list_batch_jobs(&conn, Some(1), 10, None).unwrap();
|
||||
assert_eq!(all.len(), 5);
|
||||
|
||||
let page = list_batch_jobs(&conn, Some(1), 2, None).unwrap();
|
||||
assert_eq!(page.len(), 2);
|
||||
|
||||
// Different key_id returns nothing
|
||||
let empty = list_batch_jobs(&conn, Some(999), 10, None).unwrap();
|
||||
assert!(empty.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn get_nonexistent_job() {
|
||||
let conn = test_db();
|
||||
let job = get_batch_job(&conn, "batch-nope").unwrap();
|
||||
assert!(job.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn iso8601_round_trip() {
|
||||
// 2026-03-22T10:30:00Z
|
||||
let epoch = iso8601_to_epoch("2026-03-22T10:30:00Z");
|
||||
assert!(epoch > 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn anthropic_batch_map_round_trip() {
|
||||
let conn = rusqlite::Connection::open_in_memory().unwrap();
|
||||
init_batch_tables(&conn).unwrap();
|
||||
init_anthropic_batch_map_table(&conn).unwrap();
|
||||
|
||||
insert_anthropic_batch_map(&conn, "msgbatch_our1", "batch_openai1").unwrap();
|
||||
|
||||
@@ -1,281 +1,11 @@
|
||||
// Batch processing types and JSONL validation.
|
||||
// Implements OpenAI-compatible batch file upload and job management.
|
||||
//! Batch API HTTP handlers. Types and logic live in anyllm_batch_engine.
|
||||
|
||||
pub mod anthropic_batch;
|
||||
pub mod db;
|
||||
pub mod openai_batch_client;
|
||||
pub mod routes;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashSet;
|
||||
|
||||
/// Maximum number of lines in a JSONL batch file.
|
||||
const MAX_LINE_COUNT: usize = 50_000;
|
||||
|
||||
/// Maximum file size in bytes (100 MB).
|
||||
const MAX_FILE_SIZE: usize = 100 * 1024 * 1024;
|
||||
|
||||
/// Maximum length of a custom_id field.
|
||||
const MAX_CUSTOM_ID_LEN: usize = 64;
|
||||
|
||||
/// A batch input file stored in SQLite.
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct BatchFile {
|
||||
pub id: String,
|
||||
pub object: String,
|
||||
pub bytes: i64,
|
||||
pub created_at: i64,
|
||||
pub filename: Option<String>,
|
||||
pub purpose: String,
|
||||
}
|
||||
|
||||
/// A batch processing job.
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct BatchJob {
|
||||
pub id: String,
|
||||
pub object: String,
|
||||
pub endpoint: String,
|
||||
pub status: BatchStatus,
|
||||
pub input_file_id: String,
|
||||
pub completion_window: String,
|
||||
pub created_at: i64,
|
||||
pub request_counts: RequestCounts,
|
||||
pub metadata: Option<serde_json::Value>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub output_file_id: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub error_file_id: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub completed_at: Option<i64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub expires_at: Option<i64>,
|
||||
}
|
||||
|
||||
/// Counts of requests within a batch job.
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct RequestCounts {
|
||||
pub total: i64,
|
||||
pub completed: i64,
|
||||
pub failed: i64,
|
||||
}
|
||||
|
||||
/// Batch job lifecycle status.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum BatchStatus {
|
||||
Validating,
|
||||
InProgress,
|
||||
Completed,
|
||||
Failed,
|
||||
Expired,
|
||||
Cancelling,
|
||||
Cancelled,
|
||||
}
|
||||
|
||||
impl BatchStatus {
|
||||
/// Convert from the string stored in SQLite.
|
||||
pub fn from_str_status(s: &str) -> Self {
|
||||
match s {
|
||||
"validating" => Self::Validating,
|
||||
"in_progress" => Self::InProgress,
|
||||
"completed" => Self::Completed,
|
||||
"failed" => Self::Failed,
|
||||
"expired" => Self::Expired,
|
||||
"cancelling" => Self::Cancelling,
|
||||
"cancelled" => Self::Cancelled,
|
||||
_ => Self::Failed,
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert to the string stored in SQLite.
|
||||
pub fn as_str(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Validating => "validating",
|
||||
Self::InProgress => "in_progress",
|
||||
Self::Completed => "completed",
|
||||
Self::Failed => "failed",
|
||||
Self::Expired => "expired",
|
||||
Self::Cancelling => "cancelling",
|
||||
Self::Cancelled => "cancelled",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Result of JSONL validation: line count on success, error message on failure.
|
||||
#[derive(Debug)]
|
||||
pub struct ValidatedJsonl {
|
||||
pub line_count: usize,
|
||||
}
|
||||
|
||||
/// Validate a JSONL batch file.
|
||||
///
|
||||
/// Each line must be valid JSON with a unique `custom_id` (string, max 64 chars)
|
||||
/// and a `body` object containing a `model` field. Max 50,000 lines, 100 MB.
|
||||
///
|
||||
/// Takes a `BufRead` to read line-by-line without requiring a contiguous UTF-8
|
||||
/// string for the entire file.
|
||||
pub fn validate_jsonl(mut reader: impl std::io::BufRead) -> Result<ValidatedJsonl, String> {
|
||||
let mut seen_ids = HashSet::new();
|
||||
let mut line_count = 0usize;
|
||||
let mut raw_line_num = 0usize;
|
||||
let mut bytes_read = 0usize;
|
||||
let mut line_buf = String::new();
|
||||
|
||||
loop {
|
||||
line_buf.clear();
|
||||
let n = reader
|
||||
.read_line(&mut line_buf)
|
||||
.map_err(|e| format!("Read error: {e}"))?;
|
||||
if n == 0 {
|
||||
break; // EOF
|
||||
}
|
||||
raw_line_num += 1;
|
||||
bytes_read += n;
|
||||
if bytes_read > MAX_FILE_SIZE {
|
||||
return Err(format!(
|
||||
"File exceeds maximum size of {} bytes",
|
||||
MAX_FILE_SIZE
|
||||
));
|
||||
}
|
||||
|
||||
let line = line_buf.trim();
|
||||
if line.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
line_count += 1;
|
||||
if line_count > MAX_LINE_COUNT {
|
||||
return Err(format!("File exceeds maximum of {MAX_LINE_COUNT} lines"));
|
||||
}
|
||||
|
||||
let parsed: serde_json::Value = serde_json::from_str(line)
|
||||
.map_err(|e| format!("Line {raw_line_num}: invalid JSON: {e}"))?;
|
||||
|
||||
let obj = parsed
|
||||
.as_object()
|
||||
.ok_or_else(|| format!("Line {raw_line_num}: expected JSON object"))?;
|
||||
|
||||
let custom_id = obj
|
||||
.get("custom_id")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| format!("Line {raw_line_num}: missing or non-string 'custom_id'"))?;
|
||||
|
||||
if custom_id.len() > MAX_CUSTOM_ID_LEN {
|
||||
return Err(format!(
|
||||
"Line {raw_line_num}: custom_id exceeds maximum length of {MAX_CUSTOM_ID_LEN} characters"
|
||||
));
|
||||
}
|
||||
|
||||
if !seen_ids.insert(custom_id.to_string()) {
|
||||
return Err(format!(
|
||||
"Line {raw_line_num}: duplicate custom_id '{custom_id}'"
|
||||
));
|
||||
}
|
||||
|
||||
let body = obj
|
||||
.get("body")
|
||||
.and_then(|v| v.as_object())
|
||||
.ok_or_else(|| format!("Line {raw_line_num}: missing or non-object 'body'"))?;
|
||||
|
||||
if !body.contains_key("model") {
|
||||
return Err(format!("Line {raw_line_num}: body missing 'model' field"));
|
||||
}
|
||||
}
|
||||
|
||||
if line_count == 0 {
|
||||
return Err("File is empty".to_string());
|
||||
}
|
||||
|
||||
Ok(ValidatedJsonl { line_count })
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::io::{BufReader, Cursor};
|
||||
|
||||
fn check(data: &str) -> Result<ValidatedJsonl, String> {
|
||||
validate_jsonl(BufReader::new(Cursor::new(data.as_bytes())))
|
||||
}
|
||||
|
||||
fn check_bytes(data: &[u8]) -> Result<ValidatedJsonl, String> {
|
||||
validate_jsonl(BufReader::new(Cursor::new(data)))
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn valid_jsonl() {
|
||||
let data = r#"{"custom_id": "req-1", "body": {"model": "gpt-4o", "messages": []}}
|
||||
{"custom_id": "req-2", "body": {"model": "gpt-4o", "messages": []}}"#;
|
||||
let result = check(data);
|
||||
assert!(result.is_ok());
|
||||
assert_eq!(result.unwrap().line_count, 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_custom_id() {
|
||||
let data = r#"{"body": {"model": "gpt-4o"}}"#;
|
||||
let result = check(data);
|
||||
assert!(result.is_err());
|
||||
assert!(result.unwrap_err().contains("custom_id"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_body_model() {
|
||||
let data = r#"{"custom_id": "req-1", "body": {"messages": []}}"#;
|
||||
let result = check(data);
|
||||
assert!(result.is_err());
|
||||
assert!(result.unwrap_err().contains("model"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn duplicate_custom_id() {
|
||||
let data = r#"{"custom_id": "req-1", "body": {"model": "gpt-4o"}}
|
||||
{"custom_id": "req-1", "body": {"model": "gpt-4o"}}"#;
|
||||
let result = check(data);
|
||||
assert!(result.is_err());
|
||||
assert!(result.unwrap_err().contains("duplicate"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn oversized_custom_id() {
|
||||
let long_id = "a".repeat(65);
|
||||
let data = format!(r#"{{"custom_id": "{long_id}", "body": {{"model": "gpt-4o"}}}}"#);
|
||||
let result = check(&data);
|
||||
assert!(result.is_err());
|
||||
assert!(result.unwrap_err().contains("maximum length"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_file() {
|
||||
let result = check_bytes(b"");
|
||||
assert!(result.is_err());
|
||||
assert!(result.unwrap_err().contains("empty"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_json_line() {
|
||||
let data = b"not json at all";
|
||||
let result = check_bytes(data);
|
||||
assert!(result.is_err());
|
||||
assert!(result.unwrap_err().contains("invalid JSON"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn blank_lines_skipped() {
|
||||
let data = r#"{"custom_id": "req-1", "body": {"model": "gpt-4o"}}
|
||||
|
||||
{"custom_id": "req-2", "body": {"model": "gpt-4o"}}"#;
|
||||
let result = check(data);
|
||||
assert!(result.is_ok());
|
||||
assert_eq!(result.unwrap().line_count, 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn error_reports_absolute_line_number_with_blank_lines() {
|
||||
// Blank line at position 1, bad JSON at position 2.
|
||||
// Should report "Line 2", not "Line 1".
|
||||
let data = "\n{\"custom_id\": \"ok\", \"body\": INVALID}";
|
||||
let err = check(data).unwrap_err();
|
||||
assert!(err.contains("Line 2"), "expected 'Line 2' in: {err}");
|
||||
}
|
||||
}
|
||||
// Re-export engine types for handler use.
|
||||
pub use anyllm_batch_engine::job::{BatchId, BatchJob, BatchStatus, RequestCounts};
|
||||
pub use anyllm_batch_engine::validation::{validate_jsonl, ValidatedJsonl};
|
||||
pub use anyllm_batch_engine::BatchEngine;
|
||||
|
||||
+214
-138
@@ -1,10 +1,9 @@
|
||||
// Axum handlers for batch file upload and job management.
|
||||
// POST /v1/files, POST /v1/batches, GET /v1/batches/{id}, GET /v1/batches
|
||||
|
||||
use super::db;
|
||||
use super::validate_jsonl;
|
||||
use crate::backend::BackendClient;
|
||||
use crate::server::routes::AppState;
|
||||
use anyllm_batch_engine::job::{BatchSubmission, ExecutionMode, SourceFormat, SubmissionItem};
|
||||
use anyllm_translate::anthropic;
|
||||
use anyllm_translate::mapping::errors_map::create_anthropic_error;
|
||||
use axum::{
|
||||
@@ -16,11 +15,9 @@ use serde::Deserialize;
|
||||
use std::io::{BufReader, Cursor};
|
||||
|
||||
/// POST /v1/files - Upload a JSONL batch file via multipart/form-data.
|
||||
///
|
||||
/// Expects fields: `purpose` (must be "batch") and `file` (the JSONL content).
|
||||
pub async fn upload_file(State(state): State<AppState>, mut multipart: Multipart) -> Response {
|
||||
let db = match state.shared.as_ref().map(|s| s.db.clone()) {
|
||||
Some(db) => db,
|
||||
let engine = match state.batch_engine.as_ref() {
|
||||
Some(e) => e.clone(),
|
||||
None => return service_unavailable("Batch storage not available"),
|
||||
};
|
||||
|
||||
@@ -31,9 +28,7 @@ pub async fn upload_file(State(state): State<AppState>, mut multipart: Multipart
|
||||
while let Ok(Some(field)) = multipart.next_field().await {
|
||||
let field_name = field.name().unwrap_or("").to_string();
|
||||
match field_name.as_str() {
|
||||
"purpose" => {
|
||||
purpose = field.text().await.ok();
|
||||
}
|
||||
"purpose" => purpose = field.text().await.ok(),
|
||||
"file" => {
|
||||
filename = field.file_name().map(|s| s.to_string());
|
||||
file_data = field.bytes().await.ok();
|
||||
@@ -42,82 +37,61 @@ pub async fn upload_file(State(state): State<AppState>, mut multipart: Multipart
|
||||
}
|
||||
}
|
||||
|
||||
let purpose = match purpose.as_deref() {
|
||||
Some("batch") => "batch",
|
||||
match purpose.as_deref() {
|
||||
Some("batch") => {}
|
||||
Some(other) => {
|
||||
return bad_request(&format!(
|
||||
"Unsupported purpose: '{other}'. Only 'batch' is supported."
|
||||
));
|
||||
}
|
||||
None => {
|
||||
return bad_request("Missing required field 'purpose'");
|
||||
}
|
||||
};
|
||||
None => return bad_request("Missing required field 'purpose'"),
|
||||
}
|
||||
|
||||
let data = match file_data {
|
||||
Some(d) if !d.is_empty() => d,
|
||||
_ => {
|
||||
return bad_request("Missing or empty 'file' field");
|
||||
}
|
||||
_ => return bad_request("Missing or empty 'file' field"),
|
||||
};
|
||||
|
||||
// Validate JSONL structure
|
||||
let validated = match validate_jsonl(BufReader::new(Cursor::new(data.as_ref()))) {
|
||||
Ok(v) => v,
|
||||
Err(e) => {
|
||||
return bad_request(&format!("Invalid JSONL: {e}"));
|
||||
}
|
||||
};
|
||||
let validated =
|
||||
match anyllm_batch_engine::validate_jsonl(BufReader::new(Cursor::new(data.as_ref()))) {
|
||||
Ok(v) => v,
|
||||
Err(e) => return bad_request(&format!("Invalid JSONL: {e}")),
|
||||
};
|
||||
|
||||
let file_id = format!("file-{}", uuid::Uuid::new_v4());
|
||||
let byte_size = data.len() as i64;
|
||||
let line_count = validated.line_count as i64;
|
||||
|
||||
// Insert into SQLite on the blocking threadpool
|
||||
let file_id_clone = file_id.clone();
|
||||
let filename_clone = filename.clone();
|
||||
let data_ref = data.as_ref().to_vec();
|
||||
let result = tokio::task::spawn_blocking(move || {
|
||||
let conn = db.lock().unwrap_or_else(|e| e.into_inner());
|
||||
db::init_batch_tables(&conn)?;
|
||||
db::insert_batch_file(
|
||||
&conn,
|
||||
&file_id_clone,
|
||||
match engine
|
||||
.file_store
|
||||
.insert(
|
||||
&file_id,
|
||||
None,
|
||||
purpose,
|
||||
filename_clone.as_deref(),
|
||||
byte_size,
|
||||
filename.as_deref(),
|
||||
data.as_ref(),
|
||||
line_count,
|
||||
&data_ref,
|
||||
)
|
||||
})
|
||||
.await;
|
||||
|
||||
match result {
|
||||
Ok(Ok(())) => {
|
||||
.await
|
||||
{
|
||||
Ok(()) => {
|
||||
let now_epoch = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs() as i64;
|
||||
|
||||
let file_obj = super::BatchFile {
|
||||
id: file_id,
|
||||
object: "file".to_string(),
|
||||
bytes: byte_size,
|
||||
created_at: now_epoch,
|
||||
filename,
|
||||
purpose: purpose.to_string(),
|
||||
};
|
||||
let file_obj = serde_json::json!({
|
||||
"id": file_id,
|
||||
"object": "file",
|
||||
"bytes": byte_size,
|
||||
"created_at": now_epoch,
|
||||
"filename": filename,
|
||||
"purpose": "batch",
|
||||
});
|
||||
(StatusCode::OK, Json(file_obj)).into_response()
|
||||
}
|
||||
Ok(Err(e)) => {
|
||||
Err(e) => {
|
||||
tracing::error!(error = %e, "failed to store batch file");
|
||||
internal_error("Failed to store file")
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::error!(error = %e, "spawn_blocking panicked");
|
||||
internal_error("Internal error")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -141,8 +115,6 @@ fn default_completion_window() -> String {
|
||||
}
|
||||
|
||||
/// POST /v1/batches - Create a new batch job.
|
||||
///
|
||||
/// Returns 501 for unsupported backends (vertex, gemini, anthropic, bedrock).
|
||||
pub async fn create_batch(
|
||||
State(state): State<AppState>,
|
||||
Json(req): Json<CreateBatchRequest>,
|
||||
@@ -155,75 +127,67 @@ pub async fn create_batch(
|
||||
));
|
||||
}
|
||||
|
||||
let db = match state.shared.as_ref().map(|s| s.db.clone()) {
|
||||
Some(db) => db,
|
||||
let engine = match state.batch_engine.as_ref() {
|
||||
Some(e) => e.clone(),
|
||||
None => return service_unavailable("Batch storage not available"),
|
||||
};
|
||||
|
||||
let input_file_id = req.input_file_id.clone();
|
||||
let batch_id = format!("batch-{}", uuid::Uuid::new_v4());
|
||||
let backend_name = state.backend_name.clone();
|
||||
let metadata = req.metadata.clone();
|
||||
// Read file content from file store.
|
||||
let content = match engine.file_store.get_content(&req.input_file_id).await {
|
||||
Ok(Some(c)) => c,
|
||||
Ok(None) => return bad_request(&format!("Input file '{}' not found", req.input_file_id)),
|
||||
Err(e) => {
|
||||
tracing::error!(error = %e, "failed to read batch file content");
|
||||
return internal_error("Failed to read file");
|
||||
}
|
||||
};
|
||||
|
||||
// Verify input file exists and get line count
|
||||
let batch_id_clone = batch_id.clone();
|
||||
let result = tokio::task::spawn_blocking(move || {
|
||||
let conn = db.lock().unwrap_or_else(|e| e.into_inner());
|
||||
db::init_batch_tables(&conn)?;
|
||||
// Parse JSONL into submission items.
|
||||
let items: Vec<SubmissionItem> = match parse_jsonl_items(&content) {
|
||||
Ok(items) => items,
|
||||
Err(e) => return bad_request(&format!("Invalid JSONL: {e}")),
|
||||
};
|
||||
|
||||
let meta = db::get_batch_file_meta(&conn, &input_file_id)?;
|
||||
let (_byte_size, line_count, _created) = match meta {
|
||||
Some(m) => m,
|
||||
None => {
|
||||
return Ok(None);
|
||||
}
|
||||
};
|
||||
let execution_mode = if is_openai_or_azure_backend(&state.backend) {
|
||||
ExecutionMode::Native {
|
||||
provider: state.backend_name.clone(),
|
||||
}
|
||||
} else {
|
||||
ExecutionMode::ProxyNative
|
||||
};
|
||||
|
||||
db::insert_batch_job(
|
||||
&conn,
|
||||
&batch_id_clone,
|
||||
None,
|
||||
&input_file_id,
|
||||
&backend_name,
|
||||
line_count,
|
||||
metadata.as_ref(),
|
||||
)?;
|
||||
let submission = BatchSubmission {
|
||||
items,
|
||||
execution_mode,
|
||||
input_file_id: req.input_file_id.clone(),
|
||||
key_id: None,
|
||||
webhook_url: None,
|
||||
metadata: req.metadata.clone(),
|
||||
priority: 0,
|
||||
};
|
||||
|
||||
db::get_batch_job(&conn, &batch_id_clone)
|
||||
})
|
||||
.await;
|
||||
|
||||
match result {
|
||||
Ok(Ok(Some(job))) => (StatusCode::OK, Json(job)).into_response(),
|
||||
Ok(Ok(None)) => bad_request(&format!("Input file '{}' not found", req.input_file_id)),
|
||||
Ok(Err(e)) => {
|
||||
tracing::error!(error = %e, "failed to create batch job");
|
||||
internal_error("Failed to create batch job")
|
||||
match engine.submit(submission).await {
|
||||
Ok(job) => (StatusCode::OK, Json(job_to_openai_response(&job))).into_response(),
|
||||
Err(anyllm_batch_engine::EngineError::FileNotFound(_)) => {
|
||||
bad_request(&format!("Input file '{}' not found", req.input_file_id))
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::error!(error = %e, "spawn_blocking panicked");
|
||||
internal_error("Internal error")
|
||||
tracing::error!(error = %e, "failed to create batch job");
|
||||
internal_error("Failed to create batch job")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// GET /v1/batches/{batch_id} - Retrieve a batch job by ID.
|
||||
/// GET /v1/batches/{batch_id}
|
||||
pub async fn get_batch(State(state): State<AppState>, Path(batch_id): Path<String>) -> Response {
|
||||
let db = match state.shared.as_ref().map(|s| s.db.clone()) {
|
||||
Some(db) => db,
|
||||
let engine = match state.batch_engine.as_ref() {
|
||||
Some(e) => e.clone(),
|
||||
None => return service_unavailable("Batch storage not available"),
|
||||
};
|
||||
|
||||
let result = tokio::task::spawn_blocking(move || {
|
||||
let conn = db.lock().unwrap_or_else(|e| e.into_inner());
|
||||
db::init_batch_tables(&conn)?;
|
||||
db::get_batch_job(&conn, &batch_id)
|
||||
})
|
||||
.await;
|
||||
|
||||
match result {
|
||||
Ok(Ok(Some(job))) => (StatusCode::OK, Json(job)).into_response(),
|
||||
Ok(Ok(None)) => {
|
||||
match engine.get(&anyllm_batch_engine::BatchId(batch_id)).await {
|
||||
Ok(Some(job)) => (StatusCode::OK, Json(job_to_openai_response(&job))).into_response(),
|
||||
Ok(None) => {
|
||||
let err = create_anthropic_error(
|
||||
anthropic::ErrorType::NotFoundError,
|
||||
"Batch not found".to_string(),
|
||||
@@ -231,14 +195,10 @@ pub async fn get_batch(State(state): State<AppState>, Path(batch_id): Path<Strin
|
||||
);
|
||||
(StatusCode::NOT_FOUND, Json(err)).into_response()
|
||||
}
|
||||
Ok(Err(e)) => {
|
||||
Err(e) => {
|
||||
tracing::error!(error = %e, "failed to fetch batch job");
|
||||
internal_error("Failed to fetch batch job")
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::error!(error = %e, "spawn_blocking panicked");
|
||||
internal_error("Internal error")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -254,50 +214,154 @@ fn default_limit() -> u32 {
|
||||
20
|
||||
}
|
||||
|
||||
/// GET /v1/batches - List batch jobs with cursor pagination.
|
||||
/// GET /v1/batches
|
||||
pub async fn list_batches(
|
||||
State(state): State<AppState>,
|
||||
Query(query): Query<ListBatchesQuery>,
|
||||
) -> Response {
|
||||
let db = match state.shared.as_ref().map(|s| s.db.clone()) {
|
||||
Some(db) => db,
|
||||
let engine = match state.batch_engine.as_ref() {
|
||||
Some(e) => e.clone(),
|
||||
None => return service_unavailable("Batch storage not available"),
|
||||
};
|
||||
|
||||
let limit = query.limit.min(100);
|
||||
let after = query.after.clone();
|
||||
|
||||
let result = tokio::task::spawn_blocking(move || {
|
||||
let conn = db.lock().unwrap_or_else(|e| e.into_inner());
|
||||
db::init_batch_tables(&conn)?;
|
||||
db::list_batch_jobs(&conn, None, limit, after.as_deref())
|
||||
})
|
||||
.await;
|
||||
|
||||
match result {
|
||||
Ok(Ok(jobs)) => {
|
||||
match engine.list(None, query.after.as_deref(), limit).await {
|
||||
Ok(jobs) => {
|
||||
let has_more = jobs.len() as u32 == limit;
|
||||
let last_id = jobs.last().map(|j| j.id.clone());
|
||||
let first_id = jobs.first().map(|j| j.id.0.clone());
|
||||
let last_id = jobs.last().map(|j| j.id.0.clone());
|
||||
let data: Vec<serde_json::Value> = jobs.iter().map(job_to_openai_response).collect();
|
||||
let response = serde_json::json!({
|
||||
"object": "list",
|
||||
"data": jobs,
|
||||
"data": data,
|
||||
"has_more": has_more,
|
||||
"first_id": jobs.first().map(|j| &j.id),
|
||||
"first_id": first_id,
|
||||
"last_id": last_id,
|
||||
});
|
||||
(StatusCode::OK, Json(response)).into_response()
|
||||
}
|
||||
Ok(Err(e)) => {
|
||||
Err(e) => {
|
||||
tracing::error!(error = %e, "failed to list batch jobs");
|
||||
internal_error("Failed to list batch jobs")
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::error!(error = %e, "spawn_blocking panicked");
|
||||
internal_error("Internal error")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// POST /v1/batches/{batch_id}/cancel
|
||||
pub async fn cancel_batch(State(state): State<AppState>, Path(batch_id): Path<String>) -> Response {
|
||||
let Some(engine) = state.batch_engine.as_ref() else {
|
||||
return not_implemented("batch engine not available");
|
||||
};
|
||||
|
||||
let id = anyllm_batch_engine::BatchId(batch_id);
|
||||
match engine.cancel(&id).await {
|
||||
Ok(_) => match engine.get(&id).await {
|
||||
Ok(Some(job)) => (StatusCode::OK, Json(job_to_openai_response(&job))).into_response(),
|
||||
Ok(None) => not_found_response("batch not found"),
|
||||
Err(e) => internal_error(&e.to_string()),
|
||||
},
|
||||
Err(anyllm_batch_engine::EngineError::Queue(anyllm_batch_engine::QueueError::NotFound)) => {
|
||||
not_found_response("batch not found")
|
||||
}
|
||||
Err(e) => internal_error(&e.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Map a BatchJob to an OpenAI-compatible batch response JSON.
|
||||
pub fn job_to_openai_response(job: &anyllm_batch_engine::BatchJob) -> serde_json::Value {
|
||||
let created_epoch = iso8601_to_epoch(&job.created_at);
|
||||
let completed_epoch = job.completed_at.as_deref().map(iso8601_to_epoch);
|
||||
|
||||
serde_json::json!({
|
||||
"id": job.id.0,
|
||||
"object": "batch",
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"status": map_batch_status(&job.status),
|
||||
"input_file_id": job.input_file_id,
|
||||
"completion_window": "24h",
|
||||
"created_at": created_epoch,
|
||||
"request_counts": {
|
||||
"total": job.request_counts.total,
|
||||
"completed": job.request_counts.succeeded,
|
||||
"failed": job.request_counts.failed,
|
||||
},
|
||||
"metadata": job.metadata,
|
||||
"output_file_id": serde_json::Value::Null,
|
||||
"error_file_id": serde_json::Value::Null,
|
||||
"completed_at": completed_epoch,
|
||||
})
|
||||
}
|
||||
|
||||
/// Map BatchEngine status to OpenAI batch status string.
|
||||
fn map_batch_status(status: &anyllm_batch_engine::BatchStatus) -> &'static str {
|
||||
match status {
|
||||
anyllm_batch_engine::BatchStatus::Queued => "validating",
|
||||
anyllm_batch_engine::BatchStatus::Processing => "in_progress",
|
||||
anyllm_batch_engine::BatchStatus::Completed => "completed",
|
||||
anyllm_batch_engine::BatchStatus::Failed => "failed",
|
||||
anyllm_batch_engine::BatchStatus::Cancelling => "cancelling",
|
||||
anyllm_batch_engine::BatchStatus::Cancelled => "cancelled",
|
||||
anyllm_batch_engine::BatchStatus::Expired => "expired",
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse JSONL bytes into SubmissionItems.
|
||||
fn parse_jsonl_items(content: &[u8]) -> Result<Vec<SubmissionItem>, String> {
|
||||
let mut items = Vec::new();
|
||||
let text = std::str::from_utf8(content).map_err(|e| format!("Invalid UTF-8: {e}"))?;
|
||||
for line in text.lines() {
|
||||
let line = line.trim();
|
||||
if line.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let parsed: serde_json::Value =
|
||||
serde_json::from_str(line).map_err(|e| format!("Invalid JSON: {e}"))?;
|
||||
let obj = parsed.as_object().ok_or("Expected JSON object")?;
|
||||
let custom_id = obj
|
||||
.get("custom_id")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or("Missing custom_id")?
|
||||
.to_string();
|
||||
let body = obj.get("body").cloned().unwrap_or(serde_json::Value::Null);
|
||||
let model = body
|
||||
.get("model")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("unknown")
|
||||
.to_string();
|
||||
items.push(SubmissionItem {
|
||||
custom_id,
|
||||
model,
|
||||
body,
|
||||
source_format: SourceFormat::OpenAI,
|
||||
});
|
||||
}
|
||||
Ok(items)
|
||||
}
|
||||
|
||||
fn iso8601_to_epoch(s: &str) -> i64 {
|
||||
let parts: Vec<&str> = s.split('T').collect();
|
||||
if parts.len() != 2 {
|
||||
return 0;
|
||||
}
|
||||
let date_parts: Vec<u64> = parts[0].split('-').filter_map(|p| p.parse().ok()).collect();
|
||||
let time_str = parts[1].trim_end_matches('Z');
|
||||
let time_parts: Vec<u64> = time_str.split(':').filter_map(|p| p.parse().ok()).collect();
|
||||
if date_parts.len() != 3 || time_parts.len() != 3 {
|
||||
return 0;
|
||||
}
|
||||
let (y, m, d) = (date_parts[0], date_parts[1], date_parts[2]);
|
||||
let (hh, mm, ss) = (time_parts[0], time_parts[1], time_parts[2]);
|
||||
let y_adj = if m <= 2 { y - 1 } else { y };
|
||||
let era = y_adj / 400;
|
||||
let yoe = y_adj - era * 400;
|
||||
let m_adj = if m > 2 { m - 3 } else { m + 9 };
|
||||
let doy = (153 * m_adj + 2) / 5 + d - 1;
|
||||
let doe = yoe * 365 + yoe / 4 - yoe / 100 + doy;
|
||||
let days = era * 146097 + doe - 719468;
|
||||
(days * 86400 + hh * 3600 + mm * 60 + ss) as i64
|
||||
}
|
||||
|
||||
/// Check if the backend supports batch processing (OpenAI and Azure only).
|
||||
fn is_batch_supported(backend: &BackendClient) -> bool {
|
||||
matches!(
|
||||
@@ -306,6 +370,13 @@ fn is_batch_supported(backend: &BackendClient) -> bool {
|
||||
)
|
||||
}
|
||||
|
||||
fn is_openai_or_azure_backend(backend: &BackendClient) -> bool {
|
||||
matches!(
|
||||
backend,
|
||||
BackendClient::OpenAI(_) | BackendClient::AzureOpenAI(_)
|
||||
)
|
||||
}
|
||||
|
||||
fn bad_request(msg: &str) -> Response {
|
||||
let err = create_anthropic_error(
|
||||
anthropic::ErrorType::InvalidRequestError,
|
||||
@@ -333,3 +404,8 @@ fn internal_error(msg: &str) -> Response {
|
||||
let err = create_anthropic_error(anthropic::ErrorType::ApiError, msg.to_string(), None);
|
||||
(StatusCode::INTERNAL_SERVER_ERROR, Json(err)).into_response()
|
||||
}
|
||||
|
||||
fn not_found_response(msg: &str) -> Response {
|
||||
let err = create_anthropic_error(anthropic::ErrorType::NotFoundError, msg.to_string(), None);
|
||||
(StatusCode::NOT_FOUND, Json(err)).into_response()
|
||||
}
|
||||
|
||||
@@ -154,9 +154,7 @@ impl ToolStartupConfig {
|
||||
/// Returns true when at least one tool-related section was present in the config.
|
||||
/// Used to decide whether to construct a ToolEngineState at all.
|
||||
pub fn has_any(&self) -> bool {
|
||||
self.tool_execution.is_some()
|
||||
|| self.builtin_tools.is_some()
|
||||
|| self.mcp_servers.is_some()
|
||||
self.tool_execution.is_some() || self.builtin_tools.is_some() || self.mcp_servers.is_some()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -395,9 +393,7 @@ impl SimpleConfig {
|
||||
crate::tools::LoopConfig {
|
||||
max_iterations: te.max_iterations.unwrap_or(1),
|
||||
tool_timeout: std::time::Duration::from_secs(te.tool_timeout_secs.unwrap_or(30)),
|
||||
total_timeout: std::time::Duration::from_secs(
|
||||
te.total_timeout_secs.unwrap_or(300),
|
||||
),
|
||||
total_timeout: std::time::Duration::from_secs(te.total_timeout_secs.unwrap_or(300)),
|
||||
max_tool_calls_per_turn: te.max_tool_calls_per_turn.unwrap_or(16),
|
||||
}
|
||||
} else {
|
||||
@@ -903,7 +899,10 @@ mcp_servers:
|
||||
assert_eq!(rule.timeout, Some(std::time::Duration::from_secs(10)));
|
||||
|
||||
// MCP glob rule
|
||||
assert_eq!(policy.resolve("mcp_github_search_repos"), PolicyAction::Allow);
|
||||
assert_eq!(
|
||||
policy.resolve("mcp_github_search_repos"),
|
||||
PolicyAction::Allow
|
||||
);
|
||||
|
||||
// Default loop config (no tool_execution section)
|
||||
assert_eq!(loop_config.max_iterations, 1);
|
||||
@@ -922,10 +921,7 @@ tool_execution:
|
||||
let config: SimpleConfig = serde_yaml::from_str(yaml).unwrap();
|
||||
let (_policy, loop_config) = config.build_tool_config();
|
||||
assert_eq!(loop_config.max_iterations, 5);
|
||||
assert_eq!(
|
||||
loop_config.tool_timeout,
|
||||
std::time::Duration::from_secs(45)
|
||||
);
|
||||
assert_eq!(loop_config.tool_timeout, std::time::Duration::from_secs(45));
|
||||
assert_eq!(
|
||||
loop_config.total_timeout,
|
||||
std::time::Duration::from_secs(300)
|
||||
|
||||
+119
-78
@@ -160,96 +160,97 @@ async fn main() {
|
||||
// Build tool engine state from config, if tool sections were present.
|
||||
// Only constructed when at least one of tool_execution / builtin_tools / mcp_servers
|
||||
// is present in the config file, to avoid overhead when tools are unused.
|
||||
let tool_engine_state: Option<Arc<routes::ToolEngineState>> =
|
||||
if let Some(tc) = load_result.tool_config.filter(|tc| tc.has_any()) {
|
||||
let simple_config_shell = config::simple::SimpleConfig {
|
||||
routing_strategy: None,
|
||||
listen_port: None,
|
||||
log_bodies: None,
|
||||
models: vec![],
|
||||
tool_execution: tc.tool_execution,
|
||||
builtin_tools: tc.builtin_tools,
|
||||
mcp_servers: tc.mcp_servers,
|
||||
};
|
||||
let (policy, loop_config) = simple_config_shell.build_tool_config();
|
||||
let tool_engine_state: Option<Arc<routes::ToolEngineState>> = if let Some(tc) =
|
||||
load_result.tool_config.filter(|tc| tc.has_any())
|
||||
{
|
||||
let simple_config_shell = config::simple::SimpleConfig {
|
||||
routing_strategy: None,
|
||||
listen_port: None,
|
||||
log_bodies: None,
|
||||
models: vec![],
|
||||
tool_execution: tc.tool_execution,
|
||||
builtin_tools: tc.builtin_tools,
|
||||
mcp_servers: tc.mcp_servers,
|
||||
};
|
||||
let (policy, loop_config) = simple_config_shell.build_tool_config();
|
||||
|
||||
let mut registry = tools::ToolRegistry::new();
|
||||
// Register built-in tools (gated behind the dangerous-builtin-tools feature).
|
||||
anyllm_proxy::tools::builtin::register_all(
|
||||
&mut registry,
|
||||
simple_config_shell.builtin_tools.as_ref(),
|
||||
);
|
||||
let mut registry = tools::ToolRegistry::new();
|
||||
// Register built-in tools (gated behind the dangerous-builtin-tools feature).
|
||||
anyllm_proxy::tools::builtin::register_all(
|
||||
&mut registry,
|
||||
simple_config_shell.builtin_tools.as_ref(),
|
||||
);
|
||||
|
||||
// Build MCP manager and discover tools from configured servers.
|
||||
let mcp_manager = if let Some(ref servers) = simple_config_shell.mcp_servers {
|
||||
let manager = Arc::new(tools::McpServerManager::new());
|
||||
for server_cfg in servers {
|
||||
// SSRF protection: skip servers with private/loopback URLs.
|
||||
if let Err(e) = crate::config::validate_base_url(&server_cfg.url) {
|
||||
tracing::error!(
|
||||
// Build MCP manager and discover tools from configured servers.
|
||||
let mcp_manager = if let Some(ref servers) = simple_config_shell.mcp_servers {
|
||||
let manager = Arc::new(tools::McpServerManager::new());
|
||||
for server_cfg in servers {
|
||||
// SSRF protection: skip servers with private/loopback URLs.
|
||||
if let Err(e) = crate::config::validate_base_url(&server_cfg.url) {
|
||||
tracing::error!(
|
||||
server = %server_cfg.name,
|
||||
url = %server_cfg.url,
|
||||
error = %e,
|
||||
"MCP server URL rejected (SSRF protection); skipping"
|
||||
);
|
||||
continue;
|
||||
}
|
||||
match tools::McpServerManager::discover_tools(&server_cfg.url).await {
|
||||
Ok(discovered) => {
|
||||
tracing::info!(
|
||||
server = %server_cfg.name,
|
||||
url = %server_cfg.url,
|
||||
tools = discovered.len(),
|
||||
"MCP server connected and tools discovered"
|
||||
);
|
||||
if let Err(e) = manager.register_server_blocking(
|
||||
&server_cfg.name,
|
||||
&server_cfg.url,
|
||||
discovered,
|
||||
) {
|
||||
tracing::error!(
|
||||
server = %server_cfg.name,
|
||||
error = %e,
|
||||
"MCP server registration failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
server = %server_cfg.name,
|
||||
url = %server_cfg.url,
|
||||
error = %e,
|
||||
"MCP server URL rejected (SSRF protection); skipping"
|
||||
"MCP server unreachable at startup; tools from this server will be unavailable"
|
||||
);
|
||||
continue;
|
||||
}
|
||||
match tools::McpServerManager::discover_tools(&server_cfg.url).await {
|
||||
Ok(discovered) => {
|
||||
tracing::info!(
|
||||
server = %server_cfg.name,
|
||||
url = %server_cfg.url,
|
||||
tools = discovered.len(),
|
||||
"MCP server connected and tools discovered"
|
||||
);
|
||||
if let Err(e) = manager.register_server_blocking(
|
||||
&server_cfg.name,
|
||||
&server_cfg.url,
|
||||
discovered,
|
||||
) {
|
||||
tracing::error!(
|
||||
server = %server_cfg.name,
|
||||
error = %e,
|
||||
"MCP server registration failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
server = %server_cfg.name,
|
||||
url = %server_cfg.url,
|
||||
error = %e,
|
||||
"MCP server unreachable at startup; tools from this server will be unavailable"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
// Register all discovered MCP tools into the registry.
|
||||
tools::mcp::register_mcp_tools(&manager, &mut registry);
|
||||
Some(manager)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
tracing::info!(
|
||||
registered_tools = registry.list_names().len(),
|
||||
mcp_servers = mcp_manager
|
||||
.as_ref()
|
||||
.map(|m| m.list_servers_blocking().len())
|
||||
.unwrap_or(0),
|
||||
"tool execution engine initialized"
|
||||
);
|
||||
|
||||
Some(Arc::new(routes::ToolEngineState {
|
||||
registry: Arc::new(registry),
|
||||
policy: Arc::new(policy),
|
||||
loop_config,
|
||||
mcp_manager,
|
||||
}))
|
||||
}
|
||||
// Register all discovered MCP tools into the registry.
|
||||
tools::mcp::register_mcp_tools(&manager, &mut registry);
|
||||
Some(manager)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
tracing::info!(
|
||||
registered_tools = registry.list_names().len(),
|
||||
mcp_servers = mcp_manager
|
||||
.as_ref()
|
||||
.map(|m| m.list_servers_blocking().len())
|
||||
.unwrap_or(0),
|
||||
"tool execution engine initialized"
|
||||
);
|
||||
|
||||
Some(Arc::new(routes::ToolEngineState {
|
||||
registry: Arc::new(registry),
|
||||
policy: Arc::new(policy),
|
||||
loop_config,
|
||||
mcp_manager,
|
||||
}))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// Admin web UI is opt-in: pass --webui or --admin to enable.
|
||||
// DISABLE_ADMIN=1 overrides the flag (useful in container/scripted environments).
|
||||
let flag_set = args.iter().any(|a| a == "--webui" || a == "--admin");
|
||||
@@ -543,12 +544,52 @@ async fn main() {
|
||||
None
|
||||
};
|
||||
|
||||
// Initialize batch engine with its own connection to the same DB file.
|
||||
// Only available when admin is enabled (requires a DB path).
|
||||
let batch_engine: Option<
|
||||
std::sync::Arc<
|
||||
anyllm_batch_engine::BatchEngine<
|
||||
anyllm_batch_engine::queue::sqlite::SqliteQueue,
|
||||
anyllm_batch_engine::webhook::sqlite::SqliteWebhookQueue,
|
||||
>,
|
||||
>,
|
||||
> = if enable_admin {
|
||||
let db_path = std::env::var("ADMIN_DB_PATH").unwrap_or_else(|_| "admin.db".into());
|
||||
let batch_conn = rusqlite::Connection::open(&db_path)
|
||||
.expect("failed to open second SQLite connection for batch engine");
|
||||
anyllm_batch_engine::db::migrate_old_tables(&batch_conn)
|
||||
.expect("failed to migrate old batch tables");
|
||||
anyllm_batch_engine::db::init_batch_engine_tables(&batch_conn)
|
||||
.expect("failed to initialize batch engine tables");
|
||||
let batch_db = std::sync::Arc::new(tokio::sync::Mutex::new(batch_conn));
|
||||
let global_webhook_urls: Vec<String> = std::env::var("BATCH_WEBHOOK_URLS")
|
||||
.unwrap_or_default()
|
||||
.split(',')
|
||||
.map(|s| s.trim().to_string())
|
||||
.filter(|s| !s.is_empty())
|
||||
.collect();
|
||||
Some(std::sync::Arc::new(anyllm_batch_engine::BatchEngine {
|
||||
queue: std::sync::Arc::new(anyllm_batch_engine::queue::sqlite::SqliteQueue::new(
|
||||
batch_db.clone(),
|
||||
)),
|
||||
file_store: anyllm_batch_engine::file_store::FileStore::new(batch_db.clone()),
|
||||
webhook_queue: std::sync::Arc::new(
|
||||
anyllm_batch_engine::webhook::sqlite::SqliteWebhookQueue::new(batch_db),
|
||||
),
|
||||
global_webhook_urls,
|
||||
webhook_signing_secret: std::env::var("BATCH_WEBHOOK_SIGNING_SECRET").ok(),
|
||||
}))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// Build proxy router with optional shared admin state and tool engine.
|
||||
let app = routes::app_multi_with_shared(
|
||||
multi_config,
|
||||
admin_parts.as_ref().map(|(s, _, _)| s.clone()),
|
||||
model_router,
|
||||
tool_engine_state,
|
||||
batch_engine,
|
||||
);
|
||||
|
||||
// --- Start servers ---
|
||||
|
||||
@@ -258,11 +258,11 @@ pub(crate) async fn chat_completions(
|
||||
);
|
||||
oai_req.model = m;
|
||||
match c.chat_completion(&oai_req).await {
|
||||
Ok((resp, _, _)) => Ok(
|
||||
mapping::message_map::openai_to_anthropic_response(
|
||||
Ok((resp, _, _)) => {
|
||||
Ok(mapping::message_map::openai_to_anthropic_response(
|
||||
&resp, &om,
|
||||
),
|
||||
),
|
||||
))
|
||||
}
|
||||
Err(e) => Err(format!("{e}")),
|
||||
}
|
||||
}
|
||||
@@ -583,21 +583,29 @@ async fn chat_completions_stream(
|
||||
for tc in tc_list {
|
||||
let idx = tc.index as usize;
|
||||
while accumulated_tool_calls.len() <= idx {
|
||||
accumulated_tool_calls.push((String::new(), String::new(), String::new()));
|
||||
accumulated_tool_calls.push((
|
||||
String::new(),
|
||||
String::new(),
|
||||
String::new(),
|
||||
));
|
||||
}
|
||||
if let Some(ref id) = tc.id {
|
||||
if !id.is_empty() {
|
||||
accumulated_tool_calls[idx].0 = id.clone();
|
||||
accumulated_tool_calls[idx].0 =
|
||||
id.clone();
|
||||
}
|
||||
}
|
||||
if let Some(ref func) = tc.function {
|
||||
if let Some(ref name) = func.name {
|
||||
if !name.is_empty() {
|
||||
accumulated_tool_calls[idx].1 = name.clone();
|
||||
accumulated_tool_calls[idx].1 =
|
||||
name.clone();
|
||||
}
|
||||
}
|
||||
if let Some(ref args) = func.arguments {
|
||||
accumulated_tool_calls[idx].2.push_str(args);
|
||||
accumulated_tool_calls[idx]
|
||||
.2
|
||||
.push_str(args);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -707,8 +715,7 @@ async fn chat_completions_stream(
|
||||
&engine.registry,
|
||||
&engine.policy,
|
||||
);
|
||||
let denied_results =
|
||||
crate::tools::execution::denied_tool_results(&denied);
|
||||
let denied_results = crate::tools::execution::denied_tool_results(&denied);
|
||||
|
||||
if !auto_exec.is_empty() || !denied_results.is_empty() {
|
||||
let mut results = crate::tools::execution::execute_tool_calls(
|
||||
@@ -734,14 +741,14 @@ async fn chat_completions_stream(
|
||||
.collect();
|
||||
|
||||
let mut follow_up_req = anthropic_req_for_tools;
|
||||
follow_up_req
|
||||
.messages
|
||||
.push(anyllm_translate::anthropic::InputMessage {
|
||||
follow_up_req.messages.push(
|
||||
anyllm_translate::anthropic::InputMessage {
|
||||
role: anyllm_translate::anthropic::Role::Assistant,
|
||||
content: anyllm_translate::anthropic::Content::Blocks(
|
||||
assistant_content,
|
||||
),
|
||||
});
|
||||
},
|
||||
);
|
||||
follow_up_req.messages.push(
|
||||
crate::tools::execution::tool_results_to_user_message(&results),
|
||||
);
|
||||
@@ -781,25 +788,18 @@ async fn chat_completions_stream(
|
||||
let mut follow_buffer = BytesMut::new();
|
||||
let mut follow_search_from: usize = 0;
|
||||
|
||||
while let Some(chunk_result) =
|
||||
follow_byte_stream.next().await
|
||||
{
|
||||
while let Some(chunk_result) = follow_byte_stream.next().await {
|
||||
let bytes = match chunk_result {
|
||||
Ok(b) => b,
|
||||
Err(e) => {
|
||||
tracing::error!(
|
||||
"follow-up stream read error: {e}"
|
||||
);
|
||||
tracing::error!("follow-up stream read error: {e}");
|
||||
break;
|
||||
}
|
||||
};
|
||||
follow_buffer.extend_from_slice(&bytes);
|
||||
|
||||
while let Some((pos, delim_len)) =
|
||||
find_double_newline(
|
||||
&follow_buffer,
|
||||
follow_search_from,
|
||||
)
|
||||
find_double_newline(&follow_buffer, follow_search_from)
|
||||
{
|
||||
if let Ok(frame_str) =
|
||||
std::str::from_utf8(&follow_buffer[..pos])
|
||||
@@ -817,13 +817,11 @@ async fn chat_completions_stream(
|
||||
>(
|
||||
json_str
|
||||
) {
|
||||
let events =
|
||||
follow_stream_translator
|
||||
.process_chunk(&chunk);
|
||||
let events = follow_stream_translator
|
||||
.process_chunk(&chunk);
|
||||
for event in &events {
|
||||
let oai_chunks =
|
||||
follow_translator
|
||||
.process_event(event);
|
||||
let oai_chunks = follow_translator
|
||||
.process_event(event);
|
||||
for oai_chunk in &oai_chunks {
|
||||
if let Ok(json) =
|
||||
serde_json::to_string(
|
||||
@@ -847,23 +845,18 @@ async fn chat_completions_stream(
|
||||
}
|
||||
}
|
||||
}
|
||||
let _ =
|
||||
follow_buffer.split_to(pos + delim_len);
|
||||
let _ = follow_buffer.split_to(pos + delim_len);
|
||||
follow_search_from = 0;
|
||||
}
|
||||
follow_search_from =
|
||||
follow_buffer.len().saturating_sub(3);
|
||||
follow_search_from = follow_buffer.len().saturating_sub(3);
|
||||
}
|
||||
|
||||
// Emit finish events for the follow-up stream.
|
||||
let follow_finish = follow_stream_translator.finish();
|
||||
for event in &follow_finish {
|
||||
let oai_chunks =
|
||||
follow_translator.process_event(event);
|
||||
let oai_chunks = follow_translator.process_event(event);
|
||||
for oai_chunk in &oai_chunks {
|
||||
if let Ok(json) =
|
||||
serde_json::to_string(oai_chunk)
|
||||
{
|
||||
if let Ok(json) = serde_json::to_string(oai_chunk) {
|
||||
let _ = tx
|
||||
.send(Ok(format!("data: {}\n\n", json)))
|
||||
.await;
|
||||
|
||||
@@ -107,6 +107,15 @@ pub struct AppState {
|
||||
pub all_backends: Option<Arc<HashMap<String, AppState>>>,
|
||||
/// Tool execution engine state. None when tool execution is not configured.
|
||||
pub tool_engine: Option<Arc<ToolEngineState>>,
|
||||
/// Batch orchestration engine. None in test configs that don't need batch.
|
||||
pub batch_engine: Option<
|
||||
Arc<
|
||||
anyllm_batch_engine::BatchEngine<
|
||||
anyllm_batch_engine::queue::sqlite::SqliteQueue,
|
||||
anyllm_batch_engine::webhook::sqlite::SqliteWebhookQueue,
|
||||
>,
|
||||
>,
|
||||
>,
|
||||
}
|
||||
|
||||
impl AppState {
|
||||
@@ -207,7 +216,7 @@ pub fn app(config: Config) -> Router {
|
||||
/// Build the axum router from multi-backend configuration.
|
||||
/// Creates nested sub-routers for each configured backend.
|
||||
pub fn app_multi(config: MultiConfig) -> Router {
|
||||
app_multi_with_shared(config, None, None, None)
|
||||
app_multi_with_shared(config, None, None, None, None)
|
||||
}
|
||||
|
||||
/// Build the axum router with optional shared admin state and model router.
|
||||
@@ -216,6 +225,14 @@ pub fn app_multi_with_shared(
|
||||
shared: Option<SharedState>,
|
||||
model_router: Option<Arc<RwLock<crate::config::model_router::ModelRouter>>>,
|
||||
tool_engine: Option<Arc<ToolEngineState>>,
|
||||
batch_engine: Option<
|
||||
Arc<
|
||||
anyllm_batch_engine::BatchEngine<
|
||||
anyllm_batch_engine::queue::sqlite::SqliteQueue,
|
||||
anyllm_batch_engine::webhook::sqlite::SqliteWebhookQueue,
|
||||
>,
|
||||
>,
|
||||
>,
|
||||
) -> Router {
|
||||
let mut backend_metrics: HashMap<String, Metrics> = HashMap::new();
|
||||
let mut router = Router::new();
|
||||
@@ -269,6 +286,7 @@ pub fn app_multi_with_shared(
|
||||
// all_backends is set after the loop (needs all states built first).
|
||||
all_backends: None,
|
||||
tool_engine: tool_engine.clone(),
|
||||
batch_engine: batch_engine.clone(),
|
||||
};
|
||||
let sub = backend_router(state.clone(), mode);
|
||||
backend_states.insert(name.clone(), (state, mode));
|
||||
@@ -393,6 +411,10 @@ fn backend_router(state: AppState, mode: HandlerMode) -> Router<GlobalState> {
|
||||
.route(
|
||||
"/v1/batches/{batch_id}",
|
||||
get(crate::batch::routes::get_batch),
|
||||
)
|
||||
.route(
|
||||
"/v1/batches/{batch_id}/cancel",
|
||||
post(crate::batch::routes::cancel_batch),
|
||||
);
|
||||
|
||||
let api_routes = match mode {
|
||||
|
||||
@@ -52,21 +52,15 @@ impl Tool for BashTool {
|
||||
.stderr(Stdio::piped())
|
||||
.output();
|
||||
|
||||
let output = match tokio::time::timeout(
|
||||
std::time::Duration::from_secs(TIMEOUT_SECS),
|
||||
fut,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Ok(output)) => output,
|
||||
Ok(Err(e)) => return Err(format!("Failed to spawn process: {}", e)),
|
||||
Err(_) => {
|
||||
return Err(format!(
|
||||
"Command timed out after {} seconds",
|
||||
TIMEOUT_SECS
|
||||
))
|
||||
}
|
||||
};
|
||||
let output =
|
||||
match tokio::time::timeout(std::time::Duration::from_secs(TIMEOUT_SECS), fut).await
|
||||
{
|
||||
Ok(Ok(output)) => output,
|
||||
Ok(Err(e)) => return Err(format!("Failed to spawn process: {}", e)),
|
||||
Err(_) => {
|
||||
return Err(format!("Command timed out after {} seconds", TIMEOUT_SECS))
|
||||
}
|
||||
};
|
||||
|
||||
let stdout = String::from_utf8_lossy(&output.stdout);
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
|
||||
@@ -6,7 +6,7 @@ use tokio::task::JoinSet;
|
||||
|
||||
use crate::tools::policy::{PolicyAction, ToolExecutionPolicy};
|
||||
use crate::tools::registry::ToolRegistry;
|
||||
use crate::tools::trace::{ToolOutcome};
|
||||
use crate::tools::trace::ToolOutcome;
|
||||
|
||||
/// A tool call extracted from an LLM response.
|
||||
#[derive(Debug, Clone)]
|
||||
@@ -113,11 +113,8 @@ pub async fn execute_tool_calls(
|
||||
let input = call.input.clone();
|
||||
|
||||
join_set.spawn(async move {
|
||||
let result = tokio::time::timeout(
|
||||
timeout,
|
||||
execute_single(®istry, &name, input),
|
||||
)
|
||||
.await;
|
||||
let result =
|
||||
tokio::time::timeout(timeout, execute_single(®istry, &name, input)).await;
|
||||
|
||||
let outcome = match result {
|
||||
Ok(Ok(value)) => ToolOutcome::Success(value),
|
||||
|
||||
@@ -1,8 +1,8 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::{Arc, RwLock};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use anyllm_client::http::{build_http_client, HttpClientConfig};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::{Arc, RwLock};
|
||||
|
||||
/// An MCP tool definition discovered from a server.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
@@ -109,7 +109,11 @@ impl McpServerManager {
|
||||
}
|
||||
|
||||
pub fn find_server_for_tool_blocking(&self, prefixed_name: &str) -> Option<String> {
|
||||
self.tool_to_server.read().unwrap().get(prefixed_name).cloned()
|
||||
self.tool_to_server
|
||||
.read()
|
||||
.unwrap()
|
||||
.get(prefixed_name)
|
||||
.cloned()
|
||||
}
|
||||
|
||||
pub fn as_anthropic_tools_blocking(&self) -> Vec<anyllm_translate::anthropic::Tool> {
|
||||
@@ -201,49 +205,51 @@ impl McpServerManager {
|
||||
}
|
||||
}
|
||||
|
||||
async fn discover_tools_impl(client: &reqwest::Client, url: &str) -> Result<Vec<McpToolDef>, String> {
|
||||
let rpc_request = serde_json::json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": 1,
|
||||
"method": "tools/list",
|
||||
"params": {}
|
||||
});
|
||||
async fn discover_tools_impl(
|
||||
client: &reqwest::Client,
|
||||
url: &str,
|
||||
) -> Result<Vec<McpToolDef>, String> {
|
||||
let rpc_request = serde_json::json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": 1,
|
||||
"method": "tools/list",
|
||||
"params": {}
|
||||
});
|
||||
|
||||
let response = client
|
||||
.post(url)
|
||||
.json(&rpc_request)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("MCP discovery failed for '{}': {}", url, e))?;
|
||||
let response = client
|
||||
.post(url)
|
||||
.json(&rpc_request)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("MCP discovery failed for '{}': {}", url, e))?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
return Err(format!(
|
||||
"MCP discovery returned status {} for '{}'",
|
||||
response.status(),
|
||||
url
|
||||
));
|
||||
}
|
||||
if !response.status().is_success() {
|
||||
return Err(format!(
|
||||
"MCP discovery returned status {} for '{}'",
|
||||
response.status(),
|
||||
url
|
||||
));
|
||||
}
|
||||
|
||||
let body: Value = response
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| format!("MCP discovery parse error: {}", e))?;
|
||||
let body: Value = response
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| format!("MCP discovery parse error: {}", e))?;
|
||||
|
||||
if let Some(error) = body.get("error") {
|
||||
let msg = error
|
||||
.get("message")
|
||||
.and_then(|m| m.as_str())
|
||||
.unwrap_or("unknown error");
|
||||
return Err(format!("MCP discovery error: {}", msg));
|
||||
}
|
||||
if let Some(error) = body.get("error") {
|
||||
let msg = error
|
||||
.get("message")
|
||||
.and_then(|m| m.as_str())
|
||||
.unwrap_or("unknown error");
|
||||
return Err(format!("MCP discovery error: {}", msg));
|
||||
}
|
||||
|
||||
let tools_value = body
|
||||
.get("result")
|
||||
.and_then(|r| r.get("tools"))
|
||||
.ok_or_else(|| "MCP response missing result.tools".to_string())?;
|
||||
let tools_value = body
|
||||
.get("result")
|
||||
.and_then(|r| r.get("tools"))
|
||||
.ok_or_else(|| "MCP response missing result.tools".to_string())?;
|
||||
|
||||
serde_json::from_value(tools_value.clone())
|
||||
.map_err(|e| format!("MCP tools parse error: {}", e))
|
||||
serde_json::from_value(tools_value.clone()).map_err(|e| format!("MCP tools parse error: {}", e))
|
||||
}
|
||||
|
||||
impl Default for McpServerManager {
|
||||
@@ -284,7 +290,10 @@ impl crate::tools::registry::Tool for McpToolAdapter {
|
||||
}
|
||||
|
||||
/// Register all MCP tools from the manager into a ToolRegistry.
|
||||
pub fn register_mcp_tools(manager: &Arc<McpServerManager>, registry: &mut crate::tools::ToolRegistry) {
|
||||
pub fn register_mcp_tools(
|
||||
manager: &Arc<McpServerManager>,
|
||||
registry: &mut crate::tools::ToolRegistry,
|
||||
) {
|
||||
let servers = manager.list_servers_blocking();
|
||||
for server in &servers {
|
||||
for tool in &server.tools {
|
||||
@@ -340,7 +349,9 @@ mod tests {
|
||||
mgr.find_server_for_tool_blocking("mcp_github_create_issue"),
|
||||
Some("github".to_string())
|
||||
);
|
||||
assert!(mgr.find_server_for_tool_blocking("mcp_slack_send").is_none());
|
||||
assert!(mgr
|
||||
.find_server_for_tool_blocking("mcp_slack_send")
|
||||
.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -24,20 +24,37 @@ fn test_config() -> Config {
|
||||
}
|
||||
}
|
||||
|
||||
async fn make_test_batch_engine() -> std::sync::Arc<
|
||||
anyllm_batch_engine::BatchEngine<
|
||||
anyllm_batch_engine::queue::sqlite::SqliteQueue,
|
||||
anyllm_batch_engine::webhook::sqlite::SqliteWebhookQueue,
|
||||
>,
|
||||
> {
|
||||
use anyllm_batch_engine::{
|
||||
db::init_batch_engine_tables, file_store::FileStore, queue::sqlite::SqliteQueue,
|
||||
webhook::sqlite::SqliteWebhookQueue, BatchEngine,
|
||||
};
|
||||
let conn = rusqlite::Connection::open_in_memory().unwrap();
|
||||
init_batch_engine_tables(&conn).unwrap();
|
||||
let db = std::sync::Arc::new(tokio::sync::Mutex::new(conn));
|
||||
std::sync::Arc::new(BatchEngine {
|
||||
queue: std::sync::Arc::new(SqliteQueue::new(db.clone())),
|
||||
file_store: FileStore::new(db.clone()),
|
||||
webhook_queue: std::sync::Arc::new(SqliteWebhookQueue::new(db)),
|
||||
global_webhook_urls: vec![],
|
||||
webhook_signing_secret: None,
|
||||
})
|
||||
}
|
||||
|
||||
/// Spawn a test server with SharedState (needed for batch DB access).
|
||||
async fn spawn_test_server_with_shared() -> String {
|
||||
std::env::set_var("PROXY_OPEN_RELAY", "true");
|
||||
let config = test_config();
|
||||
let multi = MultiConfig::from_single_config(&config);
|
||||
let shared = admin::state::SharedState::new_for_test();
|
||||
let engine = make_test_batch_engine().await;
|
||||
|
||||
// Initialize batch tables in the test DB
|
||||
{
|
||||
let conn = shared.db.lock().unwrap();
|
||||
anyllm_proxy::batch::db::init_batch_tables(&conn).unwrap();
|
||||
}
|
||||
|
||||
let app = routes::app_multi_with_shared(multi, Some(shared), None, None);
|
||||
let app = routes::app_multi_with_shared(multi, Some(shared), None, None, Some(engine));
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
|
||||
@@ -94,7 +111,8 @@ async fn upload_file_and_create_batch() {
|
||||
|
||||
let batch_obj: serde_json::Value = resp.json().await.unwrap();
|
||||
assert_eq!(batch_obj["object"], "batch");
|
||||
assert!(batch_obj["id"].as_str().unwrap().starts_with("batch-"));
|
||||
// Engine generates batch IDs (e.g. "batch_<uuid>")
|
||||
assert!(batch_obj["id"].as_str().unwrap().starts_with("batch"));
|
||||
assert_eq!(batch_obj["status"], "validating");
|
||||
assert_eq!(batch_obj["input_file_id"], file_id);
|
||||
assert_eq!(batch_obj["request_counts"]["total"], 2);
|
||||
@@ -200,11 +218,8 @@ async fn unsupported_backend_returns_501() {
|
||||
|
||||
let multi = MultiConfig::from_single_config(&config);
|
||||
let shared = admin::state::SharedState::new_for_test();
|
||||
{
|
||||
let conn = shared.db.lock().unwrap();
|
||||
anyllm_proxy::batch::db::init_batch_tables(&conn).unwrap();
|
||||
}
|
||||
let app = routes::app_multi_with_shared(multi, Some(shared), None, None);
|
||||
let engine = make_test_batch_engine().await;
|
||||
let app = routes::app_multi_with_shared(multi, Some(shared), None, None, Some(engine));
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
|
||||
@@ -225,6 +240,62 @@ async fn unsupported_backend_returns_501() {
|
||||
assert_eq!(resp.status(), 501);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cancel_queued_batch() {
|
||||
let base = spawn_test_server_with_shared().await;
|
||||
let client = Client::new();
|
||||
|
||||
// Upload a file first
|
||||
let form = multipart::Form::new().text("purpose", "batch").part(
|
||||
"file",
|
||||
multipart::Part::bytes(valid_jsonl().as_bytes().to_vec())
|
||||
.file_name("test.jsonl")
|
||||
.mime_str("application/jsonl")
|
||||
.unwrap(),
|
||||
);
|
||||
let resp = client
|
||||
.post(format!("{base}/v1/files"))
|
||||
.header("x-api-key", "test")
|
||||
.multipart(form)
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(resp.status(), 200);
|
||||
let file_obj: serde_json::Value = resp.json().await.unwrap();
|
||||
let file_id = file_obj["id"].as_str().unwrap().to_string();
|
||||
|
||||
// Create a batch
|
||||
let resp = client
|
||||
.post(format!("{base}/v1/batches"))
|
||||
.header("x-api-key", "test")
|
||||
.json(&serde_json::json!({
|
||||
"input_file_id": file_id,
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"completion_window": "24h"
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(resp.status(), 200);
|
||||
let batch_obj: serde_json::Value = resp.json().await.unwrap();
|
||||
let batch_id = batch_obj["id"].as_str().unwrap().to_string();
|
||||
|
||||
// Cancel the batch
|
||||
let resp = client
|
||||
.post(format!("{base}/v1/batches/{batch_id}/cancel"))
|
||||
.header("x-api-key", "test")
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(resp.status(), 200);
|
||||
let cancelled: serde_json::Value = resp.json().await.unwrap();
|
||||
let status = cancelled["status"].as_str().unwrap();
|
||||
assert!(
|
||||
status == "cancelling" || status == "cancelled",
|
||||
"expected cancelling or cancelled, got {status}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn anthropic_batch_rejects_empty_requests() {
|
||||
let base = spawn_test_server_with_shared().await;
|
||||
|
||||
@@ -7,14 +7,14 @@ use std::future::Future;
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyllm_proxy::tools::{
|
||||
PolicyAction, PolicyRule, Tool, ToolCall, ToolExecutionPolicy, ToolRegistry, ToolResult,
|
||||
};
|
||||
use anyllm_proxy::tools::execution::{
|
||||
denied_tool_results, execute_tool_calls, extract_tool_calls, is_duplicate, maybe_execute_tools,
|
||||
partition_tool_calls, tool_results_to_user_message, LoopConfig, ToolEngineState,
|
||||
};
|
||||
use anyllm_proxy::tools::trace::ToolOutcome;
|
||||
use anyllm_proxy::tools::{
|
||||
PolicyAction, PolicyRule, Tool, ToolCall, ToolExecutionPolicy, ToolRegistry, ToolResult,
|
||||
};
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Test tool: uppercases the "text" field of the input
|
||||
@@ -231,14 +231,12 @@ fn tool_results_error_outcome_sets_is_error_true() {
|
||||
}];
|
||||
let msg = tool_results_to_user_message(&results);
|
||||
match &msg.content {
|
||||
anyllm_translate::anthropic::Content::Blocks(blocks) => {
|
||||
match &blocks[0] {
|
||||
anyllm_translate::anthropic::ContentBlock::ToolResult { is_error, .. } => {
|
||||
assert_eq!(*is_error, Some(true));
|
||||
}
|
||||
other => panic!("expected ToolResult block, got {:?}", other),
|
||||
anyllm_translate::anthropic::Content::Blocks(blocks) => match &blocks[0] {
|
||||
anyllm_translate::anthropic::ContentBlock::ToolResult { is_error, .. } => {
|
||||
assert_eq!(*is_error, Some(true));
|
||||
}
|
||||
}
|
||||
other => panic!("expected ToolResult block, got {:?}", other),
|
||||
},
|
||||
other => panic!("expected Blocks content, got {:?}", other),
|
||||
}
|
||||
}
|
||||
@@ -247,10 +245,20 @@ fn tool_results_error_outcome_sets_is_error_true() {
|
||||
fn duplicate_detection_works() {
|
||||
let a = vec![make_call("1", "upper", serde_json::json!({"text": "same"}))];
|
||||
let b = vec![make_call("2", "upper", serde_json::json!({"text": "same"}))];
|
||||
assert!(is_duplicate(&a, &b), "same name+input with different IDs should be duplicate");
|
||||
assert!(
|
||||
is_duplicate(&a, &b),
|
||||
"same name+input with different IDs should be duplicate"
|
||||
);
|
||||
|
||||
let c = vec![make_call("3", "upper", serde_json::json!({"text": "different"}))];
|
||||
assert!(!is_duplicate(&a, &c), "different input should not be duplicate");
|
||||
let c = vec![make_call(
|
||||
"3",
|
||||
"upper",
|
||||
serde_json::json!({"text": "different"}),
|
||||
)];
|
||||
assert!(
|
||||
!is_duplicate(&a, &c),
|
||||
"different input should not be duplicate"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -454,7 +454,7 @@ async fn spawn_mock_backend() -> String {
|
||||
async fn spawn_proxy_with_shared_vk(config: Config) -> String {
|
||||
let state = shared_state(); // must call before building app to ensure set_virtual_keys fires
|
||||
let multi = anyllm_proxy::config::MultiConfig::from_single_config(&config);
|
||||
let base_app = routes::app_multi_with_shared(multi, Some(state), None, None);
|
||||
let base_app = routes::app_multi_with_shared(multi, Some(state), None, None, None);
|
||||
|
||||
// Add a test /admin/ route behind the same auth middleware so RBAC can be tested.
|
||||
let admin_test = Router::new()
|
||||
|
||||
@@ -0,0 +1,630 @@
|
||||
# Security Audit Fixes Implementation Plan
|
||||
|
||||
> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.
|
||||
|
||||
**Goal:** Fix six security vulnerabilities found in the security audit: arbitrary file read, MCP name ambiguity, add_model input injection, DNS rebinding on MCP client, CSRF improvement, and bash tool hardening.
|
||||
|
||||
**Architecture:** All fixes are targeted single-file or two-file changes. No new crates or dependencies required. The read_file fix threads allowed_dirs config through existing config structs. The MCP fix swaps `reqwest::Client::new()` for `build_http_client`. The rest are validation guards.
|
||||
|
||||
**Tech Stack:** Rust stable, tokio, reqwest, axum, anyllm_client::http
|
||||
|
||||
---
|
||||
|
||||
## Task 1: Fix read_file path confinement (Vuln 2 — High)
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/proxy/src/config/simple.rs` (add `allowed_dirs` to `BuiltinToolConfig`)
|
||||
- Modify: `crates/proxy/src/tools/builtin/read_file.rs` (add allowed_dirs field, enforce check)
|
||||
- Modify: `crates/proxy/src/tools/builtin/mod.rs` (pass config to register_all)
|
||||
- Modify: `crates/proxy/src/main.rs` (pass builtin config to register_all)
|
||||
|
||||
- [ ] **Step 1: Add `allowed_dirs` to BuiltinToolConfig in `crates/proxy/src/config/simple.rs`**
|
||||
|
||||
Find the `BuiltinToolConfig` struct (around line 114) and add the field:
|
||||
|
||||
```rust
|
||||
/// Configuration for a single builtin tool.
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct BuiltinToolConfig {
|
||||
#[serde(default = "default_true")]
|
||||
pub enabled: bool,
|
||||
#[serde(default)]
|
||||
pub policy: Option<String>,
|
||||
#[serde(default)]
|
||||
pub timeout_secs: Option<u64>,
|
||||
/// For read_file: restrict reads to these absolute directory paths.
|
||||
/// If empty or absent, all paths are permitted (dangerous; set this).
|
||||
#[serde(default)]
|
||||
pub allowed_dirs: Vec<String>,
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Add allowed_dirs field to ReadFileTool and enforce it**
|
||||
|
||||
Replace the entire `crates/proxy/src/tools/builtin/read_file.rs`:
|
||||
|
||||
```rust
|
||||
use crate::tools::registry::Tool;
|
||||
use serde_json::Value;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
/// Maximum file size to read (1 MB). Prevents OOM from huge files.
|
||||
const MAX_FILE_SIZE: u64 = 1024 * 1024;
|
||||
|
||||
/// Tool for reading file contents safely.
|
||||
pub struct ReadFileTool {
|
||||
/// If non-empty, restrict reads to files under these directories.
|
||||
/// All entries must be canonical absolute paths (no symlinks, no ..).
|
||||
pub allowed_dirs: Vec<PathBuf>,
|
||||
}
|
||||
|
||||
impl Tool for ReadFileTool {
|
||||
fn name(&self) -> &str {
|
||||
"read_file"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Reads the contents of a local file and returns it as a string."
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> Value {
|
||||
serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"path": {
|
||||
"type": "string",
|
||||
"description": "The absolute path to the file to read."
|
||||
}
|
||||
},
|
||||
"required": ["path"]
|
||||
})
|
||||
}
|
||||
|
||||
fn execute<'a>(
|
||||
&'a self,
|
||||
input: Value,
|
||||
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<Value, String>> + Send + 'a>>
|
||||
{
|
||||
let allowed_dirs = self.allowed_dirs.clone();
|
||||
Box::pin(async move {
|
||||
let raw_path = input
|
||||
.get("path")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| "Missing 'path' argument".to_string())?;
|
||||
|
||||
let path = Path::new(raw_path);
|
||||
|
||||
// Require absolute paths to prevent relative path traversal.
|
||||
if !path.is_absolute() {
|
||||
return Err("Only absolute paths are allowed".to_string());
|
||||
}
|
||||
|
||||
// Resolve symlinks and .. components.
|
||||
let canonical = path
|
||||
.canonicalize()
|
||||
.map_err(|e| format!("Cannot resolve path '{}': {}", raw_path, e))?;
|
||||
|
||||
// Enforce allowed_dirs allowlist. If configured, the canonical path
|
||||
// must start with at least one of the allowed directory prefixes.
|
||||
if !allowed_dirs.is_empty() {
|
||||
let permitted = allowed_dirs
|
||||
.iter()
|
||||
.any(|base| canonical.starts_with(base));
|
||||
if !permitted {
|
||||
return Err(format!(
|
||||
"Path '{}' is outside the configured allowed directories",
|
||||
raw_path
|
||||
));
|
||||
}
|
||||
} else {
|
||||
// No allowed_dirs configured: log a warning. Operator should set this.
|
||||
tracing::warn!(
|
||||
path = %raw_path,
|
||||
"read_file executed with no allowed_dirs restriction; \
|
||||
set allowed_dirs in builtin_tools config to restrict access"
|
||||
);
|
||||
}
|
||||
|
||||
// Check file size before reading.
|
||||
let metadata = std::fs::metadata(&canonical)
|
||||
.map_err(|e| format!("Cannot stat '{}': {}", raw_path, e))?;
|
||||
if metadata.len() > MAX_FILE_SIZE {
|
||||
return Err(format!(
|
||||
"File is {} bytes, exceeds {} byte limit",
|
||||
metadata.len(),
|
||||
MAX_FILE_SIZE
|
||||
));
|
||||
}
|
||||
|
||||
match std::fs::read_to_string(&canonical) {
|
||||
Ok(content) => Ok(serde_json::json!({ "content": content })),
|
||||
Err(e) => Err(format!("Failed to read file '{}': {}", raw_path, e)),
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 3: Update `register_all` to accept config and pass allowed_dirs**
|
||||
|
||||
Replace `crates/proxy/src/tools/builtin/mod.rs`:
|
||||
|
||||
```rust
|
||||
// SAFETY: The tools in this module execute arbitrary shell commands (BashTool)
|
||||
// and read arbitrary files (ReadFileTool) as the proxy process user. They must
|
||||
// NEVER be registered in a server-side tool registry without explicit operator
|
||||
// opt-in and appropriate sandboxing. Gated behind the `dangerous-builtin-tools`
|
||||
// feature flag, which is OFF by default.
|
||||
|
||||
#[cfg(feature = "dangerous-builtin-tools")]
|
||||
pub mod bash;
|
||||
#[cfg(feature = "dangerous-builtin-tools")]
|
||||
pub mod read_file;
|
||||
|
||||
use crate::tools::registry::ToolRegistry;
|
||||
|
||||
/// Populate a registry with standard built-in tools.
|
||||
///
|
||||
/// `builtin_configs`: map of tool name -> config from PROXY_CONFIG; used to pass
|
||||
/// per-tool settings (e.g., `allowed_dirs` for `read_file`) to tool constructors.
|
||||
///
|
||||
/// # Safety
|
||||
///
|
||||
/// When `dangerous-builtin-tools` is enabled, this registers `BashTool` (arbitrary
|
||||
/// shell execution) and `ReadFileTool` (arbitrary file reads). Only call this if
|
||||
/// the tool execution engine is sandboxed or if the operator has explicitly opted in.
|
||||
///
|
||||
/// When the feature is disabled (the default), this is a no-op.
|
||||
pub fn register_all(
|
||||
_registry: &mut ToolRegistry,
|
||||
_builtin_configs: Option<&std::collections::HashMap<String, crate::config::simple::BuiltinToolConfig>>,
|
||||
) {
|
||||
#[cfg(feature = "dangerous-builtin-tools")]
|
||||
{
|
||||
_registry.register(Box::new(bash::BashTool));
|
||||
|
||||
// Build ReadFileTool with allowed_dirs from config, if present.
|
||||
let allowed_dirs = _builtin_configs
|
||||
.and_then(|m| m.get("read_file"))
|
||||
.map(|cfg| {
|
||||
cfg.allowed_dirs
|
||||
.iter()
|
||||
.filter_map(|d| {
|
||||
let p = std::path::PathBuf::from(d);
|
||||
// Canonicalize at registration time so we compare canonical paths.
|
||||
match p.canonicalize() {
|
||||
Ok(canon) => Some(canon),
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
dir = %d,
|
||||
error = %e,
|
||||
"read_file allowed_dirs entry could not be canonicalized; skipping"
|
||||
);
|
||||
None
|
||||
}
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
|
||||
_registry.register(Box::new(read_file::ReadFileTool { allowed_dirs }));
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 4: Update the call site in `crates/proxy/src/main.rs`**
|
||||
|
||||
Find the line `anyllm_proxy::tools::builtin::register_all(&mut registry);` (around line 178) and change it to:
|
||||
|
||||
```rust
|
||||
anyllm_proxy::tools::builtin::register_all(
|
||||
&mut registry,
|
||||
simple_config_shell.builtin_tools.as_ref(),
|
||||
);
|
||||
```
|
||||
|
||||
- [ ] **Step 5: Run tests to verify nothing breaks**
|
||||
|
||||
```bash
|
||||
cd /Users/whit3rabbit/Documents/GitHub/llm-translate-api
|
||||
cargo test -p anyllm_proxy 2>&1 | tail -20
|
||||
```
|
||||
|
||||
Expected: all tests pass, no compile errors.
|
||||
|
||||
- [ ] **Step 6: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/proxy/src/config/simple.rs \
|
||||
crates/proxy/src/tools/builtin/read_file.rs \
|
||||
crates/proxy/src/tools/builtin/mod.rs \
|
||||
crates/proxy/src/main.rs
|
||||
git commit -m "fix(tools): enforce allowed_dirs allowlist in ReadFileTool
|
||||
|
||||
Adds allowed_dirs config field to BuiltinToolConfig. ReadFileTool now
|
||||
rejects reads outside the configured base directories after canonicalize().
|
||||
Logs a warning when allowed_dirs is empty. Threads config through
|
||||
register_all so tool constructors receive per-tool settings."
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Task 2: Fix MCP server name ambiguity (Vuln 3 — High)
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/proxy/src/tools/mcp.rs` (validate server name, reject underscores)
|
||||
|
||||
- [ ] **Step 1: Add `is_valid_mcp_server_name` and return Result from `register_server_blocking`**
|
||||
|
||||
In `crates/proxy/src/tools/mcp.rs`, add a validation function and change the signature:
|
||||
|
||||
```rust
|
||||
/// Validate MCP server name: must be non-empty, alphanumeric + hyphens only.
|
||||
/// Underscores are forbidden because the tool name scheme uses `mcp_{server}_{tool}`;
|
||||
/// an underscore in the server name makes `parse_mcp_tool_name` ambiguous.
|
||||
pub fn is_valid_mcp_server_name(name: &str) -> bool {
|
||||
!name.is_empty()
|
||||
&& name
|
||||
.chars()
|
||||
.all(|c| c.is_alphanumeric() || c == '-')
|
||||
}
|
||||
```
|
||||
|
||||
Change `register_server_blocking` to return `Result<(), String>`:
|
||||
|
||||
```rust
|
||||
pub fn register_server_blocking(
|
||||
&self,
|
||||
name: &str,
|
||||
url: &str,
|
||||
tools: Vec<McpToolDef>,
|
||||
) -> Result<(), String> {
|
||||
if !is_valid_mcp_server_name(name) {
|
||||
return Err(format!(
|
||||
"invalid MCP server name '{}': only alphanumerics and hyphens allowed",
|
||||
name
|
||||
));
|
||||
}
|
||||
self.remove_server_blocking(name);
|
||||
let mut tool_map = self.tool_to_server.write().unwrap();
|
||||
for tool in &tools {
|
||||
tool_map.insert(mcp_tool_name(name, &tool.name), name.to_string());
|
||||
}
|
||||
let server = McpServer {
|
||||
name: name.to_string(),
|
||||
url: url.to_string(),
|
||||
tools,
|
||||
};
|
||||
self.servers.write().unwrap().insert(name.to_string(), server);
|
||||
Ok(())
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Update all callers of `register_server_blocking`**
|
||||
|
||||
In `crates/proxy/src/main.rs` (around line 187-200), the call to `manager.register_server_blocking(...)` must handle the Result:
|
||||
|
||||
```rust
|
||||
if let Err(e) = manager.register_server_blocking(
|
||||
&server_cfg.name,
|
||||
&server_cfg.url,
|
||||
tools,
|
||||
) {
|
||||
tracing::error!(
|
||||
server = %server_cfg.name,
|
||||
error = %e,
|
||||
"MCP server registration failed"
|
||||
);
|
||||
continue;
|
||||
}
|
||||
```
|
||||
|
||||
In `crates/proxy/src/admin/routes.rs`, the handler for `POST /admin/api/mcp-servers` also calls `register_server_blocking`. Find it and handle the Result:
|
||||
|
||||
```rust
|
||||
if let Err(e) = mcp_manager.register_server_blocking(&body.name, &body.url, tools) {
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(serde_json::json!({"error": e})),
|
||||
).into_response();
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 3: Update tests in mcp.rs that call register_server_blocking**
|
||||
|
||||
Find all test calls like `mgr.register_server_blocking("github", ...)` and append `.unwrap()` or `.expect(...)` since they use valid names and should succeed.
|
||||
|
||||
- [ ] **Step 4: Run tests**
|
||||
|
||||
```bash
|
||||
cargo test -p anyllm_proxy 2>&1 | tail -20
|
||||
```
|
||||
|
||||
Expected: all tests pass.
|
||||
|
||||
- [ ] **Step 5: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/proxy/src/tools/mcp.rs crates/proxy/src/main.rs crates/proxy/src/admin/routes.rs
|
||||
git commit -m "fix(mcp): validate server names to prevent tool routing ambiguity
|
||||
|
||||
MCP tool names use mcp_{server}_{tool} scheme; underscores in server names
|
||||
cause parse_mcp_tool_name to misroute calls. is_valid_mcp_server_name now
|
||||
rejects names containing underscores. register_server_blocking returns
|
||||
Result<(), String> so callers can handle invalid names at registration time."
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Task 3: Fix add_model input validation (Vuln 5 — Medium)
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/proxy/src/admin/routes.rs` (add is_safe_model_name checks to add_model)
|
||||
|
||||
- [ ] **Step 1: Add validation at the top of the `add_model` handler**
|
||||
|
||||
In `crates/proxy/src/admin/routes.rs`, find the `add_model` function (around line 1449) and add validation after the model router guard:
|
||||
|
||||
```rust
|
||||
async fn add_model(
|
||||
ConnectInfo(addr): ConnectInfo<SocketAddr>,
|
||||
State(shared): State<SharedState>,
|
||||
Json(body): Json<AddModelRequest>,
|
||||
) -> impl IntoResponse {
|
||||
let Some(ref router_lock) = shared.model_router else {
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(serde_json::json!({"error": "no model router active"})),
|
||||
)
|
||||
.into_response();
|
||||
};
|
||||
|
||||
// Validate all name fields to prevent log injection and routing issues.
|
||||
for (field, value) in [
|
||||
("model_name", &body.model_name),
|
||||
("backend_name", &body.backend_name),
|
||||
("actual_model", &body.actual_model),
|
||||
] {
|
||||
if !is_safe_model_name(value) {
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(serde_json::json!({
|
||||
"error": format!("invalid {field}: contains disallowed characters")
|
||||
})),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
}
|
||||
|
||||
// ... rest of the existing function unchanged
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Run tests**
|
||||
|
||||
```bash
|
||||
cargo test -p anyllm_proxy 2>&1 | tail -20
|
||||
```
|
||||
|
||||
Expected: all tests pass.
|
||||
|
||||
- [ ] **Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/proxy/src/admin/routes.rs
|
||||
git commit -m "fix(admin): validate model_name/backend_name/actual_model in add_model
|
||||
|
||||
Applies is_safe_model_name to all three AddModelRequest fields before use,
|
||||
preventing log injection via newlines or control characters in audit log
|
||||
detail entries. Consistent with existing validation in put_config."
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Task 4: Fix MCP DNS rebinding via SSRF-safe HTTP client (Vuln 6 — Medium)
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/proxy/src/tools/mcp.rs` (replace reqwest::Client::new() with build_http_client)
|
||||
|
||||
- [ ] **Step 1: Replace the plain reqwest client in McpServerManager with an SSRF-safe client**
|
||||
|
||||
In `crates/proxy/src/tools/mcp.rs`, update the import and `McpServerManager::new()`:
|
||||
|
||||
At the top of the file, add the import:
|
||||
```rust
|
||||
use anyllm_client::http::{build_http_client, HttpClientConfig};
|
||||
```
|
||||
|
||||
Change `McpServerManager::new()`:
|
||||
```rust
|
||||
pub fn new() -> Self {
|
||||
let client = build_http_client(&HttpClientConfig {
|
||||
ssrf_protection: true,
|
||||
..Default::default()
|
||||
});
|
||||
Self {
|
||||
servers: RwLock::new(HashMap::new()),
|
||||
tool_to_server: RwLock::new(HashMap::new()),
|
||||
client,
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Also fix `discover_tools` (the static method that creates its own client):
|
||||
```rust
|
||||
pub async fn discover_tools(url: &str) -> Result<Vec<McpToolDef>, String> {
|
||||
let client = build_http_client(&HttpClientConfig {
|
||||
ssrf_protection: true,
|
||||
..Default::default()
|
||||
});
|
||||
discover_tools_impl(&client, url).await
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Run tests**
|
||||
|
||||
```bash
|
||||
cargo test -p anyllm_proxy 2>&1 | tail -20
|
||||
```
|
||||
|
||||
Expected: all tests pass. The SSRF-safe resolver is a no-op for the mock URLs used in unit tests.
|
||||
|
||||
- [ ] **Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/proxy/src/tools/mcp.rs
|
||||
git commit -m "fix(mcp): use SSRF-safe HTTP client for MCP tool calls
|
||||
|
||||
Replaces reqwest::Client::new() in McpServerManager with build_http_client
|
||||
(ssrf_protection: true), which attaches SsrfSafeDnsResolver. This prevents
|
||||
DNS rebinding: a domain that passes the registration-time check but later
|
||||
resolves to a private/metadata IP (e.g. 169.254.169.254) will be blocked
|
||||
at connection time by the resolver."
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Task 5: CSRF one-time token tracking (Vuln 4 — Medium)
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/proxy/src/admin/state.rs` (add issued_csrf_tokens DashMap)
|
||||
- Modify: `crates/proxy/src/admin/routes.rs` (store token on issue, remove on use)
|
||||
|
||||
- [ ] **Step 1: Check state.rs structure**
|
||||
|
||||
```bash
|
||||
grep -n "struct AdminState\|SharedState\|DashMap\|csrf" \
|
||||
crates/proxy/src/admin/state.rs | head -20
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Add `issued_csrf_tokens` to admin state**
|
||||
|
||||
In `crates/proxy/src/admin/state.rs`, find the `AdminState` or `SharedState` struct. Add a field:
|
||||
|
||||
```rust
|
||||
/// Set of CSRF tokens issued by GET /admin/csrf-token.
|
||||
/// Tokens are removed on first successful use (one-time tokens).
|
||||
/// Bounded by MaxAge=86400; background cleanup not strictly needed for
|
||||
/// localhost-only admin, but the set stays small in practice.
|
||||
pub issued_csrf_tokens: Arc<DashMap<String, ()>>,
|
||||
```
|
||||
|
||||
Also update the constructor / `Default` impl to initialize the new field:
|
||||
```rust
|
||||
issued_csrf_tokens: Arc::new(DashMap::new()),
|
||||
```
|
||||
|
||||
- [ ] **Step 3: Store token on issue in `get_csrf_token`**
|
||||
|
||||
In `crates/proxy/src/admin/routes.rs`, find `get_csrf_token()`. After generating the token, insert it into the shared set. This requires the handler to accept `State(shared): State<SharedState>`:
|
||||
|
||||
```rust
|
||||
async fn get_csrf_token(State(shared): State<SharedState>) -> axum::response::Response {
|
||||
let token = generate_csrf_token();
|
||||
shared.issued_csrf_tokens.insert(token.clone(), ());
|
||||
let body = serde_json::json!({"csrf_token": token});
|
||||
axum::http::Response::builder()
|
||||
// ... existing headers unchanged
|
||||
```
|
||||
|
||||
- [ ] **Step 4: Validate and consume the token in the CSRF middleware**
|
||||
|
||||
In `crates/proxy/src/admin/routes.rs`, find the CSRF validation middleware (the function that calls `validate_csrf_tokens`). After the `validate_csrf_tokens` check passes, also verify the token was server-issued and remove it (one-time use):
|
||||
|
||||
```rust
|
||||
// Verify the token was server-issued (prevents forgery).
|
||||
// Remove it immediately to enforce one-time use.
|
||||
if shared.issued_csrf_tokens.remove(csrf_header).is_none() {
|
||||
return (
|
||||
StatusCode::FORBIDDEN,
|
||||
axum::Json(serde_json::json!({
|
||||
"error": {"type": "forbidden", "message": "CSRF token not recognized or already used"}
|
||||
})),
|
||||
).into_response();
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 5: Run tests**
|
||||
|
||||
```bash
|
||||
cargo test -p anyllm_proxy 2>&1 | tail -20
|
||||
```
|
||||
|
||||
Expected: all tests pass. Any tests that exercise CSRF must now call the csrf-token endpoint first.
|
||||
|
||||
- [ ] **Step 6: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/proxy/src/admin/state.rs crates/proxy/src/admin/routes.rs
|
||||
git commit -m "fix(admin): enforce one-time CSRF tokens tracked server-side
|
||||
|
||||
GET /admin/csrf-token now inserts the generated token into a DashMap on
|
||||
SharedState. CSRF validation middleware checks the token was server-issued
|
||||
and removes it on first use, preventing replay of previously issued tokens.
|
||||
This closes the window where any localhost process could prefetch a token
|
||||
and reuse it across multiple requests."
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Task 6: BashTool startup warning (Vuln 1 — High, design-level)
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/proxy/src/tools/builtin/mod.rs` (add warning when bash is allow-listed)
|
||||
- Modify: `crates/proxy/src/config/simple.rs` (emit warn in build_tool_config)
|
||||
|
||||
- [ ] **Step 1: Add a startup warning when execute_bash policy is `allow`**
|
||||
|
||||
In `crates/proxy/src/config/simple.rs`, in the `build_tool_config` method, after the `action` is determined for a tool named `execute_bash` with `PolicyAction::Allow`, emit:
|
||||
|
||||
```rust
|
||||
if name == "execute_bash" && action == PolicyAction::Allow {
|
||||
tracing::warn!(
|
||||
"execute_bash policy is set to Allow. This permits the LLM to execute \
|
||||
arbitrary OS commands as the proxy process user. Ensure the proxy runs \
|
||||
in an isolated environment (seccomp, read-only rootfs, no network access \
|
||||
from the sandbox) before enabling this in production."
|
||||
);
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Run tests and build**
|
||||
|
||||
```bash
|
||||
cargo build 2>&1 | tail -10
|
||||
cargo test -p anyllm_proxy 2>&1 | tail -10
|
||||
```
|
||||
|
||||
Expected: clean build, tests pass.
|
||||
|
||||
- [ ] **Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/proxy/src/config/simple.rs
|
||||
git commit -m "fix(tools): warn at startup when execute_bash policy is Allow
|
||||
|
||||
Emits a tracing::warn! when execute_bash is configured with policy: allow
|
||||
so operators see an explicit reminder that the tool executes arbitrary OS
|
||||
commands. The dangerous-builtin-tools feature flag remains the primary
|
||||
compile-time gate; this is an additional runtime visibility measure."
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Final verification
|
||||
|
||||
- [ ] **Run full test suite**
|
||||
|
||||
```bash
|
||||
cargo test 2>&1 | tail -30
|
||||
```
|
||||
|
||||
Expected: ~906+ tests pass, 8 ignored, 0 failures.
|
||||
|
||||
- [ ] **Run clippy**
|
||||
|
||||
```bash
|
||||
cargo clippy -- -D warnings 2>&1 | tail -20
|
||||
```
|
||||
|
||||
Expected: no warnings.
|
||||
Reference in New Issue
Block a user