mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
refactor: 拆分 gateway 单体为独立 crate,新增 systemd 部署方案
将 gateway 内部的 model-fetch、provider-transport、scheduler-core、 usage-runtime、video-tasks-core 模块提取为独立 crate;重构 gateway 内部模块结构(state/router/cache/data/query 等);移除大量遗留模块 文件;新增 systemd 二进制部署骨架及相关文档;更新前端 usage 相关 API 和组件。
This commit is contained in:
538
apps/aether-gateway/src/tests/ai_execute/control_execute.rs
Normal file
538
apps/aether-gateway/src/tests/ai_execute/control_execute.rs
Normal file
@@ -0,0 +1,538 @@
|
||||
use super::{
|
||||
any, build_router, build_router_with_execution_runtime_override, json, start_server, Arc, Body,
|
||||
Bytes, HeaderName, HeaderValue, Infallible, Json, Mutex, Request, Response, Router, StatusCode,
|
||||
CONTROL_EXECUTED_HEADER, CONTROL_EXECUTE_FALLBACK_HEADER, DEPENDENCY_REASON_HEADER,
|
||||
EXECUTION_PATH_HEADER, EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS, TRACE_ID_HEADER,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_locally_denies_sync_ai_control_execute_when_opted_in_and_execution_runtime_missing(
|
||||
) {
|
||||
let execute_hits = Arc::new(Mutex::new(0usize));
|
||||
let execute_hits_clone = Arc::clone(&execute_hits);
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
let public_execution_path = Arc::new(Mutex::new(None::<String>));
|
||||
let public_execution_path_clone = Arc::clone(&public_execution_path);
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/resolve",
|
||||
any(|_request: Request| async move {
|
||||
Json(json!({
|
||||
"action": "proxy_public",
|
||||
"route_class": "ai_public",
|
||||
"route_family": "openai",
|
||||
"route_kind": "chat",
|
||||
"auth_endpoint_signature": "openai:chat",
|
||||
"execution_runtime_candidate": true,
|
||||
"auth_context": {
|
||||
"user_id": "user-sync-123",
|
||||
"api_key_id": "key-sync-123",
|
||||
"access_allowed": true
|
||||
},
|
||||
"public_path": "/v1/chat/completions"
|
||||
}))
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/execute-sync",
|
||||
any(move |_request: Request| {
|
||||
let execute_hits_inner = Arc::clone(&execute_hits_clone);
|
||||
async move {
|
||||
*execute_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
let mut response = Response::builder()
|
||||
.status(StatusCode::CREATED)
|
||||
.body(Body::from("{\"ok\":true}"))
|
||||
.expect("response should build");
|
||||
response.headers_mut().insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
HeaderValue::from_static("application/json"),
|
||||
);
|
||||
response.headers_mut().insert(
|
||||
HeaderName::from_static(CONTROL_EXECUTED_HEADER),
|
||||
HeaderValue::from_static("true"),
|
||||
);
|
||||
response
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1/chat/completions",
|
||||
any(move |request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
let public_execution_path_inner = Arc::clone(&public_execution_path_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
*public_execution_path_inner
|
||||
.lock()
|
||||
.expect("mutex should lock") = Some(
|
||||
request
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
);
|
||||
let mut response = Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from("{\"public\":true}"))
|
||||
.expect("response should build");
|
||||
response.headers_mut().insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
HeaderValue::from_static("application/json"),
|
||||
);
|
||||
response
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router().expect("gateway should build");
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/chat/completions"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(CONTROL_EXECUTE_FALLBACK_HEADER, "true")
|
||||
.header(TRACE_ID_HEADER, "trace-sync-123")
|
||||
.body("{\"model\":\"gpt-5\",\"messages\":[]}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
let execution_path = response
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(ToOwned::to_owned);
|
||||
let dependency_reason = response
|
||||
.headers()
|
||||
.get(DEPENDENCY_REASON_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(ToOwned::to_owned);
|
||||
assert_eq!(dependency_reason.as_deref(), None);
|
||||
let payload: serde_json::Value = response.json().await.expect("body should parse");
|
||||
assert_eq!(
|
||||
execution_path.as_deref(),
|
||||
Some(EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS)
|
||||
);
|
||||
assert_eq!(
|
||||
payload["error"]["message"],
|
||||
"OpenAI chat execution runtime miss did not match a Rust execution path"
|
||||
);
|
||||
assert_eq!(*execute_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(
|
||||
public_execution_path
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.as_deref(),
|
||||
None
|
||||
);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_locally_denies_stream_ai_control_execute_when_opted_in_and_execution_runtime_missing(
|
||||
) {
|
||||
let execute_hits = Arc::new(Mutex::new(0usize));
|
||||
let execute_hits_clone = Arc::clone(&execute_hits);
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
let public_execution_path = Arc::new(Mutex::new(None::<String>));
|
||||
let public_execution_path_clone = Arc::clone(&public_execution_path);
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/resolve",
|
||||
any(|_request: Request| async move {
|
||||
Json(json!({
|
||||
"action": "proxy_public",
|
||||
"route_class": "ai_public",
|
||||
"route_family": "openai",
|
||||
"route_kind": "chat",
|
||||
"auth_endpoint_signature": "openai:chat",
|
||||
"execution_runtime_candidate": true,
|
||||
"auth_context": {
|
||||
"user_id": "user-stream-123",
|
||||
"api_key_id": "key-stream-123",
|
||||
"access_allowed": true
|
||||
},
|
||||
"public_path": "/v1/chat/completions"
|
||||
}))
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/execute-stream",
|
||||
any(move |_request: Request| {
|
||||
let execute_hits_inner = Arc::clone(&execute_hits_clone);
|
||||
async move {
|
||||
*execute_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
let stream = futures_util::stream::iter([
|
||||
Ok::<_, Infallible>(Bytes::from_static(b"data: one\n\n")),
|
||||
Ok::<_, Infallible>(Bytes::from_static(b"data: [DONE]\n\n")),
|
||||
]);
|
||||
let mut response = Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from_stream(stream))
|
||||
.expect("response should build");
|
||||
response.headers_mut().insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
HeaderValue::from_static("text/event-stream"),
|
||||
);
|
||||
response.headers_mut().insert(
|
||||
HeaderName::from_static(CONTROL_EXECUTED_HEADER),
|
||||
HeaderValue::from_static("true"),
|
||||
);
|
||||
response
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1/chat/completions",
|
||||
any(move |request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
let public_execution_path_inner = Arc::clone(&public_execution_path_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
*public_execution_path_inner
|
||||
.lock()
|
||||
.expect("mutex should lock") = Some(
|
||||
request
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
);
|
||||
let mut response = Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from("data: {\"public\":true}\n\ndata: [DONE]\n\n"))
|
||||
.expect("response should build");
|
||||
response.headers_mut().insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
HeaderValue::from_static("text/event-stream"),
|
||||
);
|
||||
response
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router().expect("gateway should build");
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/chat/completions"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(CONTROL_EXECUTE_FALLBACK_HEADER, "true")
|
||||
.header(TRACE_ID_HEADER, "trace-stream-123")
|
||||
.body("{\"model\":\"gpt-5\",\"messages\":[],\"stream\":true}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
let execution_path = response
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(ToOwned::to_owned);
|
||||
let dependency_reason = response
|
||||
.headers()
|
||||
.get(DEPENDENCY_REASON_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(ToOwned::to_owned);
|
||||
assert_eq!(dependency_reason.as_deref(), None);
|
||||
let payload: serde_json::Value = response.json().await.expect("body should parse");
|
||||
assert_eq!(
|
||||
execution_path.as_deref(),
|
||||
Some(EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS)
|
||||
);
|
||||
assert_eq!(
|
||||
payload["error"]["message"],
|
||||
"OpenAI chat execution runtime miss did not match a Rust execution path"
|
||||
);
|
||||
assert_eq!(*execute_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(
|
||||
public_execution_path
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.as_deref(),
|
||||
None
|
||||
);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_does_not_proxy_control_execute_over_http_when_opted_in_and_execution_runtime_misses_sync_ai_routes(
|
||||
) {
|
||||
let plan_hits = Arc::new(Mutex::new(0usize));
|
||||
let plan_hits_clone = Arc::clone(&plan_hits);
|
||||
let execute_hits = Arc::new(Mutex::new(0usize));
|
||||
let execute_hits_clone = Arc::clone(&execute_hits);
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/resolve",
|
||||
any(|_request: Request| async move {
|
||||
Json(json!({
|
||||
"action": "proxy_public",
|
||||
"route_class": "ai_public",
|
||||
"route_family": "openai",
|
||||
"route_kind": "chat",
|
||||
"auth_endpoint_signature": "openai:chat",
|
||||
"execution_runtime_candidate": true,
|
||||
"public_path": "/v1/chat/completions"
|
||||
}))
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/plan-sync",
|
||||
any(move |_request: Request| {
|
||||
let plan_hits_inner = Arc::clone(&plan_hits_clone);
|
||||
async move {
|
||||
*plan_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
let mut response = Response::builder()
|
||||
.status(StatusCode::ACCEPTED)
|
||||
.body(Body::from("{\"ok\":true,\"via\":\"plan-sync\"}"))
|
||||
.expect("response should build");
|
||||
response.headers_mut().insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
HeaderValue::from_static("application/json"),
|
||||
);
|
||||
response.headers_mut().insert(
|
||||
HeaderName::from_static(CONTROL_EXECUTED_HEADER),
|
||||
HeaderValue::from_static("true"),
|
||||
);
|
||||
response
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/execute-sync",
|
||||
any(move |_request: Request| {
|
||||
let execute_hits_inner = Arc::clone(&execute_hits_clone);
|
||||
async move {
|
||||
*execute_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
let mut response = Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from("{\"ok\":false}"))
|
||||
.expect("response should build");
|
||||
response.headers_mut().insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
HeaderValue::from_static("application/json"),
|
||||
);
|
||||
response.headers_mut().insert(
|
||||
HeaderName::from_static(CONTROL_EXECUTED_HEADER),
|
||||
HeaderValue::from_static("true"),
|
||||
);
|
||||
response
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1/chat/completions",
|
||||
any(move |_request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
let mut response = Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from("{\"public\":true}"))
|
||||
.expect("response should build");
|
||||
response.headers_mut().insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
HeaderValue::from_static("application/json"),
|
||||
);
|
||||
response
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let execution_runtime = Router::new();
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let gateway = build_router_with_execution_runtime_override(execution_runtime_url);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/chat/completions"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(CONTROL_EXECUTE_FALLBACK_HEADER, "true")
|
||||
.body("{\"model\":\"gpt-5\",\"messages\":[]}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some(EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS)
|
||||
);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(DEPENDENCY_REASON_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
None
|
||||
);
|
||||
let payload: serde_json::Value = response.json().await.expect("body should parse");
|
||||
assert_eq!(
|
||||
payload["error"]["message"],
|
||||
"OpenAI chat execution runtime miss did not match a Rust execution path"
|
||||
);
|
||||
assert_eq!(*plan_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*execute_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_does_not_proxy_control_execute_over_http_when_opted_in_and_execution_runtime_misses_stream_ai_routes(
|
||||
) {
|
||||
let plan_hits = Arc::new(Mutex::new(0usize));
|
||||
let plan_hits_clone = Arc::clone(&plan_hits);
|
||||
let execute_hits = Arc::new(Mutex::new(0usize));
|
||||
let execute_hits_clone = Arc::clone(&execute_hits);
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/resolve",
|
||||
any(|_request: Request| async move {
|
||||
Json(json!({
|
||||
"action": "proxy_public",
|
||||
"route_class": "ai_public",
|
||||
"route_family": "openai",
|
||||
"route_kind": "chat",
|
||||
"auth_endpoint_signature": "openai:chat",
|
||||
"execution_runtime_candidate": true,
|
||||
"public_path": "/v1/chat/completions"
|
||||
}))
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/plan-stream",
|
||||
any(move |_request: Request| {
|
||||
let plan_hits_inner = Arc::clone(&plan_hits_clone);
|
||||
async move {
|
||||
*plan_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
let stream = futures_util::stream::iter([
|
||||
Ok::<_, Infallible>(Bytes::from_static(b"data: one\n\n")),
|
||||
Ok::<_, Infallible>(Bytes::from_static(b"data: [DONE]\n\n")),
|
||||
]);
|
||||
let mut response = Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from_stream(stream))
|
||||
.expect("response should build");
|
||||
response.headers_mut().insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
HeaderValue::from_static("text/event-stream"),
|
||||
);
|
||||
response.headers_mut().insert(
|
||||
HeaderName::from_static(CONTROL_EXECUTED_HEADER),
|
||||
HeaderValue::from_static("true"),
|
||||
);
|
||||
response
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/execute-stream",
|
||||
any(move |_request: Request| {
|
||||
let execute_hits_inner = Arc::clone(&execute_hits_clone);
|
||||
async move {
|
||||
*execute_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
let mut response = Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from("data: fallback\n\n"))
|
||||
.expect("response should build");
|
||||
response.headers_mut().insert(
|
||||
HeaderName::from_static(CONTROL_EXECUTED_HEADER),
|
||||
HeaderValue::from_static("true"),
|
||||
);
|
||||
response
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1/chat/completions",
|
||||
any(move |_request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
let mut response = Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from("data: {\"public\":true}\n\ndata: [DONE]\n\n"))
|
||||
.expect("response should build");
|
||||
response.headers_mut().insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
HeaderValue::from_static("text/event-stream"),
|
||||
);
|
||||
response
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let execution_runtime = Router::new();
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let gateway = build_router_with_execution_runtime_override(execution_runtime_url);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/chat/completions"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(CONTROL_EXECUTE_FALLBACK_HEADER, "true")
|
||||
.body("{\"model\":\"gpt-5\",\"messages\":[],\"stream\":true}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some(EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS)
|
||||
);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(DEPENDENCY_REASON_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
None
|
||||
);
|
||||
let payload: serde_json::Value = response.json().await.expect("body should parse");
|
||||
assert_eq!(
|
||||
payload["error"]["message"],
|
||||
"OpenAI chat execution runtime miss did not match a Rust execution path"
|
||||
);
|
||||
assert_eq!(*plan_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*execute_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
731
apps/aether-gateway/src/tests/ai_execute/fallback.rs
Normal file
731
apps/aether-gateway/src/tests/ai_execute/fallback.rs
Normal file
@@ -0,0 +1,731 @@
|
||||
use super::{
|
||||
any, build_router, build_router_with_execution_runtime_override, json, start_server, Arc, Body,
|
||||
HeaderValue, Json, Mutex, Request, Response, Router, StatusCode, DEPENDENCY_REASON_HEADER,
|
||||
EXECUTION_PATH_HEADER, EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS,
|
||||
LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER, TRACE_ID_HEADER,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_locally_denies_openai_chat_after_repeated_execution_runtime_misses_without_control_execute_opt_in(
|
||||
) {
|
||||
let execute_hits = Arc::new(Mutex::new(0usize));
|
||||
let execute_hits_clone = Arc::clone(&execute_hits);
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
let public_execution_paths = Arc::new(Mutex::new(Vec::<String>::new()));
|
||||
let public_execution_paths_clone = Arc::clone(&public_execution_paths);
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/resolve",
|
||||
any(|_request: Request| async move {
|
||||
Json(json!({
|
||||
"action": "proxy_public",
|
||||
"route_class": "ai_public",
|
||||
"route_family": "openai",
|
||||
"route_kind": "chat",
|
||||
"auth_endpoint_signature": "openai:chat",
|
||||
"execution_runtime_candidate": true,
|
||||
"public_path": "/v1/chat/completions"
|
||||
}))
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/execute-sync",
|
||||
any(move |_request: Request| {
|
||||
let execute_hits_inner = Arc::clone(&execute_hits_clone);
|
||||
async move {
|
||||
*execute_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
let mut response = Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from("{\"fallback\":true}"))
|
||||
.expect("response should build");
|
||||
response.headers_mut().insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
HeaderValue::from_static("application/json"),
|
||||
);
|
||||
response
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1/chat/completions",
|
||||
any(move |request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
let public_execution_paths_inner = Arc::clone(&public_execution_paths_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
public_execution_paths_inner
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.push(
|
||||
request
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
);
|
||||
let mut response = Response::builder()
|
||||
.status(StatusCode::ACCEPTED)
|
||||
.body(Body::from("{\"proxied\":true}"))
|
||||
.expect("response should build");
|
||||
response.headers_mut().insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
HeaderValue::from_static("application/json"),
|
||||
);
|
||||
response
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let execution_runtime = Router::new();
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let gateway = build_router_with_execution_runtime_override(execution_runtime_url);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
for trace_id in ["trace-openai-chat-bypass-1", "trace-openai-chat-bypass-2"] {
|
||||
let response = client
|
||||
.post(format!("{gateway_url}/v1/chat/completions"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(TRACE_ID_HEADER, trace_id)
|
||||
.body("{\"model\":\"gpt-5\",\"messages\":[]}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some(EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS)
|
||||
);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(DEPENDENCY_REASON_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("missing_auth_context")
|
||||
);
|
||||
let payload: serde_json::Value = response.json().await.expect("body should parse");
|
||||
assert_eq!(payload["error"]["type"], "http_error");
|
||||
assert_eq!(
|
||||
payload["error"]["message"],
|
||||
"OpenAI chat execution runtime miss did not match a Rust execution path"
|
||||
);
|
||||
}
|
||||
|
||||
assert_eq!(*execute_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(
|
||||
*public_execution_paths.lock().expect("mutex should lock"),
|
||||
Vec::<String>::new()
|
||||
);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_locally_denies_openai_chat_when_control_api_is_configured_without_opt_in() {
|
||||
let execute_hits = Arc::new(Mutex::new(0usize));
|
||||
let execute_hits_clone = Arc::clone(&execute_hits);
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
let public_execution_path = Arc::new(Mutex::new(None::<String>));
|
||||
let public_execution_path_clone = Arc::clone(&public_execution_path);
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/resolve",
|
||||
any(|_request: Request| async move {
|
||||
Json(json!({
|
||||
"action": "proxy_public",
|
||||
"route_class": "ai_public",
|
||||
"route_family": "openai",
|
||||
"route_kind": "chat",
|
||||
"auth_endpoint_signature": "openai:chat",
|
||||
"execution_runtime_candidate": true,
|
||||
"public_path": "/v1/chat/completions"
|
||||
}))
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/execute-sync",
|
||||
any(move |_request: Request| {
|
||||
let execute_hits_inner = Arc::clone(&execute_hits_clone);
|
||||
async move {
|
||||
*execute_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
let mut response = Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from("{\"fallback\":true}"))
|
||||
.expect("response should build");
|
||||
response.headers_mut().insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
HeaderValue::from_static("application/json"),
|
||||
);
|
||||
response
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1/chat/completions",
|
||||
any(move |request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
let public_execution_path_inner = Arc::clone(&public_execution_path_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
*public_execution_path_inner
|
||||
.lock()
|
||||
.expect("mutex should lock") = Some(
|
||||
request
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
);
|
||||
let mut response = Response::builder()
|
||||
.status(StatusCode::ACCEPTED)
|
||||
.body(Body::from("{\"proxied\":true}"))
|
||||
.expect("response should build");
|
||||
response.headers_mut().insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
HeaderValue::from_static("application/json"),
|
||||
);
|
||||
response
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router().expect("gateway should build");
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/chat/completions"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.body("{\"model\":\"gpt-5\",\"messages\":[]}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some(EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS)
|
||||
);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(DEPENDENCY_REASON_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("missing_auth_context")
|
||||
);
|
||||
let payload: serde_json::Value = response.json().await.expect("body should parse");
|
||||
assert_eq!(payload["error"]["type"], "http_error");
|
||||
assert_eq!(
|
||||
payload["error"]["message"],
|
||||
"OpenAI chat execution runtime miss did not match a Rust execution path"
|
||||
);
|
||||
assert_eq!(*execute_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(
|
||||
public_execution_path
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.as_deref(),
|
||||
None
|
||||
);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_locally_denies_openai_chat_stream_after_execution_runtime_miss_without_control_execute_opt_in(
|
||||
) {
|
||||
let execute_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
let public_execution_path = Arc::new(Mutex::new(None::<String>));
|
||||
let public_execution_path_clone = Arc::clone(&public_execution_path);
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/resolve",
|
||||
any(|_request: Request| async move {
|
||||
Json(json!({
|
||||
"action": "proxy_public",
|
||||
"route_class": "ai_public",
|
||||
"route_family": "openai",
|
||||
"route_kind": "chat",
|
||||
"auth_endpoint_signature": "openai:chat",
|
||||
"execution_runtime_candidate": true,
|
||||
"public_path": "/v1/chat/completions"
|
||||
}))
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1/chat/completions",
|
||||
any(move |request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
let public_execution_path_inner = Arc::clone(&public_execution_path_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
*public_execution_path_inner
|
||||
.lock()
|
||||
.expect("mutex should lock") = Some(
|
||||
request
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
);
|
||||
let mut response = Response::builder()
|
||||
.status(StatusCode::ACCEPTED)
|
||||
.body(Body::from("{\"proxied\":true}"))
|
||||
.expect("response should build");
|
||||
response.headers_mut().insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
HeaderValue::from_static("application/json"),
|
||||
);
|
||||
response
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let execution_runtime = Router::new();
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let gateway = build_router_with_execution_runtime_override(execution_runtime_url);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/chat/completions"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.body("{\"model\":\"gpt-5\",\"messages\":[],\"stream\":true}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some(EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS)
|
||||
);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(DEPENDENCY_REASON_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("missing_auth_context")
|
||||
);
|
||||
let payload: serde_json::Value = response.json().await.expect("body should parse");
|
||||
assert_eq!(payload["error"]["type"], "http_error");
|
||||
assert_eq!(
|
||||
payload["error"]["message"],
|
||||
"OpenAI chat execution runtime miss did not match a Rust execution path"
|
||||
);
|
||||
assert_eq!(*execute_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(
|
||||
public_execution_path
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.as_deref(),
|
||||
None
|
||||
);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
async fn assert_ai_route_locally_denied_after_execution_runtime_miss(
|
||||
public_path: &'static str,
|
||||
route_family: &'static str,
|
||||
route_kind: &'static str,
|
||||
endpoint_signature: &'static str,
|
||||
request_body: &'static str,
|
||||
expected_message: &'static str,
|
||||
) {
|
||||
assert_ai_route_locally_denied_after_execution_runtime_miss_with_request(
|
||||
reqwest::Method::POST,
|
||||
public_path,
|
||||
public_path,
|
||||
route_family,
|
||||
route_kind,
|
||||
endpoint_signature,
|
||||
Some(request_body),
|
||||
expected_message,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
async fn assert_ai_route_locally_denied_after_execution_runtime_miss_with_request(
|
||||
method: reqwest::Method,
|
||||
route_path: &'static str,
|
||||
request_path: &'static str,
|
||||
route_family: &'static str,
|
||||
route_kind: &'static str,
|
||||
endpoint_signature: &'static str,
|
||||
request_body: Option<&'static str>,
|
||||
expected_message: &'static str,
|
||||
) {
|
||||
let control_execute_hits = Arc::new(Mutex::new(0usize));
|
||||
let control_execute_hits_clone = Arc::clone(&control_execute_hits);
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
let public_execution_path = Arc::new(Mutex::new(None::<String>));
|
||||
let public_execution_path_clone = Arc::clone(&public_execution_path);
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/resolve",
|
||||
any(move |_request: Request| async move {
|
||||
Json(json!({
|
||||
"action": "proxy_public",
|
||||
"route_class": "ai_public",
|
||||
"route_family": route_family,
|
||||
"route_kind": route_kind,
|
||||
"auth_endpoint_signature": endpoint_signature,
|
||||
"execution_runtime_candidate": true,
|
||||
"public_path": route_path
|
||||
}))
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/execute-sync",
|
||||
any(move |_request: Request| {
|
||||
let control_execute_hits_inner = Arc::clone(&control_execute_hits_clone);
|
||||
async move {
|
||||
*control_execute_hits_inner
|
||||
.lock()
|
||||
.expect("mutex should lock") += 1;
|
||||
let mut response = Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from("{\"fallback\":true}"))
|
||||
.expect("response should build");
|
||||
response.headers_mut().insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
HeaderValue::from_static("application/json"),
|
||||
);
|
||||
response
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
route_path,
|
||||
any(move |request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
let public_execution_path_inner = Arc::clone(&public_execution_path_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
*public_execution_path_inner
|
||||
.lock()
|
||||
.expect("mutex should lock") = Some(
|
||||
request
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
);
|
||||
let mut response = Response::builder()
|
||||
.status(StatusCode::ACCEPTED)
|
||||
.body(Body::from("{\"proxied\":true}"))
|
||||
.expect("response should build");
|
||||
response.headers_mut().insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
HeaderValue::from_static("application/json"),
|
||||
);
|
||||
response
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let execution_runtime = Router::new();
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let gateway = build_router_with_execution_runtime_override(execution_runtime_url);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let mut request =
|
||||
reqwest::Client::new().request(method, format!("{gateway_url}{request_path}"));
|
||||
if let Some(request_body) = request_body {
|
||||
request = request
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.body(request_body);
|
||||
}
|
||||
let response = request.send().await.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some(EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS)
|
||||
);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(DEPENDENCY_REASON_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
None
|
||||
);
|
||||
let payload: serde_json::Value = response.json().await.expect("body should parse");
|
||||
assert_eq!(payload["error"]["type"], "http_error");
|
||||
assert_eq!(payload["error"]["message"], expected_message);
|
||||
assert_eq!(*control_execute_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(
|
||||
public_execution_path
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.as_deref(),
|
||||
None
|
||||
);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_locally_denies_openai_responses_after_execution_runtime_miss_without_control_execute_opt_in(
|
||||
) {
|
||||
assert_ai_route_locally_denied_after_execution_runtime_miss(
|
||||
"/v1/responses",
|
||||
"openai",
|
||||
"cli",
|
||||
"openai:cli",
|
||||
"{\"model\":\"gpt-5\",\"input\":\"hello\"}",
|
||||
"OpenAI responses execution runtime miss did not match a Rust execution path",
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_locally_denies_claude_messages_after_execution_runtime_miss_without_control_execute_opt_in(
|
||||
) {
|
||||
assert_ai_route_locally_denied_after_execution_runtime_miss(
|
||||
"/v1/messages",
|
||||
"claude",
|
||||
"chat",
|
||||
"claude:chat",
|
||||
"{\"model\":\"claude-sonnet-4-5\",\"messages\":[]}",
|
||||
"Claude messages execution runtime miss did not match a Rust execution path",
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_locally_denies_openai_responses_stream_after_execution_runtime_miss_without_control_execute_opt_in(
|
||||
) {
|
||||
assert_ai_route_locally_denied_after_execution_runtime_miss(
|
||||
"/v1/responses",
|
||||
"openai",
|
||||
"cli",
|
||||
"openai:cli",
|
||||
"{\"model\":\"gpt-5\",\"input\":\"hello\",\"stream\":true}",
|
||||
"OpenAI responses execution runtime miss did not match a Rust execution path",
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_locally_denies_claude_messages_stream_after_execution_runtime_miss_without_control_execute_opt_in(
|
||||
) {
|
||||
assert_ai_route_locally_denied_after_execution_runtime_miss(
|
||||
"/v1/messages",
|
||||
"claude",
|
||||
"chat",
|
||||
"claude:chat",
|
||||
"{\"model\":\"claude-sonnet-4-5\",\"messages\":[],\"stream\":true}",
|
||||
"Claude messages execution runtime miss did not match a Rust execution path",
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_locally_denies_openai_compact_after_execution_runtime_miss_without_control_execute_opt_in(
|
||||
) {
|
||||
assert_ai_route_locally_denied_after_execution_runtime_miss(
|
||||
"/v1/responses/compact",
|
||||
"openai",
|
||||
"compact",
|
||||
"openai:compact",
|
||||
"{\"model\":\"gpt-5\",\"input\":\"hello\"}",
|
||||
"OpenAI compact execution runtime miss did not match a Rust execution path",
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_locally_denies_openai_compact_stream_after_execution_runtime_miss_without_control_execute_opt_in(
|
||||
) {
|
||||
assert_ai_route_locally_denied_after_execution_runtime_miss(
|
||||
"/v1/responses/compact",
|
||||
"openai",
|
||||
"compact",
|
||||
"openai:compact",
|
||||
"{\"model\":\"gpt-5\",\"input\":\"hello\",\"stream\":true}",
|
||||
"OpenAI compact execution runtime miss did not match a Rust execution path",
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_locally_denies_gemini_generate_after_execution_runtime_miss_without_control_execute_opt_in(
|
||||
) {
|
||||
assert_ai_route_locally_denied_after_execution_runtime_miss(
|
||||
"/v1beta/models/gemini-2.5-pro:generateContent",
|
||||
"gemini",
|
||||
"chat",
|
||||
"gemini:chat",
|
||||
"{\"contents\":[]}",
|
||||
"Gemini public execution runtime miss did not match a Rust execution path",
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_locally_denies_gemini_v1_generate_after_execution_runtime_miss_without_control_execute_opt_in(
|
||||
) {
|
||||
assert_ai_route_locally_denied_after_execution_runtime_miss(
|
||||
"/v1/models/gemini-2.5-pro:generateContent",
|
||||
"gemini",
|
||||
"chat",
|
||||
"gemini:chat",
|
||||
"{\"contents\":[]}",
|
||||
"Gemini public execution runtime miss did not match a Rust execution path",
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_locally_denies_gemini_stream_after_execution_runtime_miss_without_control_execute_opt_in(
|
||||
) {
|
||||
assert_ai_route_locally_denied_after_execution_runtime_miss(
|
||||
"/v1beta/models/gemini-2.5-pro:streamGenerateContent",
|
||||
"gemini",
|
||||
"chat",
|
||||
"gemini:chat",
|
||||
"{\"contents\":[]}",
|
||||
"Gemini public execution runtime miss did not match a Rust execution path",
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_locally_denies_openai_video_after_execution_runtime_miss_without_control_execute_opt_in(
|
||||
) {
|
||||
assert_ai_route_locally_denied_after_execution_runtime_miss(
|
||||
"/v1/videos",
|
||||
"openai",
|
||||
"video",
|
||||
"openai:video",
|
||||
"{\"model\":\"sora-2\"}",
|
||||
"OpenAI video execution runtime miss did not match a Rust execution path",
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_locally_denies_gemini_video_after_execution_runtime_miss_without_control_execute_opt_in(
|
||||
) {
|
||||
assert_ai_route_locally_denied_after_execution_runtime_miss(
|
||||
"/v1beta/models/veo-3:predictLongRunning",
|
||||
"gemini",
|
||||
"video",
|
||||
"gemini:video",
|
||||
"{\"instances\":[]}",
|
||||
"Gemini public execution runtime miss did not match a Rust execution path",
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_locally_denies_gemini_files_root_after_execution_runtime_miss_without_control_execute_opt_in(
|
||||
) {
|
||||
assert_ai_route_locally_denied_after_execution_runtime_miss_with_request(
|
||||
reqwest::Method::GET,
|
||||
"/v1beta/files",
|
||||
"/v1beta/files?view=BASIC",
|
||||
"gemini",
|
||||
"files",
|
||||
"gemini:chat",
|
||||
None,
|
||||
"Gemini files execution runtime miss did not match a Rust execution path",
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_locally_denies_gemini_files_download_after_execution_runtime_miss_without_control_execute_opt_in(
|
||||
) {
|
||||
assert_ai_route_locally_denied_after_execution_runtime_miss_with_request(
|
||||
reqwest::Method::GET,
|
||||
"/v1beta/files/file-123:download",
|
||||
"/v1beta/files/file-123:download?alt=media",
|
||||
"gemini",
|
||||
"files",
|
||||
"gemini:chat",
|
||||
None,
|
||||
"Gemini files execution runtime miss did not match a Rust execution path",
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_locally_denies_gemini_files_upload_after_execution_runtime_miss_without_control_execute_opt_in(
|
||||
) {
|
||||
assert_ai_route_locally_denied_after_execution_runtime_miss_with_request(
|
||||
reqwest::Method::POST,
|
||||
"/upload/v1beta/files",
|
||||
"/upload/v1beta/files?uploadType=resumable",
|
||||
"gemini",
|
||||
"files",
|
||||
"gemini:chat",
|
||||
Some("{\"file\":{}}"),
|
||||
"Gemini files execution runtime miss did not match a Rust execution path",
|
||||
)
|
||||
.await;
|
||||
}
|
||||
2473
apps/aether-gateway/src/tests/ai_execute/finalize_local.rs
Normal file
2473
apps/aether-gateway/src/tests/ai_execute/finalize_local.rs
Normal file
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,28 @@
|
||||
use std::convert::Infallible;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use axum::body::{to_bytes, Body, Bytes};
|
||||
use axum::response::Response;
|
||||
use axum::routing::any;
|
||||
use axum::{extract::Request, Json, Router};
|
||||
use http::header::{HeaderName, HeaderValue};
|
||||
use http::StatusCode;
|
||||
use serde_json::json;
|
||||
|
||||
use crate::constants::{
|
||||
CONTROL_EXECUTED_HEADER, CONTROL_EXECUTE_FALLBACK_HEADER, DEPENDENCY_REASON_HEADER,
|
||||
EXECUTION_PATH_EXECUTION_RUNTIME_STREAM, EXECUTION_PATH_EXECUTION_RUNTIME_SYNC,
|
||||
EXECUTION_PATH_HEADER, EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS,
|
||||
LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER, TRACE_ID_HEADER,
|
||||
};
|
||||
|
||||
use super::{
|
||||
build_router, build_router_with_execution_runtime_override, build_router_with_state,
|
||||
build_state_with_execution_runtime_override, start_server, wait_until, AppState,
|
||||
FrontdoorCorsConfig, FrontdoorUserRpmConfig, GatewayFallbackMetricKind, GatewayFallbackReason,
|
||||
UsageRuntimeConfig, VideoTaskTruthSourceMode,
|
||||
};
|
||||
|
||||
mod compact;
|
||||
mod cross_format;
|
||||
mod direct;
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,27 @@
|
||||
use std::convert::Infallible;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use axum::body::{to_bytes, Body, Bytes};
|
||||
use axum::response::Response;
|
||||
use axum::routing::any;
|
||||
use axum::{extract::Request, Json, Router};
|
||||
use http::header::{HeaderName, HeaderValue};
|
||||
use http::StatusCode;
|
||||
use serde_json::json;
|
||||
|
||||
use crate::constants::{
|
||||
CONTROL_EXECUTED_HEADER, CONTROL_EXECUTE_FALLBACK_HEADER, DEPENDENCY_REASON_HEADER,
|
||||
EXECUTION_PATH_EXECUTION_RUNTIME_STREAM, EXECUTION_PATH_EXECUTION_RUNTIME_SYNC,
|
||||
EXECUTION_PATH_HEADER, EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS,
|
||||
LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER, TRACE_ID_HEADER,
|
||||
};
|
||||
|
||||
use super::{
|
||||
build_router, build_router_with_execution_runtime_override, build_router_with_state,
|
||||
build_state_with_execution_runtime_override, start_server, wait_until, AppState,
|
||||
FrontdoorCorsConfig, FrontdoorUserRpmConfig, GatewayFallbackMetricKind, GatewayFallbackReason,
|
||||
UsageRuntimeConfig, VideoTaskTruthSourceMode,
|
||||
};
|
||||
|
||||
mod claude;
|
||||
mod gemini;
|
||||
391
apps/aether-gateway/src/tests/ai_execute/lifecycle.rs
Normal file
391
apps/aether-gateway/src/tests/ai_execute/lifecycle.rs
Normal file
@@ -0,0 +1,391 @@
|
||||
use super::{
|
||||
any, build_router_with_state, build_state_with_execution_runtime_override, json, start_server,
|
||||
Arc, Body, Bytes, HeaderValue, Infallible, Json, Mutex, Request, Response, Router, StatusCode,
|
||||
TRACE_ID_HEADER,
|
||||
};
|
||||
|
||||
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
|
||||
use aether_data::repository::auth::{
|
||||
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot,
|
||||
};
|
||||
use aether_data::repository::candidate_selection::{
|
||||
InMemoryMinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
|
||||
StoredProviderModelMapping,
|
||||
};
|
||||
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
|
||||
use aether_data::repository::provider_catalog::{
|
||||
InMemoryProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogProvider,
|
||||
};
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
use crate::data::GatewayDataState;
|
||||
|
||||
fn hash_api_key(value: &str) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(value.as_bytes());
|
||||
format!("{:x}", hasher.finalize())
|
||||
}
|
||||
|
||||
fn sample_local_openai_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
user_id.to_string(),
|
||||
"alice".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
Some(serde_json::json!(["openai"])),
|
||||
Some(serde_json::json!(["openai:chat"])),
|
||||
Some(serde_json::json!(["gpt-5"])),
|
||||
api_key_id.to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
Some(4_102_444_800),
|
||||
Some(serde_json::json!(["openai"])),
|
||||
Some(serde_json::json!(["openai:chat"])),
|
||||
Some(serde_json::json!(["gpt-5"])),
|
||||
)
|
||||
.expect("auth snapshot should build")
|
||||
}
|
||||
|
||||
fn sample_local_openai_candidate_row() -> StoredMinimalCandidateSelectionRow {
|
||||
StoredMinimalCandidateSelectionRow {
|
||||
provider_id: "provider-openai-lifecycle-local-1".to_string(),
|
||||
provider_name: "openai".to_string(),
|
||||
provider_type: "custom".to_string(),
|
||||
provider_priority: 10,
|
||||
provider_is_active: true,
|
||||
endpoint_id: "endpoint-openai-lifecycle-local-1".to_string(),
|
||||
endpoint_api_format: "openai:chat".to_string(),
|
||||
endpoint_api_family: Some("openai".to_string()),
|
||||
endpoint_kind: Some("chat".to_string()),
|
||||
endpoint_is_active: true,
|
||||
key_id: "key-openai-lifecycle-local-1".to_string(),
|
||||
key_name: "prod".to_string(),
|
||||
key_auth_type: "api_key".to_string(),
|
||||
key_is_active: true,
|
||||
key_api_formats: Some(vec!["openai:chat".to_string()]),
|
||||
key_allowed_models: None,
|
||||
key_capabilities: None,
|
||||
key_internal_priority: 5,
|
||||
key_global_priority_by_format: Some(serde_json::json!({"openai:chat": 1})),
|
||||
model_id: "model-openai-lifecycle-local-1".to_string(),
|
||||
global_model_id: "global-model-openai-lifecycle-local-1".to_string(),
|
||||
global_model_name: "gpt-5".to_string(),
|
||||
global_model_mappings: None,
|
||||
global_model_supports_streaming: Some(true),
|
||||
model_provider_model_name: "gpt-5-upstream".to_string(),
|
||||
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
|
||||
name: "gpt-5-upstream".to_string(),
|
||||
priority: 1,
|
||||
api_formats: Some(vec!["openai:chat".to_string()]),
|
||||
}]),
|
||||
model_supports_streaming: Some(true),
|
||||
model_is_active: true,
|
||||
model_is_available: true,
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_local_openai_provider() -> StoredProviderCatalogProvider {
|
||||
StoredProviderCatalogProvider::new(
|
||||
"provider-openai-lifecycle-local-1".to_string(),
|
||||
"openai".to_string(),
|
||||
Some("https://example.com".to_string()),
|
||||
"custom".to_string(),
|
||||
)
|
||||
.expect("provider should build")
|
||||
.with_transport_fields(
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
None,
|
||||
Some(2),
|
||||
None,
|
||||
Some(20.0),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
fn sample_local_openai_endpoint() -> StoredProviderCatalogEndpoint {
|
||||
StoredProviderCatalogEndpoint::new(
|
||||
"endpoint-openai-lifecycle-local-1".to_string(),
|
||||
"provider-openai-lifecycle-local-1".to_string(),
|
||||
"openai:chat".to_string(),
|
||||
Some("openai".to_string()),
|
||||
Some("chat".to_string()),
|
||||
true,
|
||||
)
|
||||
.expect("endpoint should build")
|
||||
.with_transport_fields(
|
||||
"https://api.openai.example".to_string(),
|
||||
None,
|
||||
None,
|
||||
Some(2),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("endpoint transport should build")
|
||||
}
|
||||
|
||||
fn sample_local_openai_key() -> StoredProviderCatalogKey {
|
||||
StoredProviderCatalogKey::new(
|
||||
"key-openai-lifecycle-local-1".to_string(),
|
||||
"provider-openai-lifecycle-local-1".to_string(),
|
||||
"prod".to_string(),
|
||||
"api_key".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
.with_transport_fields(
|
||||
Some(serde_json::json!(["openai:chat"])),
|
||||
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "sk-upstream-openai")
|
||||
.expect("api key should encrypt"),
|
||||
None,
|
||||
None,
|
||||
Some(serde_json::json!({"openai:chat": 1})),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("key transport should build")
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn gateway_completes_sync_response_on_local_execution_runtime_path() {
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/v1/chat/completions",
|
||||
any(move |_request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::IM_A_TEAPOT, Body::from("public"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let execution_runtime = Router::new().route(
|
||||
"/v1/execute/sync",
|
||||
any(|_request: Request| async move {
|
||||
Json(json!({
|
||||
"request_id": "req-openai-chat-async-report-123",
|
||||
"status_code": 200,
|
||||
"headers": {
|
||||
"content-type": "application/json"
|
||||
},
|
||||
"body": {
|
||||
"json_body": {
|
||||
"id": "chatcmpl-async-report-123",
|
||||
"object": "chat.completion",
|
||||
"model": "gpt-5",
|
||||
"choices": [],
|
||||
"usage": {
|
||||
"prompt_tokens": 1,
|
||||
"completion_tokens": 2,
|
||||
"total_tokens": 3
|
||||
}
|
||||
}
|
||||
},
|
||||
"telemetry": {
|
||||
"elapsed_ms": 12
|
||||
}
|
||||
}))
|
||||
}),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("sk-client-openai-async-report")),
|
||||
sample_local_openai_auth_snapshot(
|
||||
"api-key-openai-lifecycle-local-1",
|
||||
"user-openai-lifecycle-local-1",
|
||||
),
|
||||
)]));
|
||||
let candidate_selection_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
sample_local_openai_candidate_row(),
|
||||
]));
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_local_openai_provider()],
|
||||
vec![sample_local_openai_endpoint()],
|
||||
vec![sample_local_openai_key()],
|
||||
));
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
let gateway = build_router_with_state(
|
||||
build_state_with_execution_runtime_override(execution_runtime_url)
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
|
||||
auth_repository,
|
||||
candidate_selection_repository,
|
||||
provider_catalog_repository,
|
||||
request_candidate_repository,
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/chat/completions"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(
|
||||
http::header::AUTHORIZATION,
|
||||
"Bearer sk-client-openai-async-report",
|
||||
)
|
||||
.header(TRACE_ID_HEADER, "req-openai-chat-async-report-123")
|
||||
.body("{\"model\":\"gpt-5\",\"messages\":[]}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn gateway_stops_execution_runtime_stream_when_client_disconnects() {
|
||||
let seen_report = Arc::new(Mutex::new(0usize));
|
||||
let seen_report_clone = Arc::clone(&seen_report);
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/report-stream",
|
||||
any(move |_request: Request| {
|
||||
let seen_report_inner = Arc::clone(&seen_report_clone);
|
||||
async move {
|
||||
*seen_report_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({"ok": true}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1/chat/completions",
|
||||
any(move |_request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::IM_A_TEAPOT, Body::from("public-route-hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let execution_runtime = Router::new().route(
|
||||
"/v1/execute/stream",
|
||||
any(|_request: Request| async move {
|
||||
let body_stream = async_stream::stream! {
|
||||
yield Ok::<Bytes, Infallible>(Bytes::from_static(
|
||||
b"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n"
|
||||
));
|
||||
yield Ok::<Bytes, Infallible>(Bytes::from_static(
|
||||
b"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"id\\\":\\\"chatcmpl-first\\\"}\\n\\n\"}}\n"
|
||||
));
|
||||
tokio::time::sleep(std::time::Duration::from_millis(120)).await;
|
||||
yield Ok::<Bytes, Infallible>(Bytes::from_static(
|
||||
b"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: [DONE]\\n\\n\"}}\n"
|
||||
));
|
||||
yield Ok::<Bytes, Infallible>(Bytes::from_static(
|
||||
b"{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":41,\"ttfb_ms\":12,\"upstream_bytes\":31}}}\n"
|
||||
));
|
||||
yield Ok::<Bytes, Infallible>(Bytes::from_static(
|
||||
b"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n"
|
||||
));
|
||||
};
|
||||
let mut response = Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from_stream(body_stream))
|
||||
.expect("response should build");
|
||||
response.headers_mut().insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
HeaderValue::from_static("application/x-ndjson"),
|
||||
);
|
||||
response
|
||||
}),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("sk-client-openai-stream-disconnect")),
|
||||
sample_local_openai_auth_snapshot(
|
||||
"api-key-openai-lifecycle-local-1",
|
||||
"user-openai-lifecycle-local-1",
|
||||
),
|
||||
)]));
|
||||
let candidate_selection_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
sample_local_openai_candidate_row(),
|
||||
]));
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_local_openai_provider()],
|
||||
vec![sample_local_openai_endpoint()],
|
||||
vec![sample_local_openai_key()],
|
||||
));
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
let gateway = build_router_with_state(
|
||||
build_state_with_execution_runtime_override(execution_runtime_url)
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
|
||||
auth_repository,
|
||||
candidate_selection_repository,
|
||||
provider_catalog_repository,
|
||||
request_candidate_repository,
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/chat/completions"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(
|
||||
http::header::AUTHORIZATION,
|
||||
"Bearer sk-client-openai-stream-disconnect",
|
||||
)
|
||||
.header(TRACE_ID_HEADER, "trace-openai-chat-stream-disconnect-123")
|
||||
.body("{\"model\":\"gpt-5\",\"messages\":[],\"stream\":true}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let mut response = response;
|
||||
let first_chunk = response
|
||||
.chunk()
|
||||
.await
|
||||
.expect("first chunk should read")
|
||||
.expect("first chunk should exist");
|
||||
assert_eq!(
|
||||
first_chunk,
|
||||
Bytes::from_static(b"data: {\"id\":\"chatcmpl-first\"}\n\n")
|
||||
);
|
||||
drop(response);
|
||||
|
||||
tokio::time::sleep(std::time::Duration::from_millis(500)).await;
|
||||
|
||||
assert_eq!(*seen_report.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
36
apps/aether-gateway/src/tests/ai_execute/mod.rs
Normal file
36
apps/aether-gateway/src/tests/ai_execute/mod.rs
Normal file
@@ -0,0 +1,36 @@
|
||||
use std::convert::Infallible;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use axum::body::{to_bytes, Body, Bytes};
|
||||
use axum::response::Response;
|
||||
use axum::routing::any;
|
||||
use axum::{extract::Request, Json, Router};
|
||||
use http::header::{HeaderName, HeaderValue};
|
||||
use http::StatusCode;
|
||||
use serde_json::json;
|
||||
|
||||
use crate::constants::{
|
||||
CONTROL_EXECUTED_HEADER, CONTROL_EXECUTE_FALLBACK_HEADER, DEPENDENCY_REASON_HEADER,
|
||||
EXECUTION_PATH_EXECUTION_RUNTIME_STREAM, EXECUTION_PATH_EXECUTION_RUNTIME_SYNC,
|
||||
EXECUTION_PATH_HEADER, EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS,
|
||||
LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER, TRACE_ID_HEADER,
|
||||
};
|
||||
|
||||
use super::{
|
||||
build_router, build_router_with_execution_runtime_override, build_router_with_state,
|
||||
build_state_with_execution_runtime_override, start_server, wait_until, AppState,
|
||||
FrontdoorCorsConfig, FrontdoorUserRpmConfig, GatewayFallbackMetricKind, GatewayFallbackReason,
|
||||
UsageRuntimeConfig, VideoTaskTruthSourceMode,
|
||||
};
|
||||
|
||||
mod control_execute;
|
||||
mod fallback;
|
||||
mod finalize_local;
|
||||
mod finalize_local_cli;
|
||||
mod finalize_local_provider;
|
||||
mod lifecycle;
|
||||
mod stream;
|
||||
mod stream_cli;
|
||||
mod stream_provider;
|
||||
mod stream_provider_gemini;
|
||||
mod sync;
|
||||
1872
apps/aether-gateway/src/tests/ai_execute/stream/decision.rs
Normal file
1872
apps/aether-gateway/src/tests/ai_execute/stream/decision.rs
Normal file
File diff suppressed because it is too large
Load Diff
26
apps/aether-gateway/src/tests/ai_execute/stream/mod.rs
Normal file
26
apps/aether-gateway/src/tests/ai_execute/stream/mod.rs
Normal file
@@ -0,0 +1,26 @@
|
||||
use std::convert::Infallible;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use axum::body::{to_bytes, Body, Bytes};
|
||||
use axum::response::Response;
|
||||
use axum::routing::any;
|
||||
use axum::{extract::Request, Json, Router};
|
||||
use http::header::{HeaderName, HeaderValue};
|
||||
use http::StatusCode;
|
||||
use serde_json::json;
|
||||
|
||||
use crate::constants::{
|
||||
CONTROL_EXECUTED_HEADER, CONTROL_EXECUTE_FALLBACK_HEADER, DEPENDENCY_REASON_HEADER,
|
||||
EXECUTION_PATH_EXECUTION_RUNTIME_STREAM, EXECUTION_PATH_EXECUTION_RUNTIME_SYNC,
|
||||
EXECUTION_PATH_HEADER, EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS,
|
||||
LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER, TRACE_ID_HEADER,
|
||||
};
|
||||
|
||||
use super::{
|
||||
build_router, build_router_with_execution_runtime_override, build_router_with_state,
|
||||
build_state_with_execution_runtime_override, start_server, wait_until, AppState,
|
||||
FrontdoorCorsConfig, FrontdoorUserRpmConfig, GatewayFallbackMetricKind, GatewayFallbackReason,
|
||||
UsageRuntimeConfig, VideoTaskTruthSourceMode,
|
||||
};
|
||||
|
||||
mod decision;
|
||||
558
apps/aether-gateway/src/tests/ai_execute/stream_cli/compact.rs
Normal file
558
apps/aether-gateway/src/tests/ai_execute/stream_cli/compact.rs
Normal file
@@ -0,0 +1,558 @@
|
||||
use super::{
|
||||
any, build_router_with_state, build_state_with_execution_runtime_override, json, start_server,
|
||||
to_bytes, Arc, Body, Bytes, HeaderName, HeaderValue, Infallible, Json, Mutex, Request,
|
||||
Response, Router, StatusCode, EXECUTION_PATH_EXECUTION_RUNTIME_STREAM, EXECUTION_PATH_HEADER,
|
||||
TRACE_ID_HEADER,
|
||||
};
|
||||
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
|
||||
use aether_data::repository::auth::{
|
||||
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot,
|
||||
};
|
||||
use aether_data::repository::candidate_selection::{
|
||||
InMemoryMinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
|
||||
StoredProviderModelMapping,
|
||||
};
|
||||
use aether_data::repository::candidates::{
|
||||
InMemoryRequestCandidateRepository, RequestCandidateReadRepository, RequestCandidateStatus,
|
||||
};
|
||||
use aether_data::repository::provider_catalog::{
|
||||
InMemoryProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogProvider,
|
||||
};
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_executes_openai_compact_stream_via_local_decision_gate_with_local_stream_decision()
|
||||
{
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenExecutionRuntimeStreamRequest {
|
||||
trace_id: String,
|
||||
url: String,
|
||||
model: String,
|
||||
stream: bool,
|
||||
accept: String,
|
||||
authorization: String,
|
||||
endpoint_tag: String,
|
||||
conditional_header: String,
|
||||
renamed_header: String,
|
||||
dropped_header_present: bool,
|
||||
metadata_mode: String,
|
||||
metadata_source: String,
|
||||
metadata_origin: String,
|
||||
instructions: String,
|
||||
store_present: bool,
|
||||
proxy_node_id: String,
|
||||
tls_profile: String,
|
||||
}
|
||||
|
||||
fn hash_api_key(value: &str) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(value.as_bytes());
|
||||
format!("{:x}", hasher.finalize())
|
||||
}
|
||||
|
||||
fn sample_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
user_id.to_string(),
|
||||
"alice".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
Some(serde_json::json!(["openai"])),
|
||||
Some(serde_json::json!(["openai:compact"])),
|
||||
Some(serde_json::json!(["gpt-5"])),
|
||||
api_key_id.to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
Some(4_102_444_800_i64),
|
||||
Some(serde_json::json!(["openai"])),
|
||||
Some(serde_json::json!(["openai:compact"])),
|
||||
Some(serde_json::json!(["gpt-5"])),
|
||||
)
|
||||
.expect("auth snapshot should build")
|
||||
}
|
||||
|
||||
fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow {
|
||||
StoredMinimalCandidateSelectionRow {
|
||||
provider_id: "provider-openai-compact-local-1".to_string(),
|
||||
provider_name: "openai".to_string(),
|
||||
provider_type: "custom".to_string(),
|
||||
provider_priority: 10,
|
||||
provider_is_active: true,
|
||||
endpoint_id: "endpoint-openai-compact-local-1".to_string(),
|
||||
endpoint_api_format: "openai:compact".to_string(),
|
||||
endpoint_api_family: Some("openai".to_string()),
|
||||
endpoint_kind: Some("compact".to_string()),
|
||||
endpoint_is_active: true,
|
||||
key_id: "key-openai-compact-local-1".to_string(),
|
||||
key_name: "prod".to_string(),
|
||||
key_auth_type: "bearer".to_string(),
|
||||
key_is_active: true,
|
||||
key_api_formats: Some(vec!["openai:compact".to_string()]),
|
||||
key_allowed_models: None,
|
||||
key_capabilities: None,
|
||||
key_internal_priority: 5,
|
||||
key_global_priority_by_format: Some(serde_json::json!({"openai:compact": 1})),
|
||||
model_id: "model-openai-compact-local-1".to_string(),
|
||||
global_model_id: "global-model-openai-compact-local-1".to_string(),
|
||||
global_model_name: "gpt-5".to_string(),
|
||||
global_model_mappings: None,
|
||||
global_model_supports_streaming: Some(true),
|
||||
model_provider_model_name: "gpt-5-upstream".to_string(),
|
||||
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
|
||||
name: "gpt-5-upstream".to_string(),
|
||||
priority: 1,
|
||||
api_formats: Some(vec!["openai:compact".to_string()]),
|
||||
}]),
|
||||
model_supports_streaming: Some(true),
|
||||
model_is_active: true,
|
||||
model_is_available: true,
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
|
||||
StoredProviderCatalogProvider::new(
|
||||
"provider-openai-compact-local-1".to_string(),
|
||||
"openai".to_string(),
|
||||
Some("https://example.com".to_string()),
|
||||
"custom".to_string(),
|
||||
)
|
||||
.expect("provider should build")
|
||||
.with_transport_fields(
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
None,
|
||||
Some(2),
|
||||
Some(serde_json::json!({"url":"http://provider-proxy.internal:8080"})),
|
||||
Some(20.0),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
|
||||
StoredProviderCatalogEndpoint::new(
|
||||
"endpoint-openai-compact-local-1".to_string(),
|
||||
"provider-openai-compact-local-1".to_string(),
|
||||
"openai:compact".to_string(),
|
||||
Some("openai".to_string()),
|
||||
Some("compact".to_string()),
|
||||
true,
|
||||
)
|
||||
.expect("endpoint should build")
|
||||
.with_transport_fields(
|
||||
"https://api.openai.example".to_string(),
|
||||
Some(serde_json::json!([
|
||||
{"action":"set","key":"x-endpoint-tag","value":"openai-compact-local"},
|
||||
{"action":"set","key":"x-conditional-tag","value":"header-condition-hit","condition":{"path":"instructions","op":"exists","source":"current"}},
|
||||
{"action":"rename","from":"x-client-rename","to":"x-upstream-rename"},
|
||||
{"action":"drop","key":"x-drop-me"}
|
||||
])),
|
||||
Some(serde_json::json!([
|
||||
{"action":"set","path":"instructions","value":"You are GPT-5.","condition":{"path":"instructions","op":"not_exists","source":"current"}},
|
||||
{"action":"set","path":"metadata.mode","value":"safe","condition":{"path":"metadata.mode","op":"not_exists","source":"current"}},
|
||||
{"action":"rename","from":"metadata.client","to":"metadata.source"},
|
||||
{"action":"set","path":"metadata.origin","value":"from-original","condition":{"path":"metadata.client","op":"exists","source":"original"}},
|
||||
{"action":"drop","path":"store"}
|
||||
])),
|
||||
Some(2),
|
||||
Some("/custom/v1/responses/compact".to_string()),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("endpoint transport should build")
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
|
||||
StoredProviderCatalogKey::new(
|
||||
"key-openai-compact-local-1".to_string(),
|
||||
"provider-openai-compact-local-1".to_string(),
|
||||
"prod".to_string(),
|
||||
"bearer".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
.with_transport_fields(
|
||||
Some(serde_json::json!(["openai:compact"])),
|
||||
encrypt_python_fernet_plaintext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
"sk-upstream-openai-compact",
|
||||
)
|
||||
.expect("api key should encrypt"),
|
||||
None,
|
||||
None,
|
||||
Some(serde_json::json!({"openai:compact": 1})),
|
||||
None,
|
||||
None,
|
||||
Some(serde_json::json!({"enabled": true, "node_id":"proxy-node-openai-compact-local"})),
|
||||
Some(serde_json::json!({"tls_profile":"chrome_136"})),
|
||||
)
|
||||
.expect("key transport should build")
|
||||
}
|
||||
|
||||
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeStreamRequest>));
|
||||
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
|
||||
let seen_report = Arc::new(Mutex::new(false));
|
||||
let seen_report_clone = Arc::clone(&seen_report);
|
||||
let decision_hits = Arc::new(Mutex::new(0usize));
|
||||
let decision_hits_clone = Arc::clone(&decision_hits);
|
||||
let plan_hits = Arc::new(Mutex::new(0usize));
|
||||
let plan_hits_clone = Arc::clone(&plan_hits);
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/resolve",
|
||||
any(|_request: Request| async move {
|
||||
Json(json!({
|
||||
"action": "proxy_public",
|
||||
"route_class": "ai_public",
|
||||
"route_family": "openai",
|
||||
"route_kind": "compact",
|
||||
"auth_endpoint_signature": "openai:compact",
|
||||
"execution_runtime_candidate": true,
|
||||
"auth_context": {
|
||||
"user_id": "user-openai-compact-local-123",
|
||||
"api_key_id": "key-openai-compact-local-123",
|
||||
"access_allowed": true
|
||||
},
|
||||
"public_path": "/v1/responses/compact"
|
||||
}))
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/decision-stream",
|
||||
any(move |_request: Request| {
|
||||
let decision_hits_inner = Arc::clone(&decision_hits_clone);
|
||||
async move {
|
||||
*decision_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({"action": "proxy_public"}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/plan-stream",
|
||||
any(move |_request: Request| {
|
||||
let plan_hits_inner = Arc::clone(&plan_hits_clone);
|
||||
async move {
|
||||
*plan_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({"action": "proxy_public"}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/report-stream",
|
||||
any(move |request: Request| {
|
||||
let seen_report_inner = Arc::clone(&seen_report_clone);
|
||||
async move {
|
||||
let (_parts, body) = request.into_parts();
|
||||
let _raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
|
||||
*seen_report_inner.lock().expect("mutex should lock") = true;
|
||||
Json(json!({"ok": true}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1/responses/compact",
|
||||
any(move |_request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
let stream = futures_util::stream::iter([
|
||||
Ok::<_, Infallible>(Bytes::from_static(b"event: response.completed\n")),
|
||||
Ok::<_, Infallible>(Bytes::from_static(
|
||||
b"data: {\"type\":\"response.completed\"}\n\n",
|
||||
)),
|
||||
]);
|
||||
let mut response = Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from_stream(stream))
|
||||
.expect("response should build");
|
||||
response.headers_mut().insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
HeaderValue::from_static("text/event-stream"),
|
||||
);
|
||||
response
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let execution_runtime = Router::new().route(
|
||||
"/v1/execute/stream",
|
||||
any(move |request: Request| {
|
||||
let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone);
|
||||
async move {
|
||||
let (parts, body) = request.into_parts();
|
||||
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
|
||||
let payload: serde_json::Value =
|
||||
serde_json::from_slice(&raw_body).expect("execution runtime payload should parse");
|
||||
*seen_execution_runtime_inner.lock().expect("mutex should lock") =
|
||||
Some(SeenExecutionRuntimeStreamRequest {
|
||||
trace_id: parts
|
||||
.headers
|
||||
.get(TRACE_ID_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
url: payload
|
||||
.get("url")
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
model: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("model"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
stream: payload
|
||||
.get("stream")
|
||||
.and_then(|value| value.as_bool())
|
||||
.unwrap_or(false),
|
||||
accept: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("accept"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
authorization: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("authorization"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
endpoint_tag: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-endpoint-tag"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
conditional_header: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-conditional-tag"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
renamed_header: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-upstream-rename"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
dropped_header_present: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-drop-me"))
|
||||
.is_some(),
|
||||
metadata_mode: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("metadata"))
|
||||
.and_then(|value| value.get("mode"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
metadata_source: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("metadata"))
|
||||
.and_then(|value| value.get("source"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
metadata_origin: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("metadata"))
|
||||
.and_then(|value| value.get("origin"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
instructions: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("instructions"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
store_present: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("store"))
|
||||
.is_some(),
|
||||
proxy_node_id: payload
|
||||
.get("proxy")
|
||||
.and_then(|value| value.get("node_id"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
tls_profile: payload
|
||||
.get("tls_profile")
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
});
|
||||
let stream = concat!(
|
||||
"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n",
|
||||
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: response.completed\\ndata: {\\\"type\\\":\\\"response.completed\\\",\\\"response\\\":{\\\"id\\\":\\\"resp-compact-local-123\\\",\\\"object\\\":\\\"response\\\",\\\"model\\\":\\\"gpt-5-upstream\\\",\\\"output\\\":[]}}\\n\\n\"}}\n",
|
||||
"{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":41,\"ttfb_ms\":11}}}\n",
|
||||
"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n"
|
||||
);
|
||||
let mut response = Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from(stream))
|
||||
.expect("response should build");
|
||||
response.headers_mut().insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
HeaderValue::from_static("application/x-ndjson"),
|
||||
);
|
||||
response
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let client_api_key = "sk-client-openai-compact-local";
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key(client_api_key)),
|
||||
sample_auth_snapshot(
|
||||
"key-openai-compact-local-123",
|
||||
"user-openai-compact-local-123",
|
||||
),
|
||||
)]));
|
||||
let candidate_selection_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
sample_candidate_row(),
|
||||
]));
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider_catalog_provider()],
|
||||
vec![sample_provider_catalog_endpoint()],
|
||||
vec![sample_provider_catalog_key()],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url.clone())
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
|
||||
auth_repository,
|
||||
candidate_selection_repository,
|
||||
provider_catalog_repository,
|
||||
Arc::clone(&request_candidate_repository),
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
),
|
||||
);
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/responses/compact"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(
|
||||
http::header::AUTHORIZATION,
|
||||
format!("Bearer {client_api_key}"),
|
||||
)
|
||||
.header("x-client-rename", "rename-openai-compact")
|
||||
.header("x-drop-me", "drop-openai-compact")
|
||||
.header(TRACE_ID_HEADER, "trace-openai-compact-local-123")
|
||||
.body("{\"model\":\"gpt-5\",\"input\":\"hello\",\"stream\":true,\"metadata\":{\"client\":\"desktop-openai-compact\"},\"store\":false}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some(EXECUTION_PATH_EXECUTION_RUNTIME_STREAM)
|
||||
);
|
||||
let body = response.text().await.expect("body should read");
|
||||
assert!(body.contains("event: response.completed"));
|
||||
assert!(body.contains("\"model\":\"gpt-5-upstream\""));
|
||||
|
||||
let seen_execution_runtime_request = seen_execution_runtime
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.clone()
|
||||
.expect("execution runtime stream should be captured");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.trace_id,
|
||||
"trace-openai-compact-local-123"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.url,
|
||||
"https://api.openai.example/custom/v1/responses/compact"
|
||||
);
|
||||
assert_eq!(seen_execution_runtime_request.model, "gpt-5-upstream");
|
||||
assert!(seen_execution_runtime_request.stream);
|
||||
assert_eq!(seen_execution_runtime_request.accept, "text/event-stream");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.authorization,
|
||||
"Bearer sk-upstream-openai-compact"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.endpoint_tag,
|
||||
"openai-compact-local"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.conditional_header,
|
||||
"header-condition-hit"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.renamed_header,
|
||||
"rename-openai-compact"
|
||||
);
|
||||
assert!(!seen_execution_runtime_request.dropped_header_present);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.instructions,
|
||||
"You are GPT-5."
|
||||
);
|
||||
assert_eq!(seen_execution_runtime_request.metadata_mode, "safe");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.metadata_source,
|
||||
"desktop-openai-compact"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.metadata_origin,
|
||||
"from-original"
|
||||
);
|
||||
assert!(!seen_execution_runtime_request.store_present);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.proxy_node_id,
|
||||
"proxy-node-openai-compact-local"
|
||||
);
|
||||
assert_eq!(seen_execution_runtime_request.tls_profile, "chrome_136");
|
||||
|
||||
let stored_candidates = request_candidate_repository
|
||||
.list_by_request_id("trace-openai-compact-local-123")
|
||||
.await
|
||||
.expect("request candidate trace should read");
|
||||
assert_eq!(stored_candidates.len(), 1);
|
||||
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success);
|
||||
|
||||
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
|
||||
assert!(
|
||||
!*seen_report.lock().expect("mutex should lock"),
|
||||
"report-stream should stay local when request candidate persistence is available"
|
||||
);
|
||||
|
||||
assert_eq!(*decision_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*plan_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
516
apps/aether-gateway/src/tests/ai_execute/stream_cli/direct.rs
Normal file
516
apps/aether-gateway/src/tests/ai_execute/stream_cli/direct.rs
Normal file
@@ -0,0 +1,516 @@
|
||||
use super::{
|
||||
any, build_router_with_state, build_state_with_execution_runtime_override, json, start_server,
|
||||
to_bytes, Arc, Body, Bytes, HeaderName, HeaderValue, Infallible, Json, Mutex, Request,
|
||||
Response, Router, StatusCode, TRACE_ID_HEADER,
|
||||
};
|
||||
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
|
||||
use aether_data::repository::auth::{
|
||||
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot,
|
||||
};
|
||||
use aether_data::repository::candidate_selection::{
|
||||
InMemoryMinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
|
||||
StoredProviderModelMapping,
|
||||
};
|
||||
use aether_data::repository::candidates::{
|
||||
InMemoryRequestCandidateRepository, RequestCandidateReadRepository, RequestCandidateStatus,
|
||||
};
|
||||
use aether_data::repository::provider_catalog::{
|
||||
InMemoryProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogProvider,
|
||||
};
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_executes_codex_cli_stream_via_local_decision_gate_after_oauth_refresh() {
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenExecutionRuntimeStreamRequest {
|
||||
trace_id: String,
|
||||
url: String,
|
||||
model: String,
|
||||
stream: bool,
|
||||
accept: String,
|
||||
authorization: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenReportStreamRequest {
|
||||
trace_id: String,
|
||||
report_kind: String,
|
||||
request_id: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenRefreshRequest {
|
||||
content_type: String,
|
||||
body: String,
|
||||
}
|
||||
|
||||
fn hash_api_key(value: &str) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(value.as_bytes());
|
||||
format!("{:x}", hasher.finalize())
|
||||
}
|
||||
|
||||
fn sample_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
user_id.to_string(),
|
||||
"alice".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
Some(serde_json::json!(["openai", "codex"])),
|
||||
Some(serde_json::json!(["openai:cli"])),
|
||||
Some(serde_json::json!(["gpt-5.4"])),
|
||||
api_key_id.to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
Some(4_102_444_800_i64),
|
||||
Some(serde_json::json!(["openai", "codex"])),
|
||||
Some(serde_json::json!(["openai:cli"])),
|
||||
Some(serde_json::json!(["gpt-5.4"])),
|
||||
)
|
||||
.expect("auth snapshot should build")
|
||||
}
|
||||
|
||||
fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow {
|
||||
StoredMinimalCandidateSelectionRow {
|
||||
provider_id: "provider-codex-cli-stream-local-1".to_string(),
|
||||
provider_name: "codex".to_string(),
|
||||
provider_type: "codex".to_string(),
|
||||
provider_priority: 10,
|
||||
provider_is_active: true,
|
||||
endpoint_id: "endpoint-codex-cli-stream-local-1".to_string(),
|
||||
endpoint_api_format: "openai:cli".to_string(),
|
||||
endpoint_api_family: Some("openai".to_string()),
|
||||
endpoint_kind: Some("cli".to_string()),
|
||||
endpoint_is_active: true,
|
||||
key_id: "key-codex-cli-stream-local-1".to_string(),
|
||||
key_name: "oauth".to_string(),
|
||||
key_auth_type: "oauth".to_string(),
|
||||
key_is_active: true,
|
||||
key_api_formats: Some(vec!["openai:cli".to_string()]),
|
||||
key_allowed_models: None,
|
||||
key_capabilities: None,
|
||||
key_internal_priority: 5,
|
||||
key_global_priority_by_format: Some(serde_json::json!({"openai:cli": 1})),
|
||||
model_id: "model-codex-cli-stream-local-1".to_string(),
|
||||
global_model_id: "global-model-codex-cli-stream-local-1".to_string(),
|
||||
global_model_name: "gpt-5.4".to_string(),
|
||||
global_model_mappings: None,
|
||||
global_model_supports_streaming: Some(true),
|
||||
model_provider_model_name: "gpt-5.4".to_string(),
|
||||
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
|
||||
name: "gpt-5.4".to_string(),
|
||||
priority: 1,
|
||||
api_formats: Some(vec!["openai:cli".to_string()]),
|
||||
}]),
|
||||
model_supports_streaming: Some(true),
|
||||
model_is_active: true,
|
||||
model_is_available: true,
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
|
||||
StoredProviderCatalogProvider::new(
|
||||
"provider-codex-cli-stream-local-1".to_string(),
|
||||
"codex".to_string(),
|
||||
Some("https://chatgpt.com".to_string()),
|
||||
"codex".to_string(),
|
||||
)
|
||||
.expect("provider should build")
|
||||
.with_transport_fields(
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
None,
|
||||
Some(2),
|
||||
None,
|
||||
Some(20.0),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
|
||||
StoredProviderCatalogEndpoint::new(
|
||||
"endpoint-codex-cli-stream-local-1".to_string(),
|
||||
"provider-codex-cli-stream-local-1".to_string(),
|
||||
"openai:cli".to_string(),
|
||||
Some("openai".to_string()),
|
||||
Some("cli".to_string()),
|
||||
true,
|
||||
)
|
||||
.expect("endpoint should build")
|
||||
.with_transport_fields(
|
||||
"https://chatgpt.com/backend-api/codex".to_string(),
|
||||
None,
|
||||
None,
|
||||
Some(2),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("endpoint transport should build")
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
|
||||
let encrypted_auth_config = encrypt_python_fernet_plaintext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
r#"{"provider_type":"codex","refresh_token":"rt-codex-stream-local-123"}"#,
|
||||
)
|
||||
.expect("auth config should encrypt");
|
||||
StoredProviderCatalogKey::new(
|
||||
"key-codex-cli-stream-local-1".to_string(),
|
||||
"provider-codex-cli-stream-local-1".to_string(),
|
||||
"oauth".to_string(),
|
||||
"oauth".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
.with_transport_fields(
|
||||
Some(serde_json::json!(["openai:cli"])),
|
||||
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "__placeholder__")
|
||||
.expect("placeholder api key should encrypt"),
|
||||
Some(encrypted_auth_config),
|
||||
None,
|
||||
Some(serde_json::json!({"openai:cli": 1})),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("key transport should build")
|
||||
}
|
||||
|
||||
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeStreamRequest>));
|
||||
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
|
||||
let seen_report = Arc::new(Mutex::new(None::<SeenReportStreamRequest>));
|
||||
let seen_report_clone = Arc::clone(&seen_report);
|
||||
let seen_refresh = Arc::new(Mutex::new(None::<SeenRefreshRequest>));
|
||||
let seen_refresh_clone = Arc::clone(&seen_refresh);
|
||||
let decision_hits = Arc::new(Mutex::new(0usize));
|
||||
let decision_hits_clone = Arc::clone(&decision_hits);
|
||||
let plan_hits = Arc::new(Mutex::new(0usize));
|
||||
let plan_hits_clone = Arc::clone(&plan_hits);
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/resolve",
|
||||
any(|_request: Request| async move {
|
||||
Json(json!({
|
||||
"action": "proxy_public",
|
||||
"route_class": "ai_public",
|
||||
"route_family": "openai",
|
||||
"route_kind": "cli",
|
||||
"auth_endpoint_signature": "openai:cli",
|
||||
"execution_runtime_candidate": true,
|
||||
"auth_context": {
|
||||
"user_id": "user-codex-cli-stream-local-123",
|
||||
"api_key_id": "key-codex-cli-stream-local-123",
|
||||
"access_allowed": true
|
||||
},
|
||||
"public_path": "/v1/responses"
|
||||
}))
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/decision-stream",
|
||||
any(move |_request: Request| {
|
||||
let decision_hits_inner = Arc::clone(&decision_hits_clone);
|
||||
async move {
|
||||
*decision_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({"action": "proxy_public"}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/plan-stream",
|
||||
any(move |_request: Request| {
|
||||
let plan_hits_inner = Arc::clone(&plan_hits_clone);
|
||||
async move {
|
||||
*plan_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({"action": "proxy_public"}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/report-stream",
|
||||
any(move |request: Request| {
|
||||
let seen_report_inner = Arc::clone(&seen_report_clone);
|
||||
async move {
|
||||
let (parts, body) = request.into_parts();
|
||||
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
|
||||
let payload: serde_json::Value =
|
||||
serde_json::from_slice(&raw_body).expect("report payload should parse");
|
||||
*seen_report_inner.lock().expect("mutex should lock") =
|
||||
Some(SeenReportStreamRequest {
|
||||
trace_id: parts
|
||||
.headers
|
||||
.get(TRACE_ID_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
report_kind: payload
|
||||
.get("report_kind")
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
request_id: payload
|
||||
.get("report_context")
|
||||
.and_then(|value| value.get("request_id"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
});
|
||||
Json(json!({"ok": true}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1/responses",
|
||||
any(move |_request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
let stream = futures_util::stream::iter([Ok::<_, Infallible>(
|
||||
Bytes::from_static(b"unexpected"),
|
||||
)]);
|
||||
let mut response = Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from_stream(stream))
|
||||
.expect("response should build");
|
||||
response.headers_mut().insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
HeaderValue::from_static("text/plain"),
|
||||
);
|
||||
response
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let refresh = Router::new().route(
|
||||
"/oauth/token",
|
||||
any(move |request: Request| {
|
||||
let seen_refresh_inner = Arc::clone(&seen_refresh_clone);
|
||||
async move {
|
||||
let (parts, body) = request.into_parts();
|
||||
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
|
||||
*seen_refresh_inner.lock().expect("mutex should lock") = Some(SeenRefreshRequest {
|
||||
content_type: parts
|
||||
.headers
|
||||
.get(http::header::CONTENT_TYPE)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
body: String::from_utf8(raw_body.to_vec())
|
||||
.expect("refresh body should be utf8"),
|
||||
});
|
||||
Json(json!({
|
||||
"access_token": "refreshed-codex-stream-access-token",
|
||||
"refresh_token": "rt-codex-stream-local-456",
|
||||
"token_type": "Bearer",
|
||||
"expires_in": 3600
|
||||
}))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let execution_runtime = Router::new().route(
|
||||
"/v1/execute/stream",
|
||||
any(move |request: Request| {
|
||||
let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone);
|
||||
async move {
|
||||
let (parts, body) = request.into_parts();
|
||||
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
|
||||
let payload: serde_json::Value =
|
||||
serde_json::from_slice(&raw_body).expect("execution runtime payload should parse");
|
||||
*seen_execution_runtime_inner.lock().expect("mutex should lock") =
|
||||
Some(SeenExecutionRuntimeStreamRequest {
|
||||
trace_id: parts
|
||||
.headers
|
||||
.get(TRACE_ID_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
url: payload
|
||||
.get("url")
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
model: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("model"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
stream: payload
|
||||
.get("stream")
|
||||
.and_then(|value| value.as_bool())
|
||||
.unwrap_or(false),
|
||||
accept: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("accept"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
authorization: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("authorization"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
});
|
||||
let frames = concat!(
|
||||
"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n",
|
||||
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: response.completed\\ndata: {\\\"type\\\":\\\"response.completed\\\"}\\n\\n\"}}\n",
|
||||
"{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":41}}}\n",
|
||||
"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n"
|
||||
);
|
||||
let mut response = Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from(frames))
|
||||
.expect("response should build");
|
||||
response.headers_mut().insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
HeaderValue::from_static("application/x-ndjson"),
|
||||
);
|
||||
response
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let client_api_key = "sk-client-codex-cli-stream-local";
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key(client_api_key)),
|
||||
sample_auth_snapshot(
|
||||
"key-codex-cli-stream-local-123",
|
||||
"user-codex-cli-stream-local-123",
|
||||
),
|
||||
)]));
|
||||
let candidate_selection_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
sample_candidate_row(),
|
||||
]));
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider_catalog_provider()],
|
||||
vec![sample_provider_catalog_endpoint()],
|
||||
vec![sample_provider_catalog_key()],
|
||||
));
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (refresh_url, refresh_handle) = start_server(refresh).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let oauth_refresh =
|
||||
crate::provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(
|
||||
vec![Arc::new(
|
||||
crate::provider_transport::oauth_refresh::GenericOAuthRefreshAdapter::default()
|
||||
.with_token_url_for_tests("codex", format!("{refresh_url}/oauth/token")),
|
||||
)],
|
||||
);
|
||||
let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url.clone())
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
|
||||
auth_repository,
|
||||
candidate_selection_repository,
|
||||
provider_catalog_repository,
|
||||
Arc::clone(&request_candidate_repository),
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
),
|
||||
)
|
||||
.with_oauth_refresh_coordinator_for_tests(oauth_refresh);
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/responses"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(
|
||||
http::header::AUTHORIZATION,
|
||||
format!("Bearer {client_api_key}"),
|
||||
)
|
||||
.header(TRACE_ID_HEADER, "trace-codex-cli-stream-local-123")
|
||||
.body("{\"model\":\"gpt-5.4\",\"input\":\"hello\",\"stream\":true}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response.text().await.expect("body should read"),
|
||||
"event: response.completed\ndata: {\"type\":\"response.completed\"}\n\n"
|
||||
);
|
||||
|
||||
let seen_refresh_request = seen_refresh
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.clone()
|
||||
.expect("refresh request should be captured");
|
||||
assert_eq!(
|
||||
seen_refresh_request.content_type,
|
||||
"application/x-www-form-urlencoded"
|
||||
);
|
||||
assert!(seen_refresh_request
|
||||
.body
|
||||
.contains("grant_type=refresh_token"));
|
||||
assert!(seen_refresh_request
|
||||
.body
|
||||
.contains("client_id=app_EMoamEEZ73f0CkXaXp7hrann"));
|
||||
assert!(seen_refresh_request
|
||||
.body
|
||||
.contains("refresh_token=rt-codex-stream-local-123"));
|
||||
|
||||
let seen_execution_runtime_request = seen_execution_runtime
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.clone()
|
||||
.expect("execution runtime stream should be captured");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.trace_id,
|
||||
"trace-codex-cli-stream-local-123"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.url,
|
||||
"https://chatgpt.com/backend-api/codex/responses"
|
||||
);
|
||||
assert_eq!(seen_execution_runtime_request.model, "gpt-5.4");
|
||||
assert!(seen_execution_runtime_request.stream);
|
||||
assert_eq!(seen_execution_runtime_request.accept, "text/event-stream");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.authorization,
|
||||
"Bearer refreshed-codex-stream-access-token"
|
||||
);
|
||||
|
||||
let stored_candidates = request_candidate_repository
|
||||
.list_by_request_id("trace-codex-cli-stream-local-123")
|
||||
.await
|
||||
.expect("request candidate trace should read");
|
||||
assert_eq!(stored_candidates.len(), 1);
|
||||
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success);
|
||||
|
||||
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
|
||||
assert!(
|
||||
seen_report.lock().expect("mutex should lock").is_none(),
|
||||
"report-stream should stay local when request candidate persistence is available"
|
||||
);
|
||||
|
||||
assert_eq!(*decision_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*plan_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
refresh_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
27
apps/aether-gateway/src/tests/ai_execute/stream_cli/mod.rs
Normal file
27
apps/aether-gateway/src/tests/ai_execute/stream_cli/mod.rs
Normal file
@@ -0,0 +1,27 @@
|
||||
use std::convert::Infallible;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use axum::body::{to_bytes, Body, Bytes};
|
||||
use axum::response::Response;
|
||||
use axum::routing::any;
|
||||
use axum::{extract::Request, Json, Router};
|
||||
use http::header::{HeaderName, HeaderValue};
|
||||
use http::StatusCode;
|
||||
use serde_json::json;
|
||||
|
||||
use crate::constants::{
|
||||
CONTROL_EXECUTED_HEADER, CONTROL_EXECUTE_FALLBACK_HEADER, DEPENDENCY_REASON_HEADER,
|
||||
EXECUTION_PATH_EXECUTION_RUNTIME_STREAM, EXECUTION_PATH_EXECUTION_RUNTIME_SYNC,
|
||||
EXECUTION_PATH_HEADER, EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS,
|
||||
LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER, TRACE_ID_HEADER,
|
||||
};
|
||||
|
||||
use super::{
|
||||
build_router, build_router_with_execution_runtime_override, build_router_with_state,
|
||||
build_state_with_execution_runtime_override, start_server, wait_until, AppState,
|
||||
FrontdoorCorsConfig, FrontdoorUserRpmConfig, GatewayFallbackMetricKind, GatewayFallbackReason,
|
||||
UsageRuntimeConfig, VideoTaskTruthSourceMode,
|
||||
};
|
||||
|
||||
mod compact;
|
||||
mod direct;
|
||||
1954
apps/aether-gateway/src/tests/ai_execute/stream_provider.rs
Normal file
1954
apps/aether-gateway/src/tests/ai_execute/stream_provider.rs
Normal file
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,466 @@
|
||||
use super::{
|
||||
any, build_router_with_state, build_state_with_execution_runtime_override,
|
||||
encrypt_python_fernet_plaintext, json, start_server, to_bytes, Arc, Body, Bytes, Digest,
|
||||
HeaderName, HeaderValue, InMemoryAuthApiKeySnapshotRepository,
|
||||
InMemoryMinimalCandidateSelectionReadRepository, InMemoryProviderCatalogReadRepository,
|
||||
InMemoryRequestCandidateRepository, Json, Mutex, Request, RequestCandidateReadRepository,
|
||||
RequestCandidateStatus, Response, Router, Sha256, StatusCode, StoredAuthApiKeySnapshot,
|
||||
StoredMinimalCandidateSelectionRow, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogProvider, StoredProviderModelMapping, DEVELOPMENT_ENCRYPTION_KEY,
|
||||
TRACE_ID_HEADER,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_executes_gemini_chat_stream_via_local_decision_gate_with_local_stream_decision() {
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenExecutionRuntimeStreamRequest {
|
||||
trace_id: String,
|
||||
url: String,
|
||||
has_model_field: bool,
|
||||
accept: String,
|
||||
auth_header_value: String,
|
||||
exact_temperature: f64,
|
||||
endpoint_tag: String,
|
||||
metadata_mode: String,
|
||||
metadata_source: String,
|
||||
tool_config_present: bool,
|
||||
proxy_node_id: String,
|
||||
tls_profile: String,
|
||||
}
|
||||
|
||||
fn hash_api_key(value: &str) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(value.as_bytes());
|
||||
format!("{:x}", hasher.finalize())
|
||||
}
|
||||
|
||||
fn sample_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
user_id.to_string(),
|
||||
"alice".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
Some(serde_json::json!(["gemini"])),
|
||||
Some(serde_json::json!(["gemini:chat"])),
|
||||
Some(serde_json::json!(["gemini-2.5-pro"])),
|
||||
api_key_id.to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
Some(4_102_444_800),
|
||||
Some(serde_json::json!(["gemini"])),
|
||||
Some(serde_json::json!(["gemini:chat"])),
|
||||
Some(serde_json::json!(["gemini-2.5-pro"])),
|
||||
)
|
||||
.expect("auth snapshot should build")
|
||||
}
|
||||
|
||||
fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow {
|
||||
StoredMinimalCandidateSelectionRow {
|
||||
provider_id: "provider-gemini-chat-local-stream-1".to_string(),
|
||||
provider_name: "gemini".to_string(),
|
||||
provider_type: "custom".to_string(),
|
||||
provider_priority: 10,
|
||||
provider_is_active: true,
|
||||
endpoint_id: "endpoint-gemini-chat-local-stream-1".to_string(),
|
||||
endpoint_api_format: "gemini:chat".to_string(),
|
||||
endpoint_api_family: Some("gemini".to_string()),
|
||||
endpoint_kind: Some("chat".to_string()),
|
||||
endpoint_is_active: true,
|
||||
key_id: "key-gemini-chat-local-stream-1".to_string(),
|
||||
key_name: "prod".to_string(),
|
||||
key_auth_type: "api_key".to_string(),
|
||||
key_is_active: true,
|
||||
key_api_formats: Some(vec!["gemini:chat".to_string()]),
|
||||
key_allowed_models: None,
|
||||
key_capabilities: None,
|
||||
key_internal_priority: 5,
|
||||
key_global_priority_by_format: Some(serde_json::json!({"gemini:chat": 1})),
|
||||
model_id: "model-gemini-chat-local-stream-1".to_string(),
|
||||
global_model_id: "global-model-gemini-chat-local-stream-1".to_string(),
|
||||
global_model_name: "gemini-2.5-pro".to_string(),
|
||||
global_model_mappings: None,
|
||||
global_model_supports_streaming: Some(true),
|
||||
model_provider_model_name: "gemini-2.5-pro-upstream".to_string(),
|
||||
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
|
||||
name: "gemini-2.5-pro-upstream".to_string(),
|
||||
priority: 1,
|
||||
api_formats: Some(vec!["gemini:chat".to_string()]),
|
||||
}]),
|
||||
model_supports_streaming: Some(true),
|
||||
model_is_active: true,
|
||||
model_is_available: true,
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
|
||||
StoredProviderCatalogProvider::new(
|
||||
"provider-gemini-chat-local-stream-1".to_string(),
|
||||
"gemini".to_string(),
|
||||
Some("https://example.com".to_string()),
|
||||
"custom".to_string(),
|
||||
)
|
||||
.expect("provider should build")
|
||||
.with_transport_fields(
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
None,
|
||||
Some(2),
|
||||
Some(serde_json::json!({"url":"http://provider-proxy.internal:8080"})),
|
||||
Some(20.0),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
|
||||
StoredProviderCatalogEndpoint::new(
|
||||
"endpoint-gemini-chat-local-stream-1".to_string(),
|
||||
"provider-gemini-chat-local-stream-1".to_string(),
|
||||
"gemini:chat".to_string(),
|
||||
Some("gemini".to_string()),
|
||||
Some("chat".to_string()),
|
||||
true,
|
||||
)
|
||||
.expect("endpoint should build")
|
||||
.with_transport_fields(
|
||||
"https://generativelanguage.googleapis.com".to_string(),
|
||||
Some(serde_json::json!([
|
||||
{"action":"set","key":"x-endpoint-tag","value":"gemini-chat-local-stream"}
|
||||
])),
|
||||
Some(serde_json::json!([
|
||||
{"action":"set","path":"metadata.mode","value":"safe"},
|
||||
{"action":"rename","from":"metadata.client","to":"metadata.source"},
|
||||
{"action":"drop","path":"toolConfig"}
|
||||
])),
|
||||
Some(2),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("endpoint transport should build")
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
|
||||
StoredProviderCatalogKey::new(
|
||||
"key-gemini-chat-local-stream-1".to_string(),
|
||||
"provider-gemini-chat-local-stream-1".to_string(),
|
||||
"prod".to_string(),
|
||||
"api_key".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
.with_transport_fields(
|
||||
Some(serde_json::json!(["gemini:chat"])),
|
||||
encrypt_python_fernet_plaintext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
"sk-upstream-gemini-chat-stream",
|
||||
)
|
||||
.expect("api key should encrypt"),
|
||||
None,
|
||||
None,
|
||||
Some(serde_json::json!({"gemini:chat": 1})),
|
||||
None,
|
||||
None,
|
||||
Some(serde_json::json!({"enabled": true, "node_id":"proxy-node-gemini-chat-stream"})),
|
||||
Some(serde_json::json!({"tls_profile":"chrome_136"})),
|
||||
)
|
||||
.expect("key transport should build")
|
||||
}
|
||||
|
||||
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeStreamRequest>));
|
||||
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
|
||||
let seen_report = Arc::new(Mutex::new(false));
|
||||
let seen_report_clone = Arc::clone(&seen_report);
|
||||
let decision_hits = Arc::new(Mutex::new(0usize));
|
||||
let decision_hits_clone = Arc::clone(&decision_hits);
|
||||
let plan_hits = Arc::new(Mutex::new(0usize));
|
||||
let plan_hits_clone = Arc::clone(&plan_hits);
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/resolve",
|
||||
any(|_request: Request| async move {
|
||||
Json(json!({
|
||||
"action": "proxy_public",
|
||||
"route_class": "ai_public",
|
||||
"route_family": "gemini",
|
||||
"route_kind": "chat",
|
||||
"auth_endpoint_signature": "gemini:chat",
|
||||
"execution_runtime_candidate": true,
|
||||
"auth_context": {
|
||||
"user_id": "user-gemini-chat-local-stream-1",
|
||||
"api_key_id": "api-key-gemini-chat-local-stream-1",
|
||||
"access_allowed": true
|
||||
},
|
||||
"public_path": "/v1beta/models/gemini-2.5-pro:streamGenerateContent"
|
||||
}))
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/decision-stream",
|
||||
any(move |_request: Request| {
|
||||
let decision_hits_inner = Arc::clone(&decision_hits_clone);
|
||||
async move {
|
||||
*decision_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({"action": "proxy_public"}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/plan-stream",
|
||||
any(move |_request: Request| {
|
||||
let plan_hits_inner = Arc::clone(&plan_hits_clone);
|
||||
async move {
|
||||
*plan_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({"action": "proxy_public"}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/report-stream",
|
||||
any(move |request: Request| {
|
||||
let seen_report_inner = Arc::clone(&seen_report_clone);
|
||||
async move {
|
||||
let (_parts, body) = request.into_parts();
|
||||
let _raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
|
||||
*seen_report_inner.lock().expect("mutex should lock") = true;
|
||||
Json(json!({"ok": true}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1beta/models/gemini-2.5-pro:streamGenerateContent",
|
||||
any(move |_request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::IM_A_TEAPOT, Body::from("public-route-hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let execution_runtime = Router::new().route(
|
||||
"/v1/execute/stream",
|
||||
any(move |request: Request| {
|
||||
let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone);
|
||||
async move {
|
||||
let (parts, body) = request.into_parts();
|
||||
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
|
||||
let payload: serde_json::Value =
|
||||
serde_json::from_slice(&raw_body).expect("execution runtime payload should parse");
|
||||
*seen_execution_runtime_inner.lock().expect("mutex should lock") =
|
||||
Some(SeenExecutionRuntimeStreamRequest {
|
||||
trace_id: parts
|
||||
.headers
|
||||
.get(TRACE_ID_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
url: payload
|
||||
.get("url")
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
has_model_field: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("model"))
|
||||
.is_some(),
|
||||
accept: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("accept"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
auth_header_value: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-goog-api-key"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
exact_temperature: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("generationConfig"))
|
||||
.and_then(|value| value.get("temperature"))
|
||||
.and_then(|value| value.as_f64())
|
||||
.unwrap_or_default(),
|
||||
endpoint_tag: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-endpoint-tag"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
metadata_mode: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("metadata"))
|
||||
.and_then(|value| value.get("mode"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
metadata_source: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("metadata"))
|
||||
.and_then(|value| value.get("source"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
tool_config_present: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("toolConfig"))
|
||||
.is_some(),
|
||||
proxy_node_id: payload
|
||||
.get("proxy")
|
||||
.and_then(|value| value.get("node_id"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
tls_profile: payload
|
||||
.get("tls_profile")
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
});
|
||||
let frames = concat!(
|
||||
"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n",
|
||||
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"candidates\\\":[]}\\n\\n\"}}\n",
|
||||
"{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":33,\"upstream_bytes\":26}}}\n",
|
||||
"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n"
|
||||
);
|
||||
let mut response = Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from(frames))
|
||||
.expect("response should build");
|
||||
response.headers_mut().insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
HeaderValue::from_static("application/x-ndjson"),
|
||||
);
|
||||
response
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("client-gemini-chat-local-stream-key")),
|
||||
sample_auth_snapshot(
|
||||
"api-key-gemini-chat-local-stream-1",
|
||||
"user-gemini-chat-local-stream-1",
|
||||
),
|
||||
)]));
|
||||
let candidate_selection_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
sample_candidate_row(),
|
||||
]));
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider_catalog_provider()],
|
||||
vec![sample_provider_catalog_endpoint()],
|
||||
vec![sample_provider_catalog_key()],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url.clone())
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
|
||||
auth_repository,
|
||||
candidate_selection_repository,
|
||||
provider_catalog_repository,
|
||||
Arc::clone(&request_candidate_repository),
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
),
|
||||
);
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!(
|
||||
"{gateway_url}/v1beta/models/gemini-2.5-pro:streamGenerateContent?key=client-gemini-chat-local-stream-key"
|
||||
))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(TRACE_ID_HEADER, "trace-gemini-chat-local-stream-1")
|
||||
.body(
|
||||
"{\"contents\":[],\"generationConfig\":{\"temperature\":0.2},\"metadata\":{\"client\":\"desktop-gemini-stream\"},\"toolConfig\":{\"functionCallingConfig\":{\"mode\":\"AUTO\"}}}",
|
||||
)
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response.text().await.expect("body should read"),
|
||||
"data: {\"candidates\":[]}\n\n"
|
||||
);
|
||||
|
||||
let seen_execution_runtime_request = seen_execution_runtime
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.clone()
|
||||
.expect("execution runtime stream should be captured");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.trace_id,
|
||||
"trace-gemini-chat-local-stream-1"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.url,
|
||||
"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-pro-upstream:streamGenerateContent?alt=sse"
|
||||
);
|
||||
assert!(!seen_execution_runtime_request.has_model_field);
|
||||
assert_eq!(seen_execution_runtime_request.accept, "text/event-stream");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.auth_header_value,
|
||||
"sk-upstream-gemini-chat-stream"
|
||||
);
|
||||
assert!((seen_execution_runtime_request.exact_temperature - 0.2).abs() < f64::EPSILON);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.endpoint_tag,
|
||||
"gemini-chat-local-stream"
|
||||
);
|
||||
assert_eq!(seen_execution_runtime_request.metadata_mode, "safe");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.metadata_source,
|
||||
"desktop-gemini-stream"
|
||||
);
|
||||
assert!(!seen_execution_runtime_request.tool_config_present);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.proxy_node_id,
|
||||
"proxy-node-gemini-chat-stream"
|
||||
);
|
||||
assert_eq!(seen_execution_runtime_request.tls_profile, "chrome_136");
|
||||
|
||||
let stored_candidates = request_candidate_repository
|
||||
.list_by_request_id("trace-gemini-chat-local-stream-1")
|
||||
.await
|
||||
.expect("request candidate trace should read");
|
||||
assert_eq!(stored_candidates.len(), 1);
|
||||
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success);
|
||||
|
||||
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
|
||||
assert!(
|
||||
!*seen_report.lock().expect("mutex should lock"),
|
||||
"report-stream should stay local when request candidate persistence is available"
|
||||
);
|
||||
|
||||
assert_eq!(*decision_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*plan_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,43 @@
|
||||
use std::convert::Infallible;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use axum::body::{to_bytes, Body, Bytes};
|
||||
use axum::response::Response;
|
||||
use axum::routing::any;
|
||||
use axum::{extract::Request, Json, Router};
|
||||
use http::header::{HeaderName, HeaderValue};
|
||||
use http::StatusCode;
|
||||
use serde_json::json;
|
||||
|
||||
use crate::constants::{
|
||||
CONTROL_EXECUTED_HEADER, CONTROL_EXECUTE_FALLBACK_HEADER, DEPENDENCY_REASON_HEADER,
|
||||
EXECUTION_PATH_EXECUTION_RUNTIME_STREAM, EXECUTION_PATH_EXECUTION_RUNTIME_SYNC,
|
||||
EXECUTION_PATH_HEADER, EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS,
|
||||
LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER, TRACE_ID_HEADER,
|
||||
};
|
||||
|
||||
use super::{
|
||||
build_router, build_router_with_execution_runtime_override, build_router_with_state,
|
||||
build_state_with_execution_runtime_override, start_server, wait_until, AppState,
|
||||
FrontdoorCorsConfig, FrontdoorUserRpmConfig, GatewayFallbackMetricKind, GatewayFallbackReason,
|
||||
UsageRuntimeConfig, VideoTaskTruthSourceMode,
|
||||
};
|
||||
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
|
||||
use aether_data::repository::auth::{
|
||||
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot,
|
||||
};
|
||||
use aether_data::repository::candidate_selection::{
|
||||
InMemoryMinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
|
||||
StoredProviderModelMapping,
|
||||
};
|
||||
use aether_data::repository::candidates::{
|
||||
InMemoryRequestCandidateRepository, RequestCandidateReadRepository, RequestCandidateStatus,
|
||||
};
|
||||
use aether_data::repository::provider_catalog::{
|
||||
InMemoryProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogProvider,
|
||||
};
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
mod local_chat;
|
||||
mod local_cli;
|
||||
1152
apps/aether-gateway/src/tests/ai_execute/sync/chat/failover.rs
Normal file
1152
apps/aether-gateway/src/tests/ai_execute/sync/chat/failover.rs
Normal file
File diff suppressed because it is too large
Load Diff
1744
apps/aether-gateway/src/tests/ai_execute/sync/chat/local_decision.rs
Normal file
1744
apps/aether-gateway/src/tests/ai_execute/sync/chat/local_decision.rs
Normal file
File diff suppressed because it is too large
Load Diff
43
apps/aether-gateway/src/tests/ai_execute/sync/chat/mod.rs
Normal file
43
apps/aether-gateway/src/tests/ai_execute/sync/chat/mod.rs
Normal file
@@ -0,0 +1,43 @@
|
||||
use std::convert::Infallible;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use axum::body::{to_bytes, Body, Bytes};
|
||||
use axum::response::Response;
|
||||
use axum::routing::any;
|
||||
use axum::{extract::Request, Json, Router};
|
||||
use http::header::{HeaderName, HeaderValue};
|
||||
use http::StatusCode;
|
||||
use serde_json::json;
|
||||
|
||||
use crate::constants::{
|
||||
CONTROL_EXECUTED_HEADER, CONTROL_EXECUTE_FALLBACK_HEADER, DEPENDENCY_REASON_HEADER,
|
||||
EXECUTION_PATH_EXECUTION_RUNTIME_STREAM, EXECUTION_PATH_EXECUTION_RUNTIME_SYNC,
|
||||
EXECUTION_PATH_HEADER, EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS,
|
||||
LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER, TRACE_ID_HEADER,
|
||||
};
|
||||
|
||||
use super::{
|
||||
build_router, build_router_with_execution_runtime_override, build_router_with_state,
|
||||
build_state_with_execution_runtime_override, start_server, wait_until, AppState,
|
||||
FrontdoorCorsConfig, FrontdoorUserRpmConfig, GatewayFallbackMetricKind, GatewayFallbackReason,
|
||||
UsageRuntimeConfig, VideoTaskTruthSourceMode,
|
||||
};
|
||||
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
|
||||
use aether_data::repository::auth::{
|
||||
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot,
|
||||
};
|
||||
use aether_data::repository::candidate_selection::{
|
||||
InMemoryMinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
|
||||
StoredProviderModelMapping,
|
||||
};
|
||||
use aether_data::repository::candidates::{
|
||||
InMemoryRequestCandidateRepository, RequestCandidateReadRepository, RequestCandidateStatus,
|
||||
};
|
||||
use aether_data::repository::provider_catalog::{
|
||||
InMemoryProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogProvider,
|
||||
};
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
mod failover;
|
||||
mod local_decision;
|
||||
@@ -0,0 +1,554 @@
|
||||
use super::{
|
||||
any, build_router_with_state, build_state_with_execution_runtime_override,
|
||||
encrypt_python_fernet_plaintext, json, start_server, to_bytes, Arc, Body, Digest,
|
||||
InMemoryAuthApiKeySnapshotRepository, InMemoryMinimalCandidateSelectionReadRepository,
|
||||
InMemoryProviderCatalogReadRepository, InMemoryRequestCandidateRepository, Json, Mutex,
|
||||
Request, RequestCandidateReadRepository, RequestCandidateStatus, Router, Sha256, StatusCode,
|
||||
StoredAuthApiKeySnapshot, StoredMinimalCandidateSelectionRow, StoredProviderCatalogEndpoint,
|
||||
StoredProviderCatalogKey, StoredProviderCatalogProvider, StoredProviderModelMapping,
|
||||
DEVELOPMENT_ENCRYPTION_KEY, TRACE_ID_HEADER,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_executes_claude_code_cli_sync_via_local_decision_gate_with_local_sync_decision() {
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenExecutionRuntimeSyncRequest {
|
||||
trace_id: String,
|
||||
url: String,
|
||||
model: String,
|
||||
authorization: String,
|
||||
accept: String,
|
||||
anthropic_version: String,
|
||||
anthropic_beta: String,
|
||||
x_app: String,
|
||||
x_stainless_helper_method: String,
|
||||
x_stainless_package_version: String,
|
||||
user_agent: String,
|
||||
endpoint_tag: String,
|
||||
metadata_mode: String,
|
||||
metadata_source: String,
|
||||
assistant_content: serde_json::Value,
|
||||
proxy_node_id: String,
|
||||
tls_profile: String,
|
||||
}
|
||||
|
||||
fn hash_api_key(value: &str) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(value.as_bytes());
|
||||
format!("{:x}", hasher.finalize())
|
||||
}
|
||||
|
||||
fn sample_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
user_id.to_string(),
|
||||
"alice".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
Some(serde_json::json!(["claude", "claude_code"])),
|
||||
Some(serde_json::json!(["claude:cli"])),
|
||||
Some(serde_json::json!(["claude-code"])),
|
||||
api_key_id.to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
Some(4_102_444_800),
|
||||
Some(serde_json::json!(["claude", "claude_code"])),
|
||||
Some(serde_json::json!(["claude:cli"])),
|
||||
Some(serde_json::json!(["claude-code"])),
|
||||
)
|
||||
.expect("auth snapshot should build")
|
||||
}
|
||||
|
||||
fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow {
|
||||
StoredMinimalCandidateSelectionRow {
|
||||
provider_id: "provider-claude-code-cli-local-1".to_string(),
|
||||
provider_name: "claude_code".to_string(),
|
||||
provider_type: "claude_code".to_string(),
|
||||
provider_priority: 10,
|
||||
provider_is_active: true,
|
||||
endpoint_id: "endpoint-claude-code-cli-local-1".to_string(),
|
||||
endpoint_api_format: "claude:cli".to_string(),
|
||||
endpoint_api_family: Some("claude".to_string()),
|
||||
endpoint_kind: Some("cli".to_string()),
|
||||
endpoint_is_active: true,
|
||||
key_id: "key-claude-code-cli-local-1".to_string(),
|
||||
key_name: "prod".to_string(),
|
||||
key_auth_type: "oauth".to_string(),
|
||||
key_is_active: true,
|
||||
key_api_formats: Some(vec!["claude:cli".to_string()]),
|
||||
key_allowed_models: None,
|
||||
key_capabilities: None,
|
||||
key_internal_priority: 5,
|
||||
key_global_priority_by_format: Some(serde_json::json!({"claude:cli": 1})),
|
||||
model_id: "model-claude-code-cli-local-1".to_string(),
|
||||
global_model_id: "global-model-claude-code-cli-local-1".to_string(),
|
||||
global_model_name: "claude-code".to_string(),
|
||||
global_model_mappings: None,
|
||||
global_model_supports_streaming: Some(true),
|
||||
model_provider_model_name: "claude-code-upstream".to_string(),
|
||||
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
|
||||
name: "claude-code-upstream".to_string(),
|
||||
priority: 1,
|
||||
api_formats: Some(vec!["claude:cli".to_string()]),
|
||||
}]),
|
||||
model_supports_streaming: Some(true),
|
||||
model_is_active: true,
|
||||
model_is_available: true,
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
|
||||
StoredProviderCatalogProvider::new(
|
||||
"provider-claude-code-cli-local-1".to_string(),
|
||||
"claude_code".to_string(),
|
||||
Some("https://example.com".to_string()),
|
||||
"claude_code".to_string(),
|
||||
)
|
||||
.expect("provider should build")
|
||||
.with_transport_fields(
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
None,
|
||||
Some(2),
|
||||
Some(serde_json::json!({"url":"http://provider-proxy.internal:8080"})),
|
||||
Some(20.0),
|
||||
None,
|
||||
Some(serde_json::json!({
|
||||
"claude_code_advanced": {
|
||||
"cli_only_enabled": false
|
||||
}
|
||||
})),
|
||||
)
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
|
||||
StoredProviderCatalogEndpoint::new(
|
||||
"endpoint-claude-code-cli-local-1".to_string(),
|
||||
"provider-claude-code-cli-local-1".to_string(),
|
||||
"claude:cli".to_string(),
|
||||
Some("claude".to_string()),
|
||||
Some("cli".to_string()),
|
||||
true,
|
||||
)
|
||||
.expect("endpoint should build")
|
||||
.with_transport_fields(
|
||||
"https://api.anthropic.example/v1/messages".to_string(),
|
||||
Some(serde_json::json!([
|
||||
{"action":"set","key":"x-endpoint-tag","value":"claude-code-cli-local"}
|
||||
])),
|
||||
Some(serde_json::json!([
|
||||
{"action":"set","path":"metadata.mode","value":"safe"},
|
||||
{"action":"rename","from":"metadata.client","to":"metadata.source"}
|
||||
])),
|
||||
Some(2),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("endpoint transport should build")
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
|
||||
StoredProviderCatalogKey::new(
|
||||
"key-claude-code-cli-local-1".to_string(),
|
||||
"provider-claude-code-cli-local-1".to_string(),
|
||||
"prod".to_string(),
|
||||
"oauth".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
.with_transport_fields(
|
||||
Some(serde_json::json!(["claude:cli"])),
|
||||
encrypt_python_fernet_plaintext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
"sk-upstream-claude-code-oauth",
|
||||
)
|
||||
.expect("api key should encrypt"),
|
||||
None,
|
||||
None,
|
||||
Some(serde_json::json!({"claude:cli": 1})),
|
||||
None,
|
||||
None,
|
||||
Some(
|
||||
serde_json::json!({"enabled": true, "node_id":"proxy-node-claude-code-cli-local"}),
|
||||
),
|
||||
Some(serde_json::json!({
|
||||
"tls_profile":"claude_code_nodejs",
|
||||
"user_agent":"Claude-Code/9.9",
|
||||
"stainless_package_version":"1.0.5",
|
||||
"stainless_runtime_version":"v22.12.0",
|
||||
"stainless_timeout":"900"
|
||||
})),
|
||||
)
|
||||
.expect("key transport should build")
|
||||
}
|
||||
|
||||
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeSyncRequest>));
|
||||
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
|
||||
let seen_report = Arc::new(Mutex::new(false));
|
||||
let seen_report_clone = Arc::clone(&seen_report);
|
||||
let decision_hits = Arc::new(Mutex::new(0usize));
|
||||
let decision_hits_clone = Arc::clone(&decision_hits);
|
||||
let plan_hits = Arc::new(Mutex::new(0usize));
|
||||
let plan_hits_clone = Arc::clone(&plan_hits);
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/decision-sync",
|
||||
any(move |_request: Request| {
|
||||
let decision_hits_inner = Arc::clone(&decision_hits_clone);
|
||||
async move {
|
||||
*decision_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({"action": "proxy_public"}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/plan-sync",
|
||||
any(move |_request: Request| {
|
||||
let plan_hits_inner = Arc::clone(&plan_hits_clone);
|
||||
async move {
|
||||
*plan_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({"action": "proxy_public"}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/report-sync",
|
||||
any(move |request: Request| {
|
||||
let seen_report_inner = Arc::clone(&seen_report_clone);
|
||||
async move {
|
||||
let (_parts, body) = request.into_parts();
|
||||
let _raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
|
||||
*seen_report_inner.lock().expect("mutex should lock") = true;
|
||||
Json(json!({"ok": true}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1/messages",
|
||||
any(move |_request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::IM_A_TEAPOT, Body::from("public-route-hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let execution_runtime = Router::new().route(
|
||||
"/v1/execute/sync",
|
||||
any(move |request: Request| {
|
||||
let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone);
|
||||
async move {
|
||||
let (parts, body) = request.into_parts();
|
||||
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
|
||||
let payload: serde_json::Value = serde_json::from_slice(&raw_body)
|
||||
.expect("execution runtime payload should parse");
|
||||
*seen_execution_runtime_inner
|
||||
.lock()
|
||||
.expect("mutex should lock") = Some(SeenExecutionRuntimeSyncRequest {
|
||||
trace_id: parts
|
||||
.headers
|
||||
.get(TRACE_ID_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
url: payload
|
||||
.get("url")
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
model: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("model"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
authorization: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("authorization"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
accept: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("accept"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
anthropic_version: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("anthropic-version"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
anthropic_beta: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("anthropic-beta"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
x_app: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-app"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
x_stainless_helper_method: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-stainless-helper-method"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
x_stainless_package_version: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-stainless-package-version"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
user_agent: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("user-agent"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
endpoint_tag: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-endpoint-tag"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
metadata_mode: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("metadata"))
|
||||
.and_then(|value| value.get("mode"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
metadata_source: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("metadata"))
|
||||
.and_then(|value| value.get("source"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
assistant_content: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("messages"))
|
||||
.and_then(|value| value.get(0))
|
||||
.and_then(|value| value.get("content"))
|
||||
.cloned()
|
||||
.unwrap_or(serde_json::Value::Null),
|
||||
proxy_node_id: payload
|
||||
.get("proxy")
|
||||
.and_then(|value| value.get("node_id"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
tls_profile: payload
|
||||
.get("tls_profile")
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
});
|
||||
Json(json!({
|
||||
"request_id": "trace-claude-code-cli-local-sync-123",
|
||||
"status_code": 200,
|
||||
"headers": {
|
||||
"content-type": "application/json"
|
||||
},
|
||||
"body": {
|
||||
"json_body": {
|
||||
"id": "msg-local-claude-code-cli-123",
|
||||
"type": "message",
|
||||
"model": "claude-code-upstream",
|
||||
"role": "assistant",
|
||||
"content": [],
|
||||
"usage": {
|
||||
"input_tokens": 2,
|
||||
"output_tokens": 3
|
||||
}
|
||||
}
|
||||
},
|
||||
"telemetry": {
|
||||
"elapsed_ms": 29
|
||||
}
|
||||
}))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("sk-client-claude-code-cli-local")),
|
||||
sample_auth_snapshot(
|
||||
"api-key-claude-code-cli-local-1",
|
||||
"user-claude-code-cli-local-1",
|
||||
),
|
||||
)]));
|
||||
let candidate_selection_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
sample_candidate_row(),
|
||||
]));
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider_catalog_provider()],
|
||||
vec![sample_provider_catalog_endpoint()],
|
||||
vec![sample_provider_catalog_key()],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url.clone())
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
|
||||
auth_repository,
|
||||
candidate_selection_repository,
|
||||
provider_catalog_repository,
|
||||
Arc::clone(&request_candidate_repository),
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
),
|
||||
);
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/messages"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(
|
||||
http::header::AUTHORIZATION,
|
||||
"Bearer sk-client-claude-code-cli-local",
|
||||
)
|
||||
.header("anthropic-beta", "context-1m-2025-08-07,custom-beta")
|
||||
.header(TRACE_ID_HEADER, "trace-claude-code-cli-local-sync-123")
|
||||
.body(
|
||||
serde_json::json!({
|
||||
"model": "claude-code",
|
||||
"thinking": {"type":"enabled"},
|
||||
"messages": [{
|
||||
"role":"assistant",
|
||||
"content":[
|
||||
{"type":"thinking","thinking":"keep","signature":"sig_valid"},
|
||||
{"type":"thinking","thinking":"drop-empty-signature","signature":""},
|
||||
{"type":"redacted_thinking","data":"keep-redacted","signature":"sig_redacted"},
|
||||
{"type":"redacted_thinking","data":"drop-no-signature"},
|
||||
{"type":"text","text":"ok"}
|
||||
]
|
||||
}],
|
||||
"metadata":{"client":"desktop-claude-code-cli"}
|
||||
})
|
||||
.to_string(),
|
||||
)
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
let status = response.status();
|
||||
let response_body = response.text().await.expect("body should read");
|
||||
assert!(
|
||||
status == StatusCode::OK,
|
||||
"unexpected status={status} body={response_body} decision_hits={} plan_hits={} public_hits={}",
|
||||
*decision_hits.lock().expect("mutex should lock"),
|
||||
*plan_hits.lock().expect("mutex should lock"),
|
||||
*public_hits.lock().expect("mutex should lock"),
|
||||
);
|
||||
let response_json: serde_json::Value =
|
||||
serde_json::from_str(&response_body).expect("body should parse");
|
||||
assert_eq!(response_json["model"], "claude-code-upstream");
|
||||
|
||||
let seen_execution_runtime_request = seen_execution_runtime
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.clone()
|
||||
.expect("execution runtime sync should be captured");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.trace_id,
|
||||
"trace-claude-code-cli-local-sync-123"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.url,
|
||||
"https://api.anthropic.example/v1/messages"
|
||||
);
|
||||
assert_eq!(seen_execution_runtime_request.model, "claude-code-upstream");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.authorization,
|
||||
"Bearer sk-upstream-claude-code-oauth"
|
||||
);
|
||||
assert_eq!(seen_execution_runtime_request.accept, "application/json");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.anthropic_version,
|
||||
"2023-06-01"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.anthropic_beta,
|
||||
"claude-code-20250219,oauth-2025-04-20,interleaved-thinking-2025-05-14,custom-beta"
|
||||
);
|
||||
assert_eq!(seen_execution_runtime_request.x_app, "cli");
|
||||
assert_eq!(seen_execution_runtime_request.x_stainless_helper_method, "");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.x_stainless_package_version,
|
||||
"1.0.5"
|
||||
);
|
||||
assert_eq!(seen_execution_runtime_request.user_agent, "Claude-Code/9.9");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.endpoint_tag,
|
||||
"claude-code-cli-local"
|
||||
);
|
||||
assert_eq!(seen_execution_runtime_request.metadata_mode, "safe");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.metadata_source,
|
||||
"desktop-claude-code-cli"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.assistant_content,
|
||||
json!([
|
||||
{"type":"thinking","thinking":"keep","signature":"sig_valid"},
|
||||
{"type":"redacted_thinking","data":"keep-redacted","signature":"sig_redacted"},
|
||||
{"type":"text","text":"ok"}
|
||||
])
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.proxy_node_id,
|
||||
"proxy-node-claude-code-cli-local"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.tls_profile,
|
||||
"claude_code_nodejs"
|
||||
);
|
||||
|
||||
let stored_candidates = request_candidate_repository
|
||||
.list_by_request_id("trace-claude-code-cli-local-sync-123")
|
||||
.await
|
||||
.expect("request candidate trace should read");
|
||||
assert_eq!(stored_candidates.len(), 1);
|
||||
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success);
|
||||
|
||||
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
|
||||
assert!(
|
||||
!*seen_report.lock().expect("mutex should lock"),
|
||||
"report-sync should stay local when request candidate persistence is available"
|
||||
);
|
||||
|
||||
assert_eq!(*decision_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*plan_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
1032
apps/aether-gateway/src/tests/ai_execute/sync/claude/kiro.rs
Normal file
1032
apps/aether-gateway/src/tests/ai_execute/sync/claude/kiro.rs
Normal file
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,427 @@
|
||||
use super::{
|
||||
any, build_router_with_state, build_state_with_execution_runtime_override,
|
||||
encrypt_python_fernet_plaintext, json, start_server, to_bytes, Arc, Body, Digest,
|
||||
InMemoryAuthApiKeySnapshotRepository, InMemoryMinimalCandidateSelectionReadRepository,
|
||||
InMemoryProviderCatalogReadRepository, InMemoryRequestCandidateRepository, Json, Mutex,
|
||||
Request, RequestCandidateReadRepository, RequestCandidateStatus, Router, Sha256, StatusCode,
|
||||
StoredAuthApiKeySnapshot, StoredMinimalCandidateSelectionRow, StoredProviderCatalogEndpoint,
|
||||
StoredProviderCatalogKey, StoredProviderCatalogProvider, StoredProviderModelMapping,
|
||||
DEVELOPMENT_ENCRYPTION_KEY, TRACE_ID_HEADER,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_executes_claude_chat_sync_via_local_decision_gate_with_local_sync_decision() {
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenExecutionRuntimeSyncRequest {
|
||||
trace_id: String,
|
||||
url: String,
|
||||
model: String,
|
||||
auth_header_value: String,
|
||||
endpoint_tag: String,
|
||||
metadata_mode: String,
|
||||
metadata_source: String,
|
||||
proxy_node_id: String,
|
||||
tls_profile: String,
|
||||
}
|
||||
|
||||
fn hash_api_key(value: &str) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(value.as_bytes());
|
||||
format!("{:x}", hasher.finalize())
|
||||
}
|
||||
|
||||
fn sample_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
user_id.to_string(),
|
||||
"alice".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
Some(serde_json::json!(["claude"])),
|
||||
Some(serde_json::json!(["claude:chat"])),
|
||||
Some(serde_json::json!(["claude-sonnet-4-5"])),
|
||||
api_key_id.to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
Some(4_102_444_800),
|
||||
Some(serde_json::json!(["claude"])),
|
||||
Some(serde_json::json!(["claude:chat"])),
|
||||
Some(serde_json::json!(["claude-sonnet-4-5"])),
|
||||
)
|
||||
.expect("auth snapshot should build")
|
||||
}
|
||||
|
||||
fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow {
|
||||
StoredMinimalCandidateSelectionRow {
|
||||
provider_id: "provider-claude-local-1".to_string(),
|
||||
provider_name: "claude".to_string(),
|
||||
provider_type: "custom".to_string(),
|
||||
provider_priority: 10,
|
||||
provider_is_active: true,
|
||||
endpoint_id: "endpoint-claude-local-1".to_string(),
|
||||
endpoint_api_format: "claude:chat".to_string(),
|
||||
endpoint_api_family: Some("claude".to_string()),
|
||||
endpoint_kind: Some("chat".to_string()),
|
||||
endpoint_is_active: true,
|
||||
key_id: "key-claude-local-1".to_string(),
|
||||
key_name: "prod".to_string(),
|
||||
key_auth_type: "api_key".to_string(),
|
||||
key_is_active: true,
|
||||
key_api_formats: Some(vec!["claude:chat".to_string()]),
|
||||
key_allowed_models: None,
|
||||
key_capabilities: None,
|
||||
key_internal_priority: 5,
|
||||
key_global_priority_by_format: Some(serde_json::json!({"claude:chat": 1})),
|
||||
model_id: "model-claude-local-1".to_string(),
|
||||
global_model_id: "global-model-claude-local-1".to_string(),
|
||||
global_model_name: "claude-sonnet-4-5".to_string(),
|
||||
global_model_mappings: None,
|
||||
global_model_supports_streaming: Some(true),
|
||||
model_provider_model_name: "claude-sonnet-4-5-upstream".to_string(),
|
||||
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
|
||||
name: "claude-sonnet-4-5-upstream".to_string(),
|
||||
priority: 1,
|
||||
api_formats: Some(vec!["claude:chat".to_string()]),
|
||||
}]),
|
||||
model_supports_streaming: Some(true),
|
||||
model_is_active: true,
|
||||
model_is_available: true,
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
|
||||
StoredProviderCatalogProvider::new(
|
||||
"provider-claude-local-1".to_string(),
|
||||
"claude".to_string(),
|
||||
Some("https://example.com".to_string()),
|
||||
"custom".to_string(),
|
||||
)
|
||||
.expect("provider should build")
|
||||
.with_transport_fields(
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
None,
|
||||
Some(2),
|
||||
Some(serde_json::json!({"url":"http://provider-proxy.internal:8080"})),
|
||||
Some(20.0),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
|
||||
StoredProviderCatalogEndpoint::new(
|
||||
"endpoint-claude-local-1".to_string(),
|
||||
"provider-claude-local-1".to_string(),
|
||||
"claude:chat".to_string(),
|
||||
Some("claude".to_string()),
|
||||
Some("chat".to_string()),
|
||||
true,
|
||||
)
|
||||
.expect("endpoint should build")
|
||||
.with_transport_fields(
|
||||
"https://api.anthropic.example".to_string(),
|
||||
Some(serde_json::json!([
|
||||
{"action":"set","key":"x-endpoint-tag","value":"claude-chat-local"}
|
||||
])),
|
||||
Some(serde_json::json!([
|
||||
{"action":"set","path":"metadata.mode","value":"safe"},
|
||||
{"action":"rename","from":"metadata.client","to":"metadata.source"}
|
||||
])),
|
||||
Some(2),
|
||||
Some("/custom/v1/messages".to_string()),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("endpoint transport should build")
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
|
||||
StoredProviderCatalogKey::new(
|
||||
"key-claude-local-1".to_string(),
|
||||
"provider-claude-local-1".to_string(),
|
||||
"prod".to_string(),
|
||||
"api_key".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
.with_transport_fields(
|
||||
Some(serde_json::json!(["claude:chat"])),
|
||||
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "sk-upstream-claude-chat")
|
||||
.expect("api key should encrypt"),
|
||||
None,
|
||||
None,
|
||||
Some(serde_json::json!({"claude:chat": 1})),
|
||||
None,
|
||||
None,
|
||||
Some(serde_json::json!({"enabled": true, "node_id":"proxy-node-claude-chat-local"})),
|
||||
Some(serde_json::json!({"tls_profile":"chrome_136"})),
|
||||
)
|
||||
.expect("key transport should build")
|
||||
}
|
||||
|
||||
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeSyncRequest>));
|
||||
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
|
||||
let seen_report = Arc::new(Mutex::new(false));
|
||||
let seen_report_clone = Arc::clone(&seen_report);
|
||||
let decision_hits = Arc::new(Mutex::new(0usize));
|
||||
let decision_hits_clone = Arc::clone(&decision_hits);
|
||||
let plan_hits = Arc::new(Mutex::new(0usize));
|
||||
let plan_hits_clone = Arc::clone(&plan_hits);
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/decision-sync",
|
||||
any(move |_request: Request| {
|
||||
let decision_hits_inner = Arc::clone(&decision_hits_clone);
|
||||
async move {
|
||||
*decision_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({"action": "proxy_public"}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/plan-sync",
|
||||
any(move |_request: Request| {
|
||||
let plan_hits_inner = Arc::clone(&plan_hits_clone);
|
||||
async move {
|
||||
*plan_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({"action": "proxy_public"}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/report-sync",
|
||||
any(move |request: Request| {
|
||||
let seen_report_inner = Arc::clone(&seen_report_clone);
|
||||
async move {
|
||||
let (_parts, body) = request.into_parts();
|
||||
let _raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
|
||||
*seen_report_inner.lock().expect("mutex should lock") = true;
|
||||
Json(json!({"ok": true}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1/messages",
|
||||
any(move |_request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::IM_A_TEAPOT, Body::from("public-route-hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let execution_runtime = Router::new().route(
|
||||
"/v1/execute/sync",
|
||||
any(move |request: Request| {
|
||||
let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone);
|
||||
async move {
|
||||
let (parts, body) = request.into_parts();
|
||||
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
|
||||
let payload: serde_json::Value = serde_json::from_slice(&raw_body)
|
||||
.expect("execution runtime payload should parse");
|
||||
*seen_execution_runtime_inner
|
||||
.lock()
|
||||
.expect("mutex should lock") = Some(SeenExecutionRuntimeSyncRequest {
|
||||
trace_id: parts
|
||||
.headers
|
||||
.get(TRACE_ID_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
url: payload
|
||||
.get("url")
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
model: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("model"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
auth_header_value: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-api-key"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
endpoint_tag: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-endpoint-tag"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
metadata_mode: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("metadata"))
|
||||
.and_then(|value| value.get("mode"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
metadata_source: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("metadata"))
|
||||
.and_then(|value| value.get("source"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
proxy_node_id: payload
|
||||
.get("proxy")
|
||||
.and_then(|value| value.get("node_id"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
tls_profile: payload
|
||||
.get("tls_profile")
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
});
|
||||
Json(json!({
|
||||
"request_id": "trace-claude-chat-local-123",
|
||||
"status_code": 200,
|
||||
"headers": {
|
||||
"content-type": "application/json"
|
||||
},
|
||||
"body": {
|
||||
"json_body": {
|
||||
"id": "msg-local-claude-123",
|
||||
"type": "message",
|
||||
"model": "claude-sonnet-4-5-upstream",
|
||||
"role": "assistant",
|
||||
"content": [],
|
||||
"usage": {
|
||||
"input_tokens": 2,
|
||||
"output_tokens": 3
|
||||
}
|
||||
}
|
||||
},
|
||||
"telemetry": {
|
||||
"elapsed_ms": 29
|
||||
}
|
||||
}))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("sk-client-claude-chat-local")),
|
||||
sample_auth_snapshot("api-key-claude-local-1", "user-claude-local-1"),
|
||||
)]));
|
||||
let candidate_selection_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
sample_candidate_row(),
|
||||
]));
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider_catalog_provider()],
|
||||
vec![sample_provider_catalog_endpoint()],
|
||||
vec![sample_provider_catalog_key()],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let gateway_state =
|
||||
build_state_with_execution_runtime_override(execution_runtime_url.clone())
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
|
||||
auth_repository,
|
||||
candidate_selection_repository,
|
||||
provider_catalog_repository,
|
||||
Arc::clone(&request_candidate_repository),
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
),
|
||||
);
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/messages"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header("x-api-key", "sk-client-claude-chat-local")
|
||||
.header(TRACE_ID_HEADER, "trace-claude-chat-local-123")
|
||||
.body(
|
||||
"{\"model\":\"claude-sonnet-4-5\",\"messages\":[],\"metadata\":{\"client\":\"desktop-claude\"}}",
|
||||
)
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let response_json: serde_json::Value = response.json().await.expect("body should parse");
|
||||
assert_eq!(response_json["model"], "claude-sonnet-4-5-upstream");
|
||||
|
||||
let seen_execution_runtime_request = seen_execution_runtime
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.clone()
|
||||
.expect("execution runtime sync should be captured");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.trace_id,
|
||||
"trace-claude-chat-local-123"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.url,
|
||||
"https://api.anthropic.example/custom/v1/messages"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.model,
|
||||
"claude-sonnet-4-5-upstream"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.auth_header_value,
|
||||
"sk-upstream-claude-chat"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.endpoint_tag,
|
||||
"claude-chat-local"
|
||||
);
|
||||
assert_eq!(seen_execution_runtime_request.metadata_mode, "safe");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.metadata_source,
|
||||
"desktop-claude"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.proxy_node_id,
|
||||
"proxy-node-claude-chat-local"
|
||||
);
|
||||
assert_eq!(seen_execution_runtime_request.tls_profile, "chrome_136");
|
||||
|
||||
let stored_candidates = request_candidate_repository
|
||||
.list_by_request_id("trace-claude-chat-local-123")
|
||||
.await
|
||||
.expect("request candidate trace should read");
|
||||
assert_eq!(stored_candidates.len(), 1);
|
||||
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success);
|
||||
|
||||
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
|
||||
assert!(
|
||||
!*seen_report.lock().expect("mutex should lock"),
|
||||
"report-sync should stay local when request candidate persistence is available"
|
||||
);
|
||||
|
||||
assert_eq!(*decision_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*plan_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
@@ -0,0 +1,435 @@
|
||||
use super::{
|
||||
any, build_router_with_state, build_state_with_execution_runtime_override,
|
||||
encrypt_python_fernet_plaintext, json, start_server, to_bytes, Arc, Body, Digest,
|
||||
InMemoryAuthApiKeySnapshotRepository, InMemoryMinimalCandidateSelectionReadRepository,
|
||||
InMemoryProviderCatalogReadRepository, InMemoryRequestCandidateRepository, Json, Mutex,
|
||||
Request, RequestCandidateReadRepository, RequestCandidateStatus, Router, Sha256, StatusCode,
|
||||
StoredAuthApiKeySnapshot, StoredMinimalCandidateSelectionRow, StoredProviderCatalogEndpoint,
|
||||
StoredProviderCatalogKey, StoredProviderCatalogProvider, StoredProviderModelMapping,
|
||||
DEVELOPMENT_ENCRYPTION_KEY, TRACE_ID_HEADER,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_executes_claude_cli_sync_via_local_decision_gate_with_local_sync_decision() {
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenExecutionRuntimeSyncRequest {
|
||||
trace_id: String,
|
||||
url: String,
|
||||
model: String,
|
||||
authorization: String,
|
||||
endpoint_tag: String,
|
||||
metadata_mode: String,
|
||||
metadata_source: String,
|
||||
proxy_node_id: String,
|
||||
tls_profile: String,
|
||||
}
|
||||
|
||||
fn hash_api_key(value: &str) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(value.as_bytes());
|
||||
format!("{:x}", hasher.finalize())
|
||||
}
|
||||
|
||||
fn sample_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
user_id.to_string(),
|
||||
"alice".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
Some(serde_json::json!(["claude"])),
|
||||
Some(serde_json::json!(["claude:cli"])),
|
||||
Some(serde_json::json!(["claude-code"])),
|
||||
api_key_id.to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
Some(4_102_444_800),
|
||||
Some(serde_json::json!(["claude"])),
|
||||
Some(serde_json::json!(["claude:cli"])),
|
||||
Some(serde_json::json!(["claude-code"])),
|
||||
)
|
||||
.expect("auth snapshot should build")
|
||||
}
|
||||
|
||||
fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow {
|
||||
StoredMinimalCandidateSelectionRow {
|
||||
provider_id: "provider-claude-cli-local-1".to_string(),
|
||||
provider_name: "claude".to_string(),
|
||||
provider_type: "custom".to_string(),
|
||||
provider_priority: 10,
|
||||
provider_is_active: true,
|
||||
endpoint_id: "endpoint-claude-cli-local-1".to_string(),
|
||||
endpoint_api_format: "claude:cli".to_string(),
|
||||
endpoint_api_family: Some("claude".to_string()),
|
||||
endpoint_kind: Some("cli".to_string()),
|
||||
endpoint_is_active: true,
|
||||
key_id: "key-claude-cli-local-1".to_string(),
|
||||
key_name: "prod".to_string(),
|
||||
key_auth_type: "bearer".to_string(),
|
||||
key_is_active: true,
|
||||
key_api_formats: Some(vec!["claude:cli".to_string()]),
|
||||
key_allowed_models: None,
|
||||
key_capabilities: None,
|
||||
key_internal_priority: 5,
|
||||
key_global_priority_by_format: Some(serde_json::json!({"claude:cli": 1})),
|
||||
model_id: "model-claude-cli-local-1".to_string(),
|
||||
global_model_id: "global-model-claude-cli-local-1".to_string(),
|
||||
global_model_name: "claude-code".to_string(),
|
||||
global_model_mappings: None,
|
||||
global_model_supports_streaming: Some(true),
|
||||
model_provider_model_name: "claude-code-upstream".to_string(),
|
||||
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
|
||||
name: "claude-code-upstream".to_string(),
|
||||
priority: 1,
|
||||
api_formats: Some(vec!["claude:cli".to_string()]),
|
||||
}]),
|
||||
model_supports_streaming: Some(true),
|
||||
model_is_active: true,
|
||||
model_is_available: true,
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
|
||||
StoredProviderCatalogProvider::new(
|
||||
"provider-claude-cli-local-1".to_string(),
|
||||
"claude".to_string(),
|
||||
Some("https://example.com".to_string()),
|
||||
"custom".to_string(),
|
||||
)
|
||||
.expect("provider should build")
|
||||
.with_transport_fields(
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
None,
|
||||
Some(2),
|
||||
Some(serde_json::json!({"url":"http://provider-proxy.internal:8080"})),
|
||||
Some(20.0),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
|
||||
StoredProviderCatalogEndpoint::new(
|
||||
"endpoint-claude-cli-local-1".to_string(),
|
||||
"provider-claude-cli-local-1".to_string(),
|
||||
"claude:cli".to_string(),
|
||||
Some("claude".to_string()),
|
||||
Some("cli".to_string()),
|
||||
true,
|
||||
)
|
||||
.expect("endpoint should build")
|
||||
.with_transport_fields(
|
||||
"https://api.anthropic.example".to_string(),
|
||||
Some(serde_json::json!([
|
||||
{"action":"set","key":"x-endpoint-tag","value":"claude-cli-local"}
|
||||
])),
|
||||
Some(serde_json::json!([
|
||||
{"action":"set","path":"metadata.mode","value":"safe"},
|
||||
{"action":"rename","from":"metadata.client","to":"metadata.source"}
|
||||
])),
|
||||
Some(2),
|
||||
Some("/custom/v1/messages".to_string()),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("endpoint transport should build")
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
|
||||
StoredProviderCatalogKey::new(
|
||||
"key-claude-cli-local-1".to_string(),
|
||||
"provider-claude-cli-local-1".to_string(),
|
||||
"prod".to_string(),
|
||||
"bearer".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
.with_transport_fields(
|
||||
Some(serde_json::json!(["claude:cli"])),
|
||||
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "sk-upstream-claude-cli")
|
||||
.expect("api key should encrypt"),
|
||||
None,
|
||||
None,
|
||||
Some(serde_json::json!({"claude:cli": 1})),
|
||||
None,
|
||||
None,
|
||||
Some(serde_json::json!({"enabled": true, "node_id":"proxy-node-claude-cli-local"})),
|
||||
Some(serde_json::json!({"tls_profile":"chrome_136"})),
|
||||
)
|
||||
.expect("key transport should build")
|
||||
}
|
||||
|
||||
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeSyncRequest>));
|
||||
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
|
||||
let seen_report = Arc::new(Mutex::new(false));
|
||||
let seen_report_clone = Arc::clone(&seen_report);
|
||||
let decision_hits = Arc::new(Mutex::new(0usize));
|
||||
let decision_hits_clone = Arc::clone(&decision_hits);
|
||||
let plan_hits = Arc::new(Mutex::new(0usize));
|
||||
let plan_hits_clone = Arc::clone(&plan_hits);
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/decision-sync",
|
||||
any(move |_request: Request| {
|
||||
let decision_hits_inner = Arc::clone(&decision_hits_clone);
|
||||
async move {
|
||||
*decision_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({"action": "proxy_public"}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/plan-sync",
|
||||
any(move |_request: Request| {
|
||||
let plan_hits_inner = Arc::clone(&plan_hits_clone);
|
||||
async move {
|
||||
*plan_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({"action": "proxy_public"}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/report-sync",
|
||||
any(move |request: Request| {
|
||||
let seen_report_inner = Arc::clone(&seen_report_clone);
|
||||
async move {
|
||||
let (_parts, body) = request.into_parts();
|
||||
let _raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
|
||||
*seen_report_inner.lock().expect("mutex should lock") = true;
|
||||
Json(json!({"ok": true}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1/messages",
|
||||
any(move |_request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::IM_A_TEAPOT, Body::from("public-route-hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let execution_runtime = Router::new().route(
|
||||
"/v1/execute/sync",
|
||||
any(move |request: Request| {
|
||||
let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone);
|
||||
async move {
|
||||
let (parts, body) = request.into_parts();
|
||||
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
|
||||
let payload: serde_json::Value = serde_json::from_slice(&raw_body)
|
||||
.expect("execution runtime payload should parse");
|
||||
*seen_execution_runtime_inner
|
||||
.lock()
|
||||
.expect("mutex should lock") = Some(SeenExecutionRuntimeSyncRequest {
|
||||
trace_id: parts
|
||||
.headers
|
||||
.get(TRACE_ID_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
url: payload
|
||||
.get("url")
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
model: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("model"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
authorization: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("authorization"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
endpoint_tag: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-endpoint-tag"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
metadata_mode: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("metadata"))
|
||||
.and_then(|value| value.get("mode"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
metadata_source: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("metadata"))
|
||||
.and_then(|value| value.get("source"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
proxy_node_id: payload
|
||||
.get("proxy")
|
||||
.and_then(|value| value.get("node_id"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
tls_profile: payload
|
||||
.get("tls_profile")
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
});
|
||||
Json(json!({
|
||||
"request_id": "trace-claude-cli-local-sync-123",
|
||||
"status_code": 200,
|
||||
"headers": {
|
||||
"content-type": "application/json"
|
||||
},
|
||||
"body": {
|
||||
"json_body": {
|
||||
"id": "msg-local-claude-cli-123",
|
||||
"type": "message",
|
||||
"model": "claude-code-upstream",
|
||||
"role": "assistant",
|
||||
"content": [],
|
||||
"usage": {
|
||||
"input_tokens": 2,
|
||||
"output_tokens": 3
|
||||
}
|
||||
}
|
||||
},
|
||||
"telemetry": {
|
||||
"elapsed_ms": 29
|
||||
}
|
||||
}))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("sk-client-claude-cli-local")),
|
||||
sample_auth_snapshot("api-key-claude-cli-local-1", "user-claude-cli-local-1"),
|
||||
)]));
|
||||
let candidate_selection_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
sample_candidate_row(),
|
||||
]));
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider_catalog_provider()],
|
||||
vec![sample_provider_catalog_endpoint()],
|
||||
vec![sample_provider_catalog_key()],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url.clone())
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
|
||||
auth_repository,
|
||||
candidate_selection_repository,
|
||||
provider_catalog_repository,
|
||||
Arc::clone(&request_candidate_repository),
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
),
|
||||
);
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/messages"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(
|
||||
http::header::AUTHORIZATION,
|
||||
"Bearer sk-client-claude-cli-local",
|
||||
)
|
||||
.header(TRACE_ID_HEADER, "trace-claude-cli-local-sync-123")
|
||||
.body(
|
||||
"{\"model\":\"claude-code\",\"messages\":[],\"metadata\":{\"client\":\"desktop-claude-cli\"}}",
|
||||
)
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
let status = response.status();
|
||||
let response_body = response.text().await.expect("body should read");
|
||||
assert!(
|
||||
status == StatusCode::OK,
|
||||
"unexpected status={status} body={response_body} decision_hits={} plan_hits={} public_hits={}",
|
||||
*decision_hits.lock().expect("mutex should lock"),
|
||||
*plan_hits.lock().expect("mutex should lock"),
|
||||
*public_hits.lock().expect("mutex should lock"),
|
||||
);
|
||||
let response_json: serde_json::Value =
|
||||
serde_json::from_str(&response_body).expect("body should parse");
|
||||
assert_eq!(response_json["model"], "claude-code-upstream");
|
||||
|
||||
let seen_execution_runtime_request = seen_execution_runtime
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.clone()
|
||||
.expect("execution runtime sync should be captured");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.trace_id,
|
||||
"trace-claude-cli-local-sync-123"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.url,
|
||||
"https://api.anthropic.example/custom/v1/messages"
|
||||
);
|
||||
assert_eq!(seen_execution_runtime_request.model, "claude-code-upstream");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.authorization,
|
||||
"Bearer sk-upstream-claude-cli"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.endpoint_tag,
|
||||
"claude-cli-local"
|
||||
);
|
||||
assert_eq!(seen_execution_runtime_request.metadata_mode, "safe");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.metadata_source,
|
||||
"desktop-claude-cli"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.proxy_node_id,
|
||||
"proxy-node-claude-cli-local"
|
||||
);
|
||||
assert_eq!(seen_execution_runtime_request.tls_profile, "chrome_136");
|
||||
|
||||
let stored_candidates = request_candidate_repository
|
||||
.list_by_request_id("trace-claude-cli-local-sync-123")
|
||||
.await
|
||||
.expect("request candidate trace should read");
|
||||
assert_eq!(stored_candidates.len(), 1);
|
||||
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success);
|
||||
|
||||
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
|
||||
assert!(
|
||||
!*seen_report.lock().expect("mutex should lock"),
|
||||
"report-sync should stay local when request candidate persistence is available"
|
||||
);
|
||||
|
||||
assert_eq!(*decision_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*plan_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
45
apps/aether-gateway/src/tests/ai_execute/sync/claude/mod.rs
Normal file
45
apps/aether-gateway/src/tests/ai_execute/sync/claude/mod.rs
Normal file
@@ -0,0 +1,45 @@
|
||||
use std::convert::Infallible;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use axum::body::{to_bytes, Body, Bytes};
|
||||
use axum::response::Response;
|
||||
use axum::routing::any;
|
||||
use axum::{extract::Request, Json, Router};
|
||||
use http::header::{HeaderName, HeaderValue};
|
||||
use http::StatusCode;
|
||||
use serde_json::json;
|
||||
|
||||
use crate::constants::{
|
||||
CONTROL_EXECUTED_HEADER, CONTROL_EXECUTE_FALLBACK_HEADER, DEPENDENCY_REASON_HEADER,
|
||||
EXECUTION_PATH_EXECUTION_RUNTIME_STREAM, EXECUTION_PATH_EXECUTION_RUNTIME_SYNC,
|
||||
EXECUTION_PATH_HEADER, EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS,
|
||||
LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER, TRACE_ID_HEADER,
|
||||
};
|
||||
|
||||
use super::{
|
||||
build_router, build_router_with_execution_runtime_override, build_router_with_state,
|
||||
build_state_with_execution_runtime_override, start_server, wait_until, AppState,
|
||||
FrontdoorCorsConfig, FrontdoorUserRpmConfig, GatewayFallbackMetricKind, GatewayFallbackReason,
|
||||
UsageRuntimeConfig, VideoTaskTruthSourceMode,
|
||||
};
|
||||
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
|
||||
use aether_data::repository::auth::{
|
||||
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot,
|
||||
};
|
||||
use aether_data::repository::candidate_selection::{
|
||||
InMemoryMinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
|
||||
StoredProviderModelMapping,
|
||||
};
|
||||
use aether_data::repository::candidates::{
|
||||
InMemoryRequestCandidateRepository, RequestCandidateReadRepository, RequestCandidateStatus,
|
||||
};
|
||||
use aether_data::repository::provider_catalog::{
|
||||
InMemoryProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogProvider,
|
||||
};
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
mod claude_code;
|
||||
mod kiro;
|
||||
mod local_chat;
|
||||
mod local_cli;
|
||||
1042
apps/aether-gateway/src/tests/ai_execute/sync/cli.rs
Normal file
1042
apps/aether-gateway/src/tests/ai_execute/sync/cli.rs
Normal file
File diff suppressed because it is too large
Load Diff
1951
apps/aether-gateway/src/tests/ai_execute/sync/gemini/cli.rs
Normal file
1951
apps/aether-gateway/src/tests/ai_execute/sync/gemini/cli.rs
Normal file
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,435 @@
|
||||
use super::{
|
||||
any, build_router_with_state, build_state_with_execution_runtime_override,
|
||||
encrypt_python_fernet_plaintext, json, start_server, to_bytes, Arc, Body, Digest,
|
||||
InMemoryAuthApiKeySnapshotRepository, InMemoryMinimalCandidateSelectionReadRepository,
|
||||
InMemoryProviderCatalogReadRepository, InMemoryRequestCandidateRepository, Json, Mutex,
|
||||
Request, RequestCandidateReadRepository, RequestCandidateStatus, Router, Sha256, StatusCode,
|
||||
StoredAuthApiKeySnapshot, StoredMinimalCandidateSelectionRow, StoredProviderCatalogEndpoint,
|
||||
StoredProviderCatalogKey, StoredProviderCatalogProvider, StoredProviderModelMapping,
|
||||
DEVELOPMENT_ENCRYPTION_KEY, TRACE_ID_HEADER,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_executes_gemini_chat_sync_via_local_decision_gate_with_local_sync_decision() {
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenExecutionRuntimeSyncRequest {
|
||||
trace_id: String,
|
||||
url: String,
|
||||
has_model_field: bool,
|
||||
auth_header_value: String,
|
||||
exact_temperature: f64,
|
||||
endpoint_tag: String,
|
||||
metadata_mode: String,
|
||||
metadata_source: String,
|
||||
tool_config_present: bool,
|
||||
proxy_node_id: String,
|
||||
tls_profile: String,
|
||||
}
|
||||
|
||||
fn hash_api_key(value: &str) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(value.as_bytes());
|
||||
format!("{:x}", hasher.finalize())
|
||||
}
|
||||
|
||||
fn sample_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
user_id.to_string(),
|
||||
"alice".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
Some(serde_json::json!(["gemini"])),
|
||||
Some(serde_json::json!(["gemini:chat"])),
|
||||
Some(serde_json::json!(["gemini-2.5-pro"])),
|
||||
api_key_id.to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
Some(4_102_444_800),
|
||||
Some(serde_json::json!(["gemini"])),
|
||||
Some(serde_json::json!(["gemini:chat"])),
|
||||
Some(serde_json::json!(["gemini-2.5-pro"])),
|
||||
)
|
||||
.expect("auth snapshot should build")
|
||||
}
|
||||
|
||||
fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow {
|
||||
StoredMinimalCandidateSelectionRow {
|
||||
provider_id: "provider-gemini-local-1".to_string(),
|
||||
provider_name: "gemini".to_string(),
|
||||
provider_type: "custom".to_string(),
|
||||
provider_priority: 10,
|
||||
provider_is_active: true,
|
||||
endpoint_id: "endpoint-gemini-local-1".to_string(),
|
||||
endpoint_api_format: "gemini:chat".to_string(),
|
||||
endpoint_api_family: Some("gemini".to_string()),
|
||||
endpoint_kind: Some("chat".to_string()),
|
||||
endpoint_is_active: true,
|
||||
key_id: "key-gemini-local-1".to_string(),
|
||||
key_name: "prod".to_string(),
|
||||
key_auth_type: "api_key".to_string(),
|
||||
key_is_active: true,
|
||||
key_api_formats: Some(vec!["gemini:chat".to_string()]),
|
||||
key_allowed_models: None,
|
||||
key_capabilities: None,
|
||||
key_internal_priority: 5,
|
||||
key_global_priority_by_format: Some(serde_json::json!({"gemini:chat": 1})),
|
||||
model_id: "model-gemini-local-1".to_string(),
|
||||
global_model_id: "global-model-gemini-local-1".to_string(),
|
||||
global_model_name: "gemini-2.5-pro".to_string(),
|
||||
global_model_mappings: None,
|
||||
global_model_supports_streaming: Some(true),
|
||||
model_provider_model_name: "gemini-2.5-pro-upstream".to_string(),
|
||||
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
|
||||
name: "gemini-2.5-pro-upstream".to_string(),
|
||||
priority: 1,
|
||||
api_formats: Some(vec!["gemini:chat".to_string()]),
|
||||
}]),
|
||||
model_supports_streaming: Some(true),
|
||||
model_is_active: true,
|
||||
model_is_available: true,
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
|
||||
StoredProviderCatalogProvider::new(
|
||||
"provider-gemini-local-1".to_string(),
|
||||
"gemini".to_string(),
|
||||
Some("https://example.com".to_string()),
|
||||
"custom".to_string(),
|
||||
)
|
||||
.expect("provider should build")
|
||||
.with_transport_fields(
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
None,
|
||||
Some(2),
|
||||
Some(serde_json::json!({"url":"http://provider-proxy.internal:8080"})),
|
||||
Some(20.0),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
|
||||
StoredProviderCatalogEndpoint::new(
|
||||
"endpoint-gemini-local-1".to_string(),
|
||||
"provider-gemini-local-1".to_string(),
|
||||
"gemini:chat".to_string(),
|
||||
Some("gemini".to_string()),
|
||||
Some("chat".to_string()),
|
||||
true,
|
||||
)
|
||||
.expect("endpoint should build")
|
||||
.with_transport_fields(
|
||||
"https://generativelanguage.googleapis.com".to_string(),
|
||||
Some(serde_json::json!([
|
||||
{"action":"set","key":"x-endpoint-tag","value":"gemini-chat-local"}
|
||||
])),
|
||||
Some(serde_json::json!([
|
||||
{"action":"set","path":"metadata.mode","value":"safe"},
|
||||
{"action":"rename","from":"metadata.client","to":"metadata.source"},
|
||||
{"action":"drop","path":"toolConfig"}
|
||||
])),
|
||||
Some(2),
|
||||
Some("/custom/v1beta/models/gemini-2.5-pro-upstream:generateContent".to_string()),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("endpoint transport should build")
|
||||
}
|
||||
|
||||
fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
|
||||
StoredProviderCatalogKey::new(
|
||||
"key-gemini-local-1".to_string(),
|
||||
"provider-gemini-local-1".to_string(),
|
||||
"prod".to_string(),
|
||||
"api_key".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
.with_transport_fields(
|
||||
Some(serde_json::json!(["gemini:chat"])),
|
||||
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "sk-upstream-gemini-chat")
|
||||
.expect("api key should encrypt"),
|
||||
None,
|
||||
None,
|
||||
Some(serde_json::json!({"gemini:chat": 1})),
|
||||
None,
|
||||
None,
|
||||
Some(serde_json::json!({"enabled": true, "node_id":"proxy-node-gemini-chat-local"})),
|
||||
Some(serde_json::json!({"tls_profile":"chrome_136"})),
|
||||
)
|
||||
.expect("key transport should build")
|
||||
}
|
||||
|
||||
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeSyncRequest>));
|
||||
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
|
||||
let seen_report = Arc::new(Mutex::new(false));
|
||||
let seen_report_clone = Arc::clone(&seen_report);
|
||||
let decision_hits = Arc::new(Mutex::new(0usize));
|
||||
let decision_hits_clone = Arc::clone(&decision_hits);
|
||||
let plan_hits = Arc::new(Mutex::new(0usize));
|
||||
let plan_hits_clone = Arc::clone(&plan_hits);
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/decision-sync",
|
||||
any(move |_request: Request| {
|
||||
let decision_hits_inner = Arc::clone(&decision_hits_clone);
|
||||
async move {
|
||||
*decision_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({"action": "proxy_public"}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/plan-sync",
|
||||
any(move |_request: Request| {
|
||||
let plan_hits_inner = Arc::clone(&plan_hits_clone);
|
||||
async move {
|
||||
*plan_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({"action": "proxy_public"}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/internal/gateway/report-sync",
|
||||
any(move |request: Request| {
|
||||
let seen_report_inner = Arc::clone(&seen_report_clone);
|
||||
async move {
|
||||
let (_parts, body) = request.into_parts();
|
||||
let _raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
|
||||
*seen_report_inner.lock().expect("mutex should lock") = true;
|
||||
Json(json!({"ok": true}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1beta/models/gemini-2.5-pro:generateContent",
|
||||
any(move |_request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::IM_A_TEAPOT, Body::from("public-route-hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let execution_runtime = Router::new().route(
|
||||
"/v1/execute/sync",
|
||||
any(move |request: Request| {
|
||||
let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone);
|
||||
async move {
|
||||
let (parts, body) = request.into_parts();
|
||||
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
|
||||
let payload: serde_json::Value = serde_json::from_slice(&raw_body)
|
||||
.expect("execution runtime payload should parse");
|
||||
*seen_execution_runtime_inner
|
||||
.lock()
|
||||
.expect("mutex should lock") = Some(SeenExecutionRuntimeSyncRequest {
|
||||
trace_id: parts
|
||||
.headers
|
||||
.get(TRACE_ID_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
url: payload
|
||||
.get("url")
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
has_model_field: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("model"))
|
||||
.is_some(),
|
||||
auth_header_value: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-goog-api-key"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
exact_temperature: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("generationConfig"))
|
||||
.and_then(|value| value.get("temperature"))
|
||||
.and_then(|value| value.as_f64())
|
||||
.unwrap_or_default(),
|
||||
endpoint_tag: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-endpoint-tag"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
metadata_mode: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("metadata"))
|
||||
.and_then(|value| value.get("mode"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
metadata_source: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("metadata"))
|
||||
.and_then(|value| value.get("source"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
tool_config_present: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("json_body"))
|
||||
.and_then(|value| value.get("toolConfig"))
|
||||
.is_some(),
|
||||
proxy_node_id: payload
|
||||
.get("proxy")
|
||||
.and_then(|value| value.get("node_id"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
tls_profile: payload
|
||||
.get("tls_profile")
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
});
|
||||
Json(json!({
|
||||
"request_id": "trace-gemini-chat-local-123",
|
||||
"status_code": 200,
|
||||
"headers": {
|
||||
"content-type": "application/json"
|
||||
},
|
||||
"body": {
|
||||
"json_body": {
|
||||
"candidates": [],
|
||||
"usageMetadata": {
|
||||
"promptTokenCount": 1,
|
||||
"candidatesTokenCount": 2,
|
||||
"totalTokenCount": 3
|
||||
}
|
||||
}
|
||||
},
|
||||
"telemetry": {
|
||||
"elapsed_ms": 27
|
||||
}
|
||||
}))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("client-gemini-chat-local-key")),
|
||||
sample_auth_snapshot("api-key-gemini-local-1", "user-gemini-local-1"),
|
||||
)]));
|
||||
let candidate_selection_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
sample_candidate_row(),
|
||||
]));
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider_catalog_provider()],
|
||||
vec![sample_provider_catalog_endpoint()],
|
||||
vec![sample_provider_catalog_key()],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let gateway_state =
|
||||
build_state_with_execution_runtime_override(execution_runtime_url.clone())
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
|
||||
auth_repository,
|
||||
candidate_selection_repository,
|
||||
provider_catalog_repository,
|
||||
Arc::clone(&request_candidate_repository),
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
),
|
||||
);
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!(
|
||||
"{gateway_url}/v1beta/models/gemini-2.5-pro:generateContent?key=client-gemini-chat-local-key&alt=sse"
|
||||
))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(TRACE_ID_HEADER, "trace-gemini-chat-local-123")
|
||||
.body(
|
||||
"{\"contents\":[],\"generationConfig\":{\"temperature\":0.2},\"metadata\":{\"client\":\"desktop-gemini\"},\"toolConfig\":{\"functionCallingConfig\":{\"mode\":\"AUTO\"}}}",
|
||||
)
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
|
||||
let seen_execution_runtime_request = seen_execution_runtime
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.clone()
|
||||
.expect("execution runtime sync should be captured");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.trace_id,
|
||||
"trace-gemini-chat-local-123"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.url,
|
||||
"https://generativelanguage.googleapis.com/custom/v1beta/models/gemini-2.5-pro-upstream:generateContent?alt=sse"
|
||||
);
|
||||
assert!(!seen_execution_runtime_request.has_model_field);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.auth_header_value,
|
||||
"sk-upstream-gemini-chat"
|
||||
);
|
||||
assert!((seen_execution_runtime_request.exact_temperature - 0.2).abs() < f64::EPSILON);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.endpoint_tag,
|
||||
"gemini-chat-local"
|
||||
);
|
||||
assert_eq!(seen_execution_runtime_request.metadata_mode, "safe");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.metadata_source,
|
||||
"desktop-gemini"
|
||||
);
|
||||
assert!(!seen_execution_runtime_request.tool_config_present);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.proxy_node_id,
|
||||
"proxy-node-gemini-chat-local"
|
||||
);
|
||||
assert_eq!(seen_execution_runtime_request.tls_profile, "chrome_136");
|
||||
|
||||
let stored_candidates = request_candidate_repository
|
||||
.list_by_request_id("trace-gemini-chat-local-123")
|
||||
.await
|
||||
.expect("request candidate trace should read");
|
||||
assert_eq!(stored_candidates.len(), 1);
|
||||
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success);
|
||||
|
||||
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
|
||||
assert!(
|
||||
!*seen_report.lock().expect("mutex should lock"),
|
||||
"report-sync should stay local when request candidate persistence is available"
|
||||
);
|
||||
|
||||
assert_eq!(*decision_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*plan_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
43
apps/aether-gateway/src/tests/ai_execute/sync/gemini/mod.rs
Normal file
43
apps/aether-gateway/src/tests/ai_execute/sync/gemini/mod.rs
Normal file
@@ -0,0 +1,43 @@
|
||||
use std::convert::Infallible;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use axum::body::{to_bytes, Body, Bytes};
|
||||
use axum::response::Response;
|
||||
use axum::routing::any;
|
||||
use axum::{extract::Request, Json, Router};
|
||||
use http::header::{HeaderName, HeaderValue};
|
||||
use http::StatusCode;
|
||||
use serde_json::json;
|
||||
|
||||
use crate::constants::{
|
||||
CONTROL_EXECUTED_HEADER, CONTROL_EXECUTE_FALLBACK_HEADER, DEPENDENCY_REASON_HEADER,
|
||||
EXECUTION_PATH_EXECUTION_RUNTIME_STREAM, EXECUTION_PATH_EXECUTION_RUNTIME_SYNC,
|
||||
EXECUTION_PATH_HEADER, EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS,
|
||||
LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER, TRACE_ID_HEADER,
|
||||
};
|
||||
|
||||
use super::{
|
||||
build_router, build_router_with_execution_runtime_override, build_router_with_state,
|
||||
build_state_with_execution_runtime_override, start_server, wait_until, AppState,
|
||||
FrontdoorCorsConfig, FrontdoorUserRpmConfig, GatewayFallbackMetricKind, GatewayFallbackReason,
|
||||
UsageRuntimeConfig, VideoTaskTruthSourceMode,
|
||||
};
|
||||
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
|
||||
use aether_data::repository::auth::{
|
||||
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot,
|
||||
};
|
||||
use aether_data::repository::candidate_selection::{
|
||||
InMemoryMinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
|
||||
StoredProviderModelMapping,
|
||||
};
|
||||
use aether_data::repository::candidates::{
|
||||
InMemoryRequestCandidateRepository, RequestCandidateReadRepository, RequestCandidateStatus,
|
||||
};
|
||||
use aether_data::repository::provider_catalog::{
|
||||
InMemoryProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogProvider,
|
||||
};
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
mod cli;
|
||||
mod local_chat;
|
||||
29
apps/aether-gateway/src/tests/ai_execute/sync/mod.rs
Normal file
29
apps/aether-gateway/src/tests/ai_execute/sync/mod.rs
Normal file
@@ -0,0 +1,29 @@
|
||||
use std::convert::Infallible;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use axum::body::{to_bytes, Body, Bytes};
|
||||
use axum::response::Response;
|
||||
use axum::routing::any;
|
||||
use axum::{extract::Request, Json, Router};
|
||||
use http::header::{HeaderName, HeaderValue};
|
||||
use http::StatusCode;
|
||||
use serde_json::json;
|
||||
|
||||
use crate::constants::{
|
||||
CONTROL_EXECUTED_HEADER, CONTROL_EXECUTE_FALLBACK_HEADER, DEPENDENCY_REASON_HEADER,
|
||||
EXECUTION_PATH_EXECUTION_RUNTIME_STREAM, EXECUTION_PATH_EXECUTION_RUNTIME_SYNC,
|
||||
EXECUTION_PATH_HEADER, EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS,
|
||||
LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER, TRACE_ID_HEADER,
|
||||
};
|
||||
|
||||
use super::{
|
||||
build_router, build_router_with_execution_runtime_override, build_router_with_state,
|
||||
build_state_with_execution_runtime_override, start_server, wait_until, AppState,
|
||||
FrontdoorCorsConfig, FrontdoorUserRpmConfig, GatewayFallbackMetricKind, GatewayFallbackReason,
|
||||
UsageRuntimeConfig, VideoTaskTruthSourceMode,
|
||||
};
|
||||
|
||||
mod chat;
|
||||
mod claude;
|
||||
mod cli;
|
||||
mod gemini;
|
||||
Reference in New Issue
Block a user