Merge feat/batch-engine-phase1 into main

Resolves conflicts: take HEAD (security audit) for mcp.rs imports,
register_server_blocking error handling, and maybe_execute_tools loop.

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