starred/anyllm-proxy
1
0
Fork 0
mirror of https://github.com/whit3rabbit/anyllm-proxy.git synced 2026-09-22 00:00:50 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge feat/batch-engine-phase1 into main

Browse Source 
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
parent e6dd396b05 ef5f0f596c
commit 8ac9ef2677
34 changed files with 3825 additions and 1051 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/proxy/Cargo.toml
+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"] }
crates/proxy/src/admin/db.rs
-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
);
",
)?;
crates/proxy/src/admin/routes.rs
+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 }.
crates/proxy/src/backend/mod.rs
+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")
}
}
}
crates/proxy/src/batch/anthropic_batch.rs
+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;
crates/proxy/src/batch/db.rs
+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();
crates/proxy/src/batch/mod.rs
+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;
crates/proxy/src/batch/routes.rs
+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()
}
crates/proxy/src/config/simple.rs
+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)
crates/proxy/src/main.rs
+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 ---
crates/proxy/src/server/chat_completions.rs
+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;
crates/proxy/src/server/routes.rs
+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 {
crates/proxy/src/tools/builtin/bash.rs
+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);
crates/proxy/src/tools/execution.rs
+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),
crates/proxy/src/tools/mcp.rs
+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]
crates/proxy/tests/batch_api.rs
+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;
crates/proxy/tests/tool_execution.rs
+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]
crates/proxy/tests/virtual_keys.rs
+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()