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/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()