diff --git a/Cargo.lock b/Cargo.lock index ee35b8e..fd2f1b7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -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", ] diff --git a/Cargo.toml b/Cargo.toml index 2169970..7c17090 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -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] diff --git a/crates/batch_engine/Cargo.toml b/crates/batch_engine/Cargo.toml new file mode 100644 index 0000000..16f47db --- /dev/null +++ b/crates/batch_engine/Cargo.toml @@ -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" diff --git a/crates/batch_engine/src/db.rs b/crates/batch_engine/src/db.rs new file mode 100644 index 0000000..37b1a25 --- /dev/null +++ b/crates/batch_engine/src/db.rs @@ -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" + } +} diff --git a/crates/batch_engine/src/engine.rs b/crates/batch_engine/src/engine.rs new file mode 100644 index 0000000..418274e --- /dev/null +++ b/crates/batch_engine/src/engine.rs @@ -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 { + pub queue: Arc, + pub file_store: FileStore, + pub webhook_queue: Arc, + pub global_webhook_urls: Vec, + pub webhook_signing_secret: Option, +} + +impl BatchEngine { + /// Submit a new batch job. + pub async fn submit(&self, submission: BatchSubmission) -> Result { + // 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 = 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, EngineError> { + self.queue.get(id).await.map_err(EngineError::Queue) + } + + /// List batch jobs. + pub async fn list( + &self, + key_id: Option, + cursor: Option<&str>, + limit: u32, + ) -> Result, 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 { + 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, 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)> = 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 { + 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); + } +} diff --git a/crates/batch_engine/src/error.rs b/crates/batch_engine/src/error.rs new file mode 100644 index 0000000..f7ff547 --- /dev/null +++ b/crates/batch_engine/src/error.rs @@ -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 for QueueError { + fn from(e: rusqlite::Error) -> Self { + QueueError::Storage(e.to_string()) + } +} diff --git a/crates/batch_engine/src/file_store.rs b/crates/batch_engine/src/file_store.rs new file mode 100644 index 0000000..48819bf --- /dev/null +++ b/crates/batch_engine/src/file_store.rs @@ -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, + pub created_at: String, +} + +/// Manages batch file storage in SQLite. +#[derive(Clone)] +pub struct FileStore { + db: Arc>, +} + +impl FileStore { + pub fn new(db: Arc>) -> Self { + Self { db } + } + + /// Store a batch file. Returns the file_id. + pub async fn insert( + &self, + file_id: &str, + key_id: Option, + 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, 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>, 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 = 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()); + } +} diff --git a/crates/batch_engine/src/job.rs b/crates/batch_engine/src/job.rs new file mode 100644 index 0000000..0ab1c2c --- /dev/null +++ b/crates/batch_engine/src/job.rs @@ -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, + pub webhook_url: Option, + pub metadata: Option, + pub request_counts: RequestCounts, + pub input_file_id: String, + pub created_at: String, + pub started_at: Option, + pub completed_at: Option, + 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, + pub attempts: u8, + pub max_retries: u8, + pub last_error: Option, + pub next_retry_at: Option, + pub lease_id: Option, + pub lease_expires_at: Option, + pub idempotency_key: Option, + pub created_at: String, + pub completed_at: Option, +} + +#[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, + pub execution_mode: ExecutionMode, + pub input_file_id: String, + pub key_id: Option, + pub webhook_url: Option, + pub metadata: Option, + 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); + } +} diff --git a/crates/batch_engine/src/lib.rs b/crates/batch_engine/src/lib.rs new file mode 100644 index 0000000..70e4b1f --- /dev/null +++ b/crates/batch_engine/src/lib.rs @@ -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}; diff --git a/crates/batch_engine/src/queue/mod.rs b/crates/batch_engine/src/queue/mod.rs new file mode 100644 index 0000000..4995c77 --- /dev/null +++ b/crates/batch_engine/src/queue/mod.rs @@ -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, QueueError>; + async fn list( + &self, + key_id: Option, + cursor: Option<&str>, + limit: u32, + ) -> Result, QueueError>; + async fn cancel(&self, id: &BatchId) -> Result; + + // -- Item-level operations (proxy-native path, Phase 2) -- + async fn claim_next_item(&self) -> Result, 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; + async fn complete_batch(&self, id: &BatchId) -> Result<(), QueueError>; + + // -- Native batch support -- + async fn get_native_jobs_in_progress(&self) -> Result, QueueError>; + + // -- Lease management -- + async fn reclaim_expired_leases(&self) -> Result; + + // -- Progress -- + async fn update_progress(&self, id: &BatchId, counts: &RequestCounts) + -> Result<(), QueueError>; + + // -- Items query -- + async fn get_items(&self, batch_id: &BatchId) -> Result, QueueError>; +} diff --git a/crates/batch_engine/src/queue/sqlite.rs b/crates/batch_engine/src/queue/sqlite.rs new file mode 100644 index 0000000..bab86ce --- /dev/null +++ b/crates/batch_engine/src/queue/sqlite.rs @@ -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>, +} + +impl SqliteQueue { + pub fn new(db: Arc>) -> 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 = 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, 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, + cursor: Option<&str>, + limit: u32, + ) -> Result, 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> = 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::, _>>() + .map_err(QueueError::from) + }) + .await + .unwrap() + } + + async fn cancel(&self, id: &BatchId) -> Result { + 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 = 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, 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 { + 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, 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::, _>>() + .map_err(QueueError::from) + }) + .await + .unwrap() + } + + async fn reclaim_expired_leases(&self) -> Result { + 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, 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::, _>>() + .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 = row.get(3).unwrap_or(None); + let metadata_str: Option = 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 = row.get(7).unwrap_or(None); + let result_body_str: Option = row.get(8).unwrap_or(None); + + let source_format = + serde_json::from_str::(&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, 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) { + 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); + } +} diff --git a/crates/batch_engine/src/validation.rs b/crates/batch_engine/src/validation.rs new file mode 100644 index 0000000..18d796b --- /dev/null +++ b/crates/batch_engine/src/validation.rs @@ -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 { + 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 { + validate_jsonl(BufReader::new(Cursor::new(data.as_bytes()))) + } + + fn check_bytes(data: &[u8]) -> Result { + 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}"); + } +} diff --git a/crates/batch_engine/src/webhook/dispatcher.rs b/crates/batch_engine/src/webhook/dispatcher.rs new file mode 100644 index 0000000..00ff6c0 --- /dev/null +++ b/crates/batch_engine/src/webhook/dispatcher.rs @@ -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( + queue: Arc, + 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( + 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::::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"); + } + } + } + } +} diff --git a/crates/batch_engine/src/webhook/mod.rs b/crates/batch_engine/src/webhook/mod.rs new file mode 100644 index 0000000..aefaa55 --- /dev/null +++ b/crates/batch_engine/src/webhook/mod.rs @@ -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, + pub attempts: u8, + pub max_retries: u8, + pub next_retry_at: Option, +} + +/// 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, 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; +} diff --git a/crates/batch_engine/src/webhook/sqlite.rs b/crates/batch_engine/src/webhook/sqlite.rs new file mode 100644 index 0000000..19891ce --- /dev/null +++ b/crates/batch_engine/src/webhook/sqlite.rs @@ -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>, +} + +impl SqliteWebhookQueue { + pub fn new(db: Arc>) -> 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, 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 { + 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()); + } +} diff --git a/crates/proxy/Cargo.toml b/crates/proxy/Cargo.toml index aec3999..75e6d11 100644 --- a/crates/proxy/Cargo.toml +++ b/crates/proxy/Cargo.toml @@ -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"] } diff --git a/crates/proxy/src/admin/db.rs b/crates/proxy/src/admin/db.rs index d00ac49..6e55a0f 100644 --- a/crates/proxy/src/admin/db.rs +++ b/crates/proxy/src/admin/db.rs @@ -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 - ); ", )?; diff --git a/crates/proxy/src/admin/routes.rs b/crates/proxy/src/admin/routes.rs index 471329a..443d942 100644 --- a/crates/proxy/src/admin/routes.rs +++ b/crates/proxy/src/admin/routes.rs @@ -323,10 +323,7 @@ pub fn admin_router(shared: SharedState, token: Arc) -> 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) -> 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 }. diff --git a/crates/proxy/src/backend/mod.rs b/crates/proxy/src/backend/mod.rs index a2bb4b6..1143364 100644 --- a/crates/proxy/src/backend/mod.rs +++ b/crates/proxy/src/backend/mod.rs @@ -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") } } } diff --git a/crates/proxy/src/batch/anthropic_batch.rs b/crates/proxy/src/batch/anthropic_batch.rs index 890a3d0..1635c27 100644 --- a/crates/proxy/src/batch/anthropic_batch.rs +++ b/crates/proxy/src/batch/anthropic_batch.rs @@ -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; diff --git a/crates/proxy/src/batch/db.rs b/crates/proxy/src/batch/db.rs index 5ef9a25..672280f 100644 --- a/crates/proxy/src/batch/db.rs +++ b/crates/proxy/src/batch/db.rs @@ -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, - 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> { - 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, - 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> { - 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, - failed_count: Option, - output_file_id: Option<&str>, - error_file_id: Option<&str>, -) -> rusqlite::Result { - 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, - limit: u32, - after: Option<&str>, -) -> rusqlite::Result> { - 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> = 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 { - let status_str: String = row.get(3)?; - let created_at_str: String = row.get(9)?; - let metadata_str: Option = 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>(10)? - .map(|s| iso8601_to_epoch(&s)), - expires_at: row - .get::<_, Option>(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 = parts[0].split('-').filter_map(|p| p.parse().ok()).collect(); - let time_str = parts[1].trim_end_matches('Z'); - let time_parts: Vec = 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(); diff --git a/crates/proxy/src/batch/mod.rs b/crates/proxy/src/batch/mod.rs index 051bd4c..8f77f62 100644 --- a/crates/proxy/src/batch/mod.rs +++ b/crates/proxy/src/batch/mod.rs @@ -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, - 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(skip_serializing_if = "Option::is_none")] - pub output_file_id: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub error_file_id: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub completed_at: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub expires_at: Option, -} - -/// 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 { - 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 { - validate_jsonl(BufReader::new(Cursor::new(data.as_bytes()))) - } - - fn check_bytes(data: &[u8]) -> Result { - 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; diff --git a/crates/proxy/src/batch/routes.rs b/crates/proxy/src/batch/routes.rs index a25b5b3..1cc5d36 100644 --- a/crates/proxy/src/batch/routes.rs +++ b/crates/proxy/src/batch/routes.rs @@ -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, 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, 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, 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, Json(req): Json, @@ -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 = 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, Path(batch_id): Path) -> 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, Path(batch_id): Path { + 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, Query(query): Query, ) -> 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 = 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, Path(batch_id): Path) -> 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, 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 = parts[0].split('-').filter_map(|p| p.parse().ok()).collect(); + let time_str = parts[1].trim_end_matches('Z'); + let time_parts: Vec = 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() +} diff --git a/crates/proxy/src/config/simple.rs b/crates/proxy/src/config/simple.rs index d2f112f..560f60c 100644 --- a/crates/proxy/src/config/simple.rs +++ b/crates/proxy/src/config/simple.rs @@ -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) diff --git a/crates/proxy/src/main.rs b/crates/proxy/src/main.rs index ecb05dc..455bbc4 100644 --- a/crates/proxy/src/main.rs +++ b/crates/proxy/src/main.rs @@ -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> = - 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> = 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 = 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 --- diff --git a/crates/proxy/src/server/chat_completions.rs b/crates/proxy/src/server/chat_completions.rs index 7ed9af1..24a3df6 100644 --- a/crates/proxy/src/server/chat_completions.rs +++ b/crates/proxy/src/server/chat_completions.rs @@ -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; diff --git a/crates/proxy/src/server/routes.rs b/crates/proxy/src/server/routes.rs index fe7016c..5062d5e 100644 --- a/crates/proxy/src/server/routes.rs +++ b/crates/proxy/src/server/routes.rs @@ -107,6 +107,15 @@ pub struct AppState { pub all_backends: Option>>, /// Tool execution engine state. None when tool execution is not configured. pub tool_engine: Option>, + /// 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, model_router: Option>>, tool_engine: Option>, + 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 = 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 { .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 { diff --git a/crates/proxy/src/tools/builtin/bash.rs b/crates/proxy/src/tools/builtin/bash.rs index 4e5247c..2540234 100644 --- a/crates/proxy/src/tools/builtin/bash.rs +++ b/crates/proxy/src/tools/builtin/bash.rs @@ -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); diff --git a/crates/proxy/src/tools/execution.rs b/crates/proxy/src/tools/execution.rs index 1dd976f..c3a757b 100644 --- a/crates/proxy/src/tools/execution.rs +++ b/crates/proxy/src/tools/execution.rs @@ -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), diff --git a/crates/proxy/src/tools/mcp.rs b/crates/proxy/src/tools/mcp.rs index 9502e75..7623b95 100644 --- a/crates/proxy/src/tools/mcp.rs +++ b/crates/proxy/src/tools/mcp.rs @@ -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 { - 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 { @@ -201,49 +205,51 @@ impl McpServerManager { } } -async fn discover_tools_impl(client: &reqwest::Client, url: &str) -> Result, 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, 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, registry: &mut crate::tools::ToolRegistry) { +pub fn register_mcp_tools( + manager: &Arc, + 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] diff --git a/crates/proxy/tests/batch_api.rs b/crates/proxy/tests/batch_api.rs index a5bb2f3..ff422cf 100644 --- a/crates/proxy/tests/batch_api.rs +++ b/crates/proxy/tests/batch_api.rs @@ -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_") + 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; diff --git a/crates/proxy/tests/tool_execution.rs b/crates/proxy/tests/tool_execution.rs index c0907fa..f36b6b0 100644 --- a/crates/proxy/tests/tool_execution.rs +++ b/crates/proxy/tests/tool_execution.rs @@ -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] diff --git a/crates/proxy/tests/virtual_keys.rs b/crates/proxy/tests/virtual_keys.rs index 7cf227e..41852d9 100644 --- a/crates/proxy/tests/virtual_keys.rs +++ b/crates/proxy/tests/virtual_keys.rs @@ -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() diff --git a/docs/superpowers/plans/2026-03-31-security-audit-fixes.md b/docs/superpowers/plans/2026-03-31-security-audit-fixes.md new file mode 100644 index 0000000..2d81ebc --- /dev/null +++ b/docs/superpowers/plans/2026-03-31-security-audit-fixes.md @@ -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, + #[serde(default)] + pub timeout_secs: Option, + /// 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, +} +``` + +- [ ] **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, +} + +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> + 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>, +) { + #[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, +) -> 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, + State(shared): State, + Json(body): Json, +) -> 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, 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>, +``` + +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`: + +```rust +async fn get_csrf_token(State(shared): State) -> 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.