i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-13 06:30:20 +08:00
Code Issues Packages Projects Releases Wiki Activity
Files
20b27a13b247116acc209e0b601f083c1659ca0d
Aether/apps/aether-gateway/src/tests/ai_execute/finalize_local.rs
T
MMEXA 20b27a13b2 feat(codex): 按操作语义路由 Responses V2 压缩 
(cherry picked from commit 2fc604e047)
2026-07-16 00:34:42 +08:00

2727 lines
110 KiB
Rust
Raw Blame History

use super::{
any, build_router_with_state, build_state_with_execution_runtime_override, json, start_server,
to_bytes, Arc, Body, Bytes, HeaderName, HeaderValue, Json, Mutex, Request, Response, Router,
StatusCode, CONTROL_EXECUTED_HEADER, EXECUTION_PATH_HEADER,
LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER, TRACE_ID_HEADER,
};
use crate::data::GatewayDataState;
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
use aether_data::repository::auth::{
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot,
};
use aether_data::repository::candidate_selection::InMemoryMinimalCandidateSelectionReadRepository;
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data::repository::usage::InMemoryUsageReadRepository;
use aether_data_contracts::repository::candidate_selection::{
StoredMinimalCandidateSelectionRow, StoredProviderModelMapping,
};
use aether_data_contracts::repository::candidates::{
RequestCandidateReadRepository, RequestCandidateStatus,
};
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use sha2::{Digest, Sha256};
const OPENAI_CHAT_FINALIZE_TEST_STACK_BYTES: usize = 16 * 1024 * 1024;
fn run_openai_chat_finalize_test<F, Fut>(test_name: &'static str, make_future: F)
where
F: FnOnce() -> Fut + Send + 'static,
Fut: std::future::Future<Output = ()> + 'static,
{
let handle = std::thread::Builder::new()
.name(test_name.to_string())
.stack_size(OPENAI_CHAT_FINALIZE_TEST_STACK_BYTES)
.spawn(move || {
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("test runtime should build");
runtime.block_on(make_future());
})
.expect("openai chat finalize test thread should spawn");
if let Err(payload) = handle.join() {
std::panic::resume_unwind(payload);
}
}
#[test]
fn gateway_executes_openai_chat_sync_upstream_stream_via_local_finalize_response() {
run_openai_chat_finalize_test(
"gateway_executes_openai_chat_sync_upstream_stream_via_local_finalize_response",
gateway_executes_openai_chat_sync_upstream_stream_via_local_finalize_response_impl,
);
}
async fn gateway_executes_openai_chat_sync_upstream_stream_via_local_finalize_response_impl() {
use base64::Engine as _;
#[derive(Debug, Clone)]
struct SeenRemoteExecutionRuntimeRequest {
trace_id: String,
request_id: String,
url: String,
model: String,
authorization: String,
}
#[derive(Debug, Clone)]
struct SeenReportSyncRequest {
report_kind: 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: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_candidate_row() -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: "provider-openai-finalize-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-finalize-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-finalize-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-finalize-local-1".to_string(),
global_model_id: "global-model-openai-finalize-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()]),
endpoint_ids: None,
operations: None,
}]),
model_supports_streaming: Some(true),
model_is_active: true,
model_is_available: true,
}
}
fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
"provider-openai-finalize-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,
true,
None,
Some(2),
None,
Some(20.0),
None,
None,
)
}
fn sample_provider_catalog_endpoint(base_url: &str) -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
"endpoint-openai-finalize-local-1".to_string(),
"provider-openai-finalize-local-1".to_string(),
"openai:chat".to_string(),
Some("openai".to_string()),
Some("chat".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
base_url.to_string(),
None,
None,
Some(2),
None,
None,
None,
None,
)
.expect("endpoint transport should build")
}
fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
StoredProviderCatalogKey::new(
"key-openai-finalize-local-1".to_string(),
"provider-openai-finalize-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")
}
let seen_remote_execution_runtime =
Arc::new(Mutex::new(None::<SeenRemoteExecutionRuntimeRequest>));
let seen_remote_execution_runtime_clone = Arc::clone(&seen_remote_execution_runtime);
let seen_report = Arc::new(Mutex::new(None::<SeenReportSyncRequest>));
let seen_report_clone = Arc::clone(&seen_report);
let finalize_hits = Arc::new(Mutex::new(0usize));
let finalize_hits_clone = Arc::clone(&finalize_hits);
let report_hits = Arc::new(Mutex::new(0usize));
let report_hits_clone = Arc::clone(&report_hits);
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/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/finalize-sync",
any(move |_request: Request| {
let finalize_hits_inner = Arc::clone(&finalize_hits_clone);
async move {
*finalize_hits_inner.lock().expect("mutex should lock") += 1;
tokio::time::sleep(std::time::Duration::from_millis(500)).await;
let mut response = Response::builder()
.status(StatusCode::OK)
.body(Body::from(
"{\"id\":\"ignored-finalize-response\",\"object\":\"chat.completion\",\"choices\":[]}",
))
.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/report-sync",
any(move |request: Request| {
let report_hits_inner = Arc::clone(&report_hits_clone);
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");
*report_hits_inner.lock().expect("mutex should lock") += 1;
*seen_report_inner.lock().expect("mutex should lock") =
Some(SeenReportSyncRequest {
report_kind: payload
.get("report_kind")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
});
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/sync",
any(move |request: Request| {
let seen_remote_execution_runtime_inner = Arc::clone(&seen_remote_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_remote_execution_runtime_inner
.lock()
.expect("mutex should lock") = Some(SeenRemoteExecutionRuntimeRequest {
trace_id: parts
.headers
.get(TRACE_ID_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
request_id: payload
.get("request_id")
.and_then(|value| value.as_str())
.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(),
});
Json(json!({
"request_id": "trace-openai-chat-stream-sync-direct-123",
"status_code": 200,
"headers": {
"content-type": "text/event-stream"
},
"body": {
"body_bytes_b64": base64::engine::general_purpose::STANDARD.encode(
concat!(
"data: {\"id\":\"chatcmpl-stream-sync-upstream-123\",\"object\":\"chat.completion.chunk\",",
"\"created\":1,\"model\":\"gpt-5\",\"choices\":[{\"index\":0,",
"\"delta\":{\"role\":\"assistant\",\"content\":\"Hello\"},\"finish_reason\":null}]}\n\n",
"data: {\"id\":\"chatcmpl-stream-sync-upstream-123\",\"object\":\"chat.completion.chunk\",",
"\"model\":\"gpt-5\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\" world\"},",
"\"finish_reason\":null}]}\n\n",
"data: {\"id\":\"chatcmpl-stream-sync-upstream-123\",\"object\":\"chat.completion.chunk\",",
"\"model\":\"gpt-5\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n",
"data: [DONE]\n\n"
)
)
},
"telemetry": {
"elapsed_ms": 31
}
}))
}
}),
);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-client-openai-finalize-local")),
sample_auth_snapshot(
"api-key-openai-finalize-local-1",
"user-openai-finalize-local-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(
"https://api.openai.example/v1",
)],
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(
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 started_at = std::time::Instant::now();
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-finalize-local",
)
.header(TRACE_ID_HEADER, "trace-openai-chat-stream-sync-direct-123")
.body("{\"model\":\"gpt-5\",\"messages\":[]}")
.send()
.await
.expect("request should succeed");
let elapsed = started_at.elapsed();
let response_status = response.status();
let execution_path = response
.headers()
.get(EXECUTION_PATH_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string();
let miss_reason = response
.headers()
.get(LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string();
let response_body = response.text().await.expect("body should read");
let seen_remote_execution_runtime_debug = seen_remote_execution_runtime
.lock()
.expect("mutex should lock")
.clone();
let stored_candidates_debug = request_candidate_repository
.list_by_request_id("trace-openai-chat-stream-sync-direct-123")
.await
.expect("request candidate trace should read");
assert_eq!(
response_status,
StatusCode::OK,
"unexpected gateway response: path={execution_path} miss={miss_reason} body={response_body} execution_runtime_seen={seen_remote_execution_runtime_debug:?} stored_candidates={stored_candidates_debug:?}"
);
let response_json: serde_json::Value =
serde_json::from_str(&response_body).expect("body should parse");
assert_eq!(
response_json,
json!({
"id": "chatcmpl-stream-sync-upstream-123",
"object": "chat.completion",
"created": 1,
"model": "gpt-5",
"choices": [{
"index": 0,
"message": {
"role": "assistant",
"content": "Hello world"
},
"finish_reason": "stop"
}]
})
);
assert!(
elapsed < std::time::Duration::from_millis(3_000),
"response took unexpectedly long for local finalize path: elapsed={elapsed:?} finalize_hits={} report_hits={}",
*finalize_hits.lock().expect("mutex should lock"),
*report_hits.lock().expect("mutex should lock"),
);
let mut stored_candidates = Vec::new();
for _ in 0..50 {
stored_candidates = request_candidate_repository
.list_by_request_id("trace-openai-chat-stream-sync-direct-123")
.await
.expect("request candidate trace should read");
if stored_candidates.len() == 1
&& stored_candidates[0].status == RequestCandidateStatus::Success
{
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
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_eq!(
*report_hits.lock().expect("mutex should lock"),
0,
"report-sync should stay local when request candidate persistence is available"
);
assert!(
seen_report.lock().expect("mutex should lock").is_none(),
"remote report payload should not be emitted when local persistence is available"
);
let seen_remote_execution_runtime_request = seen_remote_execution_runtime
.lock()
.expect("mutex should lock")
.clone()
.expect("remote execution runtime plan should be captured");
assert_eq!(
seen_remote_execution_runtime_request.trace_id,
"trace-openai-chat-stream-sync-direct-123"
);
assert_eq!(
seen_remote_execution_runtime_request.request_id,
"trace-openai-chat-stream-sync-direct-123"
);
assert_eq!(
seen_remote_execution_runtime_request.url,
"https://api.openai.example/v1/chat/completions"
);
assert_eq!(
seen_remote_execution_runtime_request.model,
"gpt-5-upstream"
);
assert_eq!(
seen_remote_execution_runtime_request.authorization,
"Bearer sk-upstream-openai"
);
assert_eq!(
*finalize_hits.lock().expect("mutex should lock"),
0,
"finalize-sync should not be called when local finalize can downgrade to success report"
);
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();
}
#[test]
fn gateway_executes_openai_chat_cross_format_upstream_stream_via_local_finalize_response() {
run_openai_chat_finalize_test(
"gateway_executes_openai_chat_cross_format_upstream_stream_via_local_finalize_response",
gateway_executes_openai_chat_cross_format_upstream_stream_via_local_finalize_response_impl,
);
}
async fn gateway_executes_openai_chat_cross_format_upstream_stream_via_local_finalize_response_impl(
) {
use base64::Engine as _;
#[derive(Debug, Clone)]
struct SeenRemoteExecutionRuntimeRequest {
trace_id: String,
url: String,
provider_model: String,
auth_header_value: String,
endpoint_tag: 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", "gemini"])),
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", "gemini"])),
Some(serde_json::json!(["openai:chat"])),
Some(serde_json::json!(["gpt-5"])),
)
.expect("auth snapshot should build")
}
fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: "provider-openai-chat-gemini-finalize-local-1".to_string(),
provider_name: "gemini".to_string(),
provider_type: "custom".to_string(),
provider_priority: 10,
provider_is_active: true,
endpoint_id: "endpoint-openai-chat-gemini-finalize-local-1".to_string(),
endpoint_api_format: "gemini:generate_content".to_string(),
endpoint_api_family: Some("gemini".to_string()),
endpoint_kind: Some("chat".to_string()),
endpoint_is_active: true,
key_id: "key-openai-chat-gemini-finalize-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:generate_content".to_string()]),
key_allowed_models: None,
key_capabilities: None,
key_internal_priority: 5,
key_global_priority_by_format: Some(serde_json::json!({"gemini:generate_content": 1})),
model_id: "model-openai-chat-gemini-finalize-local-1".to_string(),
global_model_id: "global-model-openai-chat-gemini-finalize-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: "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:generate_content".to_string()]),
endpoint_ids: None,
operations: None,
}]),
model_supports_streaming: Some(true),
model_is_active: true,
model_is_available: true,
}
}
fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
"provider-openai-chat-gemini-finalize-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,
true,
None,
Some(2),
None,
Some(20.0),
None,
None,
)
}
fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
"endpoint-openai-chat-gemini-finalize-local-1".to_string(),
"provider-openai-chat-gemini-finalize-local-1".to_string(),
"gemini:generate_content".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":"openai-chat-gemini-finalize-cross-format"}
])),
None,
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-openai-chat-gemini-finalize-local-1".to_string(),
"provider-openai-chat-gemini-finalize-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:generate_content"])),
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
"sk-upstream-openai-chat-gemini-finalize",
)
.expect("api key should encrypt"),
None,
None,
Some(serde_json::json!({"gemini:generate_content": 1})),
None,
None,
None,
None,
)
.expect("key transport should build")
}
let seen_remote_execution_runtime =
Arc::new(Mutex::new(None::<SeenRemoteExecutionRuntimeRequest>));
let seen_remote_execution_runtime_clone = Arc::clone(&seen_remote_execution_runtime);
let report_hits = Arc::new(Mutex::new(0usize));
let report_hits_clone = Arc::clone(&report_hits);
let finalize_hits = Arc::new(Mutex::new(0usize));
let finalize_hits_clone = Arc::clone(&finalize_hits);
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 usage_repository = Arc::new(InMemoryUsageReadRepository::default());
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/finalize-sync",
any(move |_request: Request| {
let finalize_hits_inner = Arc::clone(&finalize_hits_clone);
async move {
*finalize_hits_inner.lock().expect("mutex should lock") += 1;
(
StatusCode::IM_A_TEAPOT,
Body::from("finalize-sync-should-not-be-hit"),
)
}
}),
)
.route(
"/api/internal/gateway/report-sync",
any(move |_request: Request| {
let report_hits_inner = Arc::clone(&report_hits_clone);
async move {
*report_hits_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/sync",
any(move |request: Request| {
let seen_remote_execution_runtime_inner = Arc::clone(&seen_remote_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_remote_execution_runtime_inner
.lock()
.expect("mutex should lock") = Some(SeenRemoteExecutionRuntimeRequest {
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(),
provider_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-goog-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(),
});
Json(json!({
"request_id": "trace-openai-chat-xfmt-stream-123",
"status_code": 200,
"headers": {
"content-type": "text/event-stream"
},
"body": {
"body_bytes_b64": base64::engine::general_purpose::STANDARD.encode(
concat!(
"data: {\"responseId\":\"resp-gemini-chat-stream-123\",\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"Hello \"}],\"role\":\"model\"},\"index\":0}],\"modelVersion\":\"gemini-2.5-pro-upstream\"}\n\n",
"data: {\"responseId\":\"resp-gemini-chat-stream-123\",\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"Gemini Chat\"}],\"role\":\"model\"},\"finishReason\":\"STOP\",\"index\":0}],\"modelVersion\":\"gemini-2.5-pro-upstream\",\"usageMetadata\":{\"promptTokenCount\":1,\"candidatesTokenCount\":2,\"totalTokenCount\":3}}\n\n"
)
)
},
"telemetry": {
"elapsed_ms": 33
}
}))
}
}),
);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-client-openai-chat-xfmt-stream")),
sample_auth_snapshot(
"api-key-openai-chat-xfmt-stream-1",
"user-openai-chat-xfmt-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(
GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_and_usage_for_tests(
auth_repository,
candidate_selection_repository,
provider_catalog_repository,
Arc::clone(&request_candidate_repository),
Arc::clone(&usage_repository),
DEVELOPMENT_ENCRYPTION_KEY,
),
);
let gateway = build_router_with_state(gateway_state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let started_at = std::time::Instant::now();
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-chat-xfmt-stream",
)
.header(TRACE_ID_HEADER, "trace-openai-chat-xfmt-stream-123")
.body("{\"model\":\"gpt-5\",\"messages\":[]}")
.send()
.await
.expect("request should succeed");
let elapsed = started_at.elapsed();
let response_status = response.status();
let response_body = response.text().await.expect("body should read");
assert_eq!(response_status, StatusCode::OK);
let response_json: serde_json::Value =
serde_json::from_str(&response_body).expect("body should parse");
assert_eq!(
response_json,
json!({
"id": "resp-gemini-chat-stream-123",
"object": "chat.completion",
"model": "gemini-2.5-pro-upstream",
"choices": [{
"index": 0,
"message": {
"role": "assistant",
"content": "Hello Gemini Chat"
},
"finish_reason": "stop"
}],
"usage": {
"prompt_tokens": 1,
"completion_tokens": 2,
"total_tokens": 3
}
})
);
assert!(
elapsed < std::time::Duration::from_millis(3_500),
"response took unexpectedly long for local finalize path: elapsed={elapsed:?} finalize_hits={} report_hits={}",
*finalize_hits.lock().expect("mutex should lock"),
*report_hits.lock().expect("mutex should lock"),
);
let seen_remote_execution_runtime_request = seen_remote_execution_runtime
.lock()
.expect("mutex should lock")
.clone()
.expect("remote execution runtime plan should be captured");
assert_eq!(
seen_remote_execution_runtime_request.trace_id,
"trace-openai-chat-xfmt-stream-123"
);
assert_eq!(
seen_remote_execution_runtime_request.url,
"https://generativelanguage.googleapis.com/custom/v1beta/models/gemini-2.5-pro-upstream:generateContent"
);
assert_eq!(
seen_remote_execution_runtime_request.provider_model,
"gemini-2.5-pro-upstream"
);
assert_eq!(
seen_remote_execution_runtime_request.auth_header_value,
"sk-upstream-openai-chat-gemini-finalize"
);
assert_eq!(
seen_remote_execution_runtime_request.endpoint_tag,
"openai-chat-gemini-finalize-cross-format"
);
let mut stored_candidates = Vec::new();
for _ in 0..50 {
stored_candidates = request_candidate_repository
.list_by_request_id("trace-openai-chat-xfmt-stream-123")
.await
.expect("request candidate trace should read");
if stored_candidates.len() == 1
&& stored_candidates[0].status == RequestCandidateStatus::Success
{
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
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_eq!(
*report_hits.lock().expect("mutex should lock"),
0,
"report-sync should stay local when request candidate persistence is available"
);
assert_eq!(*finalize_hits.lock().expect("mutex should lock"), 0);
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();
}
#[test]
fn gateway_executes_openai_chat_cross_format_tool_use_upstream_stream_via_local_finalize_response()
{
run_openai_chat_finalize_test(
"gateway_executes_openai_chat_cross_format_tool_use_upstream_stream_via_local_finalize_response",
gateway_executes_openai_chat_cross_format_tool_use_upstream_stream_via_local_finalize_response_impl,
);
}
async fn gateway_executes_openai_chat_cross_format_tool_use_upstream_stream_via_local_finalize_response_impl(
) {
use base64::Engine as _;
#[derive(Debug, Clone)]
struct SeenRemoteExecutionRuntimeRequest {
trace_id: String,
request_id: String,
url: String,
provider_model: String,
auth_header_value: String,
endpoint_tag: 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", "claude"])),
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", "claude"])),
Some(serde_json::json!(["openai:chat"])),
Some(serde_json::json!(["gpt-5"])),
)
.expect("auth snapshot should build")
}
fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: "provider-openai-chat-claude-tool-finalize-local-1".to_string(),
provider_name: "claude".to_string(),
provider_type: "custom".to_string(),
provider_priority: 10,
provider_is_active: true,
endpoint_id: "endpoint-openai-chat-claude-tool-finalize-local-1".to_string(),
endpoint_api_format: "claude:messages".to_string(),
endpoint_api_family: Some("claude".to_string()),
endpoint_kind: Some("chat".to_string()),
endpoint_is_active: true,
key_id: "key-openai-chat-claude-tool-finalize-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:messages".to_string()]),
key_allowed_models: None,
key_capabilities: None,
key_internal_priority: 5,
key_global_priority_by_format: Some(serde_json::json!({"claude:messages": 1})),
model_id: "model-openai-chat-claude-tool-finalize-local-1".to_string(),
global_model_id: "global-model-openai-chat-claude-tool-finalize-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: "claude-sonnet-4-upstream".to_string(),
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
name: "claude-sonnet-4-upstream".to_string(),
priority: 1,
api_formats: Some(vec!["claude:messages".to_string()]),
endpoint_ids: None,
operations: None,
}]),
model_supports_streaming: Some(true),
model_is_active: true,
model_is_available: true,
}
}
fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
"provider-openai-chat-claude-tool-finalize-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,
true,
None,
Some(2),
None,
Some(20.0),
None,
None,
)
}
fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
"endpoint-openai-chat-claude-tool-finalize-local-1".to_string(),
"provider-openai-chat-claude-tool-finalize-local-1".to_string(),
"claude:messages".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":"openai-chat-claude-tool-finalize-cross-format"}
])),
None,
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-openai-chat-claude-tool-finalize-local-1".to_string(),
"provider-openai-chat-claude-tool-finalize-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:messages"])),
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
"sk-upstream-openai-chat-claude-tool-finalize",
)
.expect("api key should encrypt"),
None,
None,
Some(serde_json::json!({"claude:messages": 1})),
None,
None,
None,
None,
)
.expect("key transport should build")
}
let seen_remote_execution_runtime =
Arc::new(Mutex::new(None::<SeenRemoteExecutionRuntimeRequest>));
let seen_remote_execution_runtime_clone = Arc::clone(&seen_remote_execution_runtime);
let report_hits = Arc::new(Mutex::new(0usize));
let report_hits_clone = Arc::clone(&report_hits);
let finalize_hits = Arc::new(Mutex::new(0usize));
let finalize_hits_clone = Arc::clone(&finalize_hits);
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 usage_repository = Arc::new(InMemoryUsageReadRepository::default());
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/finalize-sync",
any(move |_request: Request| {
let finalize_hits_inner = Arc::clone(&finalize_hits_clone);
async move {
*finalize_hits_inner.lock().expect("mutex should lock") += 1;
(
StatusCode::IM_A_TEAPOT,
Body::from("finalize-sync-should-not-be-hit"),
)
}
}),
)
.route(
"/api/internal/gateway/report-sync",
any(move |_request: Request| {
let report_hits_inner = Arc::clone(&report_hits_clone);
async move {
*report_hits_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/sync",
any(move |request: Request| {
let seen_remote_execution_runtime_inner = Arc::clone(&seen_remote_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_remote_execution_runtime_inner
.lock()
.expect("mutex should lock") = Some(SeenRemoteExecutionRuntimeRequest {
trace_id: parts
.headers
.get(TRACE_ID_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
request_id: payload
.get("request_id")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
url: payload
.get("url")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
provider_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(),
});
Json(json!({
"request_id": "trace-openai-chat-xfmt-tool-stream-123",
"status_code": 200,
"headers": {
"content-type": "text/event-stream"
},
"body": {
"body_bytes_b64": base64::engine::general_purpose::STANDARD.encode(
concat!(
"event: message_start\n",
"data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_tool_claude_123\",\"type\":\"message\",\"role\":\"assistant\",\"model\":\"claude-sonnet-4-upstream\",\"content\":[],\"stop_reason\":null,\"stop_sequence\":null}}\n\n",
"event: content_block_start\n",
"data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"Checking.\"}}\n\n",
"event: content_block_stop\n",
"data: {\"type\":\"content_block_stop\",\"index\":0}\n\n",
"event: content_block_start\n",
"data: {\"type\":\"content_block_start\",\"index\":1,\"content_block\":{\"type\":\"tool_use\",\"id\":\"toolu_123\",\"name\":\"get_weather\",\"input\":{\"location\":\"Tokyo\"}}}\n\n",
"event: content_block_stop\n",
"data: {\"type\":\"content_block_stop\",\"index\":1}\n\n",
"event: message_delta\n",
"data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"tool_use\"},\"usage\":{\"input_tokens\":5,\"output_tokens\":7}}\n\n",
"event: message_stop\n",
"data: {\"type\":\"message_stop\"}\n\n"
)
)
},
"telemetry": {
"elapsed_ms": 31
}
}))
}
}),
);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-client-openai-chat-xfmt-tool-stream")),
sample_auth_snapshot(
"api-key-openai-chat-xfmt-tool-stream-1",
"user-openai-chat-xfmt-tool-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(
GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_and_usage_for_tests(
auth_repository,
candidate_selection_repository,
provider_catalog_repository,
Arc::clone(&request_candidate_repository),
Arc::clone(&usage_repository),
DEVELOPMENT_ENCRYPTION_KEY,
),
);
let gateway = build_router_with_state(gateway_state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let started_at = std::time::Instant::now();
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-chat-xfmt-tool-stream",
)
.header(TRACE_ID_HEADER, "trace-openai-chat-xfmt-tool-stream-123")
.body("{\"model\":\"gpt-5\",\"messages\":[{\"role\":\"user\",\"content\":\"weather\"}]}")
.send()
.await
.expect("request should succeed");
let elapsed = started_at.elapsed();
let response_status = response.status();
let response_body = response.text().await.expect("body should read");
assert_eq!(response_status, StatusCode::OK);
let response_json: serde_json::Value =
serde_json::from_str(&response_body).expect("body should parse");
assert_eq!(
response_json,
json!({
"id": "msg_tool_claude_123",
"object": "chat.completion",
"model": "claude-sonnet-4-upstream",
"choices": [{
"index": 0,
"message": {
"role": "assistant",
"content": "Checking.",
"tool_calls": [{
"id": "toolu_123",
"type": "function",
"function": {
"name": "get_weather",
"arguments": "{\"location\":\"Tokyo\"}"
}
}]
},
"finish_reason": "tool_calls"
}],
"usage": {
"prompt_tokens": 5,
"completion_tokens": 7,
"total_tokens": 12
}
})
);
assert!(
elapsed < std::time::Duration::from_millis(3_500),
"response took unexpectedly long for local finalize path: elapsed={elapsed:?} finalize_hits={} report_hits={}",
*finalize_hits.lock().expect("mutex should lock"),
*report_hits.lock().expect("mutex should lock"),
);
let seen_remote_execution_runtime_request = seen_remote_execution_runtime
.lock()
.expect("mutex should lock")
.clone()
.expect("remote execution runtime plan should be captured");
assert_eq!(
seen_remote_execution_runtime_request.trace_id,
"trace-openai-chat-xfmt-tool-stream-123"
);
assert_eq!(
seen_remote_execution_runtime_request.request_id,
"trace-openai-chat-xfmt-tool-stream-123"
);
assert_eq!(
seen_remote_execution_runtime_request.url,
"https://api.anthropic.example/custom/v1/messages"
);
assert_eq!(
seen_remote_execution_runtime_request.provider_model,
"claude-sonnet-4-upstream"
);
assert_eq!(
seen_remote_execution_runtime_request.auth_header_value,
"sk-upstream-openai-chat-claude-tool-finalize"
);
assert_eq!(
seen_remote_execution_runtime_request.endpoint_tag,
"openai-chat-claude-tool-finalize-cross-format"
);
let mut stored_candidates = Vec::new();
for _ in 0..50 {
stored_candidates = request_candidate_repository
.list_by_request_id("trace-openai-chat-xfmt-tool-stream-123")
.await
.expect("request candidate trace should read");
if stored_candidates.len() == 1
&& stored_candidates[0].status == RequestCandidateStatus::Success
{
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
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_eq!(
*report_hits.lock().expect("mutex should lock"),
0,
"report-sync should stay local when request candidate persistence is available"
);
assert_eq!(*finalize_hits.lock().expect("mutex should lock"), 0);
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();
}
#[test]
fn gateway_executes_openai_chat_antigravity_cross_format_sync_via_local_finalize_response() {
run_openai_chat_finalize_test(
"gateway_executes_openai_chat_antigravity_cross_format_sync_via_local_finalize_response",
gateway_executes_openai_chat_antigravity_cross_format_sync_via_local_finalize_response_impl,
);
}
async fn gateway_executes_openai_chat_antigravity_cross_format_sync_via_local_finalize_response_impl(
) {
use base64::Engine as _;
#[derive(Debug, Clone)]
struct SeenExecutionRuntimeSyncRequest {
trace_id: String,
url: String,
accept: String,
authorization: String,
x_client_name: String,
x_client_version: String,
x_vscode_sessionid: String,
x_goog_api_client: String,
project: String,
request_id: String,
model: String,
user_agent: String,
request_type: String,
contents_len: usize,
request_has_model: bool,
}
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", "antigravity"])),
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_i64),
Some(serde_json::json!(["openai", "antigravity"])),
Some(serde_json::json!(["openai:chat"])),
Some(serde_json::json!(["gpt-5"])),
)
.expect("auth snapshot should build")
}
fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: "provider-openai-chat-antigravity-finalize-local-1".to_string(),
provider_name: "antigravity".to_string(),
provider_type: "antigravity".to_string(),
provider_priority: 10,
provider_is_active: true,
endpoint_id: "endpoint-openai-chat-antigravity-finalize-local-1".to_string(),
endpoint_api_format: "gemini:generate_content".to_string(),
endpoint_api_family: Some("gemini".to_string()),
endpoint_kind: Some("chat".to_string()),
endpoint_is_active: true,
key_id: "key-openai-chat-antigravity-finalize-local-1".to_string(),
key_name: "prod".to_string(),
key_auth_type: "oauth".to_string(),
key_is_active: true,
key_api_formats: Some(vec!["gemini:generate_content".to_string()]),
key_allowed_models: None,
key_capabilities: None,
key_internal_priority: 5,
key_global_priority_by_format: Some(serde_json::json!({"gemini:generate_content": 1})),
model_id: "model-openai-chat-antigravity-finalize-local-1".to_string(),
global_model_id: "global-model-openai-chat-antigravity-finalize-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: "claude-sonnet-4-5".to_string(),
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
name: "claude-sonnet-4-5".to_string(),
priority: 1,
api_formats: Some(vec!["gemini:generate_content".to_string()]),
endpoint_ids: None,
operations: None,
}]),
model_supports_streaming: Some(true),
model_is_active: true,
model_is_available: true,
}
}
fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
"provider-openai-chat-antigravity-finalize-local-1".to_string(),
"antigravity".to_string(),
Some("https://example.com".to_string()),
"antigravity".to_string(),
)
.expect("provider should build")
.with_transport_fields(
true,
false,
true,
None,
Some(2),
None,
Some(20.0),
None,
None,
)
}
fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
"endpoint-openai-chat-antigravity-finalize-local-1".to_string(),
"provider-openai-chat-antigravity-finalize-local-1".to_string(),
"gemini:generate_content".to_string(),
Some("gemini".to_string()),
Some("chat".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
"https://antigravity.googleapis.com".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":"antigravity","project_id":"project-antigravity-chat-local-1","client_version":"1.2.3","session_id":"sess-antigravity-chat-local-123","access_token_import_temporary":true,"headers":{"Authorization":"Bearer imported-antigravity-chat-token"}}"#,
)
.expect("auth config should encrypt");
StoredProviderCatalogKey::new(
"key-openai-chat-antigravity-finalize-local-1".to_string(),
"provider-openai-chat-antigravity-finalize-local-1".to_string(),
"prod".to_string(),
"oauth".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(serde_json::json!(["gemini:generate_content"])),
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "__placeholder__")
.expect("api key should encrypt"),
Some(encrypted_auth_config),
None,
Some(serde_json::json!({"gemini:generate_content": 1})),
None,
None,
None,
None,
)
.expect("key transport should build")
}
let report_hits = Arc::new(Mutex::new(0usize));
let report_hits_clone = Arc::clone(&report_hits);
let finalize_hits = Arc::new(Mutex::new(0usize));
let finalize_hits_clone = Arc::clone(&finalize_hits);
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 seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeSyncRequest>));
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
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/finalize-sync",
any(move |_request: Request| {
let finalize_hits_inner = Arc::clone(&finalize_hits_clone);
async move {
*finalize_hits_inner.lock().expect("mutex should lock") += 1;
(
StatusCode::IM_A_TEAPOT,
Body::from("finalize-sync-should-not-be-hit"),
)
}
}),
)
.route(
"/api/internal/gateway/report-sync",
any(move |_request: Request| {
let report_hits_inner = Arc::clone(&report_hits_clone);
async move {
*report_hits_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/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(),
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(),
x_client_name: payload
.get("headers")
.and_then(|value| value.get("x-client-name"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
x_client_version: payload
.get("headers")
.and_then(|value| value.get("x-client-version"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
x_vscode_sessionid: payload
.get("headers")
.and_then(|value| value.get("x-vscode-sessionid"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
x_goog_api_client: payload
.get("headers")
.and_then(|value| value.get("x-goog-api-client"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
project: payload
.get("body")
.and_then(|value| value.get("json_body"))
.and_then(|value| value.get("project"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
request_id: payload
.get("body")
.and_then(|value| value.get("json_body"))
.and_then(|value| value.get("requestId"))
.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(),
user_agent: payload
.get("body")
.and_then(|value| value.get("json_body"))
.and_then(|value| value.get("userAgent"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
request_type: payload
.get("body")
.and_then(|value| value.get("json_body"))
.and_then(|value| value.get("requestType"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
contents_len: payload
.get("body")
.and_then(|value| value.get("json_body"))
.and_then(|value| value.get("request"))
.and_then(|value| value.get("contents"))
.and_then(|value| value.as_array())
.map(Vec::len)
.unwrap_or_default(),
request_has_model: payload
.get("body")
.and_then(|value| value.get("json_body"))
.and_then(|value| value.get("request"))
.and_then(|value| value.get("model"))
.is_some(),
});
Json(json!({
"request_id": "trace-openai-chat-antigravity-sync-123",
"status_code": 200,
"headers": {
"content-type": "text/event-stream"
},
"body": {
"body_bytes_b64": base64::engine::general_purpose::STANDARD.encode(
concat!(
"data: {\"response\":{\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"Hello Antigravity Chat\"}],\"role\":\"model\"},\"finishReason\":\"STOP\",\"index\":0}],\"modelVersion\":\"claude-sonnet-4-5\",\"usageMetadata\":{\"promptTokenCount\":2,\"candidatesTokenCount\":3,\"totalTokenCount\":5}},\"responseId\":\"resp-antigravity-chat-sync-123\"}\n\n"
)
)
},
"telemetry": {
"elapsed_ms": 29
}
}))
}
}),
);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-client-openai-chat-antigravity-sync")),
sample_auth_snapshot(
"api-key-openai-chat-antigravity-sync-1",
"user-openai-chat-antigravity-sync-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(
GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_and_usage_for_tests(
auth_repository,
candidate_selection_repository,
provider_catalog_repository,
Arc::clone(&request_candidate_repository),
Arc::clone(&usage_repository),
DEVELOPMENT_ENCRYPTION_KEY,
),
);
let gateway = build_router_with_state(gateway_state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let started_at = std::time::Instant::now();
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-chat-antigravity-sync",
)
.header(TRACE_ID_HEADER, "trace-openai-chat-antigravity-sync-123")
.body("{\"model\":\"gpt-5\",\"messages\":[{\"role\":\"user\",\"content\":\"weather\"}]}")
.send()
.await
.expect("request should succeed");
let elapsed = started_at.elapsed();
let response_status = response.status();
let response_body = response.text().await.expect("body should read");
assert_eq!(response_status, StatusCode::OK);
let response_json: serde_json::Value =
serde_json::from_str(&response_body).expect("body should parse");
assert_eq!(
response_json,
json!({
"id": "resp-local-stream",
"object": "chat.completion",
"model": "claude-sonnet-4-5",
"choices": [{
"index": 0,
"message": {
"role": "assistant",
"content": "Hello Antigravity Chat"
},
"finish_reason": "stop"
}],
"usage": {
"prompt_tokens": 2,
"completion_tokens": 3,
"total_tokens": 5
}
})
);
assert!(
elapsed < std::time::Duration::from_millis(3_500),
"response took unexpectedly long for local finalize path: elapsed={elapsed:?} finalize_hits={} report_hits={}",
*finalize_hits.lock().expect("mutex should lock"),
*report_hits.lock().expect("mutex should lock"),
);
let seen_execution_runtime_request = seen_execution_runtime
.lock()
.expect("mutex should lock")
.clone()
.expect("execution runtime request should be captured");
assert_eq!(
seen_execution_runtime_request.trace_id,
"trace-openai-chat-antigravity-sync-123"
);
assert_eq!(
seen_execution_runtime_request.url,
"https://antigravity.googleapis.com/v1internal:streamGenerateContent?alt=sse"
);
assert_eq!(seen_execution_runtime_request.accept, "text/event-stream");
assert_eq!(
seen_execution_runtime_request.authorization,
"Bearer imported-antigravity-chat-token"
);
assert_eq!(seen_execution_runtime_request.x_client_name, "antigravity");
assert_eq!(seen_execution_runtime_request.x_client_version, "1.2.3");
assert_eq!(
seen_execution_runtime_request.x_vscode_sessionid,
"sess-antigravity-chat-local-123"
);
assert_eq!(
seen_execution_runtime_request.x_goog_api_client,
"gl-node/18.18.2 fire/0.8.6 grpc/1.10.x"
);
assert_eq!(
seen_execution_runtime_request.project,
"project-antigravity-chat-local-1"
);
assert_eq!(
seen_execution_runtime_request.request_id,
"trace-openai-chat-antigravity-sync-123"
);
assert_eq!(seen_execution_runtime_request.model, "claude-sonnet-4-5");
assert_eq!(
seen_execution_runtime_request.user_agent,
aether_provider_transport::antigravity::ANTIGRAVITY_REQUEST_USER_AGENT
);
assert_eq!(seen_execution_runtime_request.request_type, "agent");
assert_eq!(seen_execution_runtime_request.contents_len, 1);
assert!(!seen_execution_runtime_request.request_has_model);
let mut stored_candidates = Vec::new();
for _ in 0..50 {
stored_candidates = request_candidate_repository
.list_by_request_id("trace-openai-chat-antigravity-sync-123")
.await
.expect("request candidate trace should read");
if stored_candidates.len() == 1
&& stored_candidates[0].status == RequestCandidateStatus::Success
{
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
assert_eq!(stored_candidates.len(), 1);
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success);
assert_eq!(stored_candidates[0].skip_reason.as_deref(), None);
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
assert_eq!(*report_hits.lock().expect("mutex should lock"), 0);
assert_eq!(*finalize_hits.lock().expect("mutex should lock"), 0);
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();
}
#[test]
fn gateway_executes_openai_chat_cross_format_claude_upstream_sync_via_local_finalize_response() {
run_openai_chat_finalize_test(
"gateway_executes_openai_chat_cross_format_claude_upstream_sync_via_local_finalize_response",
gateway_executes_openai_chat_cross_format_claude_upstream_sync_via_local_finalize_response_impl,
);
}
async fn gateway_executes_openai_chat_cross_format_claude_upstream_sync_via_local_finalize_response_impl(
) {
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", "claude"])),
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", "claude"])),
Some(serde_json::json!(["openai:chat"])),
Some(serde_json::json!(["gpt-5"])),
)
.expect("auth snapshot should build")
}
fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: "provider-openai-chat-claude-direct-sync-1".to_string(),
provider_name: "claude".to_string(),
provider_type: "custom".to_string(),
provider_priority: 10,
provider_is_active: true,
endpoint_id: "endpoint-openai-chat-claude-direct-sync-1".to_string(),
endpoint_api_format: "claude:messages".to_string(),
endpoint_api_family: Some("claude".to_string()),
endpoint_kind: Some("chat".to_string()),
endpoint_is_active: true,
key_id: "key-openai-chat-claude-direct-sync-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:messages".to_string()]),
key_allowed_models: None,
key_capabilities: None,
key_internal_priority: 5,
key_global_priority_by_format: Some(serde_json::json!({"claude:messages": 1})),
model_id: "model-openai-chat-claude-direct-sync-1".to_string(),
global_model_id: "global-model-openai-chat-claude-direct-sync-1".to_string(),
global_model_name: "gpt-5".to_string(),
global_model_mappings: None,
global_model_supports_streaming: Some(true),
model_provider_model_name: "claude-sonnet-4-upstream".to_string(),
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
name: "claude-sonnet-4-upstream".to_string(),
priority: 1,
api_formats: Some(vec!["claude:messages".to_string()]),
endpoint_ids: None,
operations: None,
}]),
model_supports_streaming: Some(true),
model_is_active: true,
model_is_available: true,
}
}
fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
"provider-openai-chat-claude-direct-sync-1".to_string(),
"claude".to_string(),
Some("https://example.com".to_string()),
"custom".to_string(),
)
.expect("provider should build")
.with_transport_fields(
true,
false,
true,
None,
Some(2),
None,
Some(20.0),
None,
None,
)
}
fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
"endpoint-openai-chat-claude-direct-sync-1".to_string(),
"provider-openai-chat-claude-direct-sync-1".to_string(),
"claude:messages".to_string(),
Some("claude".to_string()),
Some("chat".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
"https://api.anthropic.example".to_string(),
None,
None,
Some(2),
Some("/v1/messages".to_string()),
None,
None,
None,
)
.expect("endpoint transport should build")
}
fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
StoredProviderCatalogKey::new(
"key-openai-chat-claude-direct-sync-1".to_string(),
"provider-openai-chat-claude-direct-sync-1".to_string(),
"prod".to_string(),
"api_key".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(serde_json::json!(["claude:messages"])),
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
"sk-upstream-openai-chat-claude-direct-sync",
)
.expect("api key should encrypt"),
None,
None,
Some(serde_json::json!({"claude:messages": 1})),
None,
None,
None,
None,
)
.expect("key transport should build")
}
let report_hits = Arc::new(Mutex::new(0usize));
let report_hits_clone = Arc::clone(&report_hits);
let finalize_hits = Arc::new(Mutex::new(0usize));
let finalize_hits_clone = Arc::clone(&finalize_hits);
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 usage_repository = Arc::new(InMemoryUsageReadRepository::default());
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/finalize-sync",
any(move |_request: Request| {
let finalize_hits_inner = Arc::clone(&finalize_hits_clone);
async move {
*finalize_hits_inner.lock().expect("mutex should lock") += 1;
(
StatusCode::IM_A_TEAPOT,
Body::from("finalize-sync-should-not-be-hit"),
)
}
}),
)
.route(
"/api/internal/gateway/report-sync",
any(move |_request: Request| {
let report_hits_inner = Arc::clone(&report_hits_clone);
async move {
*report_hits_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/sync",
any(|_request: Request| async move {
Json(json!({
"request_id": "trace-openai-chat-claude-direct-sync-123",
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"id": "msg_claude_direct_sync_123",
"type": "message",
"model": "claude-sonnet-4-upstream",
"role": "assistant",
"content": [{"type": "text", "text": "Hello Claude direct"}],
"stop_reason": "end_turn",
"usage": {
"input_tokens": 2,
"output_tokens": 3
}
}
},
"telemetry": {
"elapsed_ms": 27
}
}))
}),
);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-client-openai-chat-claude-direct-sync")),
sample_auth_snapshot(
"api-key-openai-chat-claude-direct-sync-1",
"user-openai-chat-claude-direct-sync-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(
GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_and_usage_for_tests(
auth_repository,
candidate_selection_repository,
provider_catalog_repository,
Arc::clone(&request_candidate_repository),
Arc::clone(&usage_repository),
DEVELOPMENT_ENCRYPTION_KEY,
),
);
let gateway = build_router_with_state(gateway_state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let started_at = std::time::Instant::now();
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-chat-claude-direct-sync",
)
.header(TRACE_ID_HEADER, "trace-openai-chat-claude-direct-sync-123")
.body("{\"model\":\"gpt-5\",\"messages\":[]}")
.send()
.await
.expect("request should succeed");
let elapsed = started_at.elapsed();
let response_status = response.status();
let response_body = response.text().await.expect("body should read");
assert_eq!(response_status, StatusCode::OK);
let response_json: serde_json::Value =
serde_json::from_str(&response_body).expect("body should parse");
assert_eq!(
response_json,
json!({
"id": "msg_claude_direct_sync_123",
"object": "chat.completion",
"model": "claude-sonnet-4-upstream",
"choices": [{
"index": 0,
"message": {
"role": "assistant",
"content": "Hello Claude direct"
},
"finish_reason": "stop"
}],
"usage": {
"prompt_tokens": 2,
"completion_tokens": 3,
"total_tokens": 5
}
})
);
assert!(
elapsed < std::time::Duration::from_millis(3_500),
"response took unexpectedly long for local finalize path"
);
let mut stored_candidates = Vec::new();
for _ in 0..50 {
stored_candidates = request_candidate_repository
.list_by_request_id("trace-openai-chat-claude-direct-sync-123")
.await
.expect("request candidate trace should read");
if stored_candidates.len() == 1
&& stored_candidates[0].status == RequestCandidateStatus::Success
{
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
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_eq!(*report_hits.lock().expect("mutex should lock"), 0);
assert_eq!(*finalize_hits.lock().expect("mutex should lock"), 0);
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();
}
#[test]
fn gateway_executes_openai_chat_cross_format_gemini_upstream_sync_via_local_finalize_response() {
run_openai_chat_finalize_test(
"gateway_executes_openai_chat_cross_format_gemini_upstream_sync_via_local_finalize_response",
gateway_executes_openai_chat_cross_format_gemini_upstream_sync_via_local_finalize_response_impl,
);
}
async fn gateway_executes_openai_chat_cross_format_gemini_upstream_sync_via_local_finalize_response_impl(
) {
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", "gemini"])),
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", "gemini"])),
Some(serde_json::json!(["openai:chat"])),
Some(serde_json::json!(["gpt-5"])),
)
.expect("auth snapshot should build")
}
fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: "provider-openai-chat-gemini-direct-sync-1".to_string(),
provider_name: "gemini".to_string(),
provider_type: "custom".to_string(),
provider_priority: 10,
provider_is_active: true,
endpoint_id: "endpoint-openai-chat-gemini-direct-sync-1".to_string(),
endpoint_api_format: "gemini:generate_content".to_string(),
endpoint_api_family: Some("gemini".to_string()),
endpoint_kind: Some("chat".to_string()),
endpoint_is_active: true,
key_id: "key-openai-chat-gemini-direct-sync-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:generate_content".to_string()]),
key_allowed_models: None,
key_capabilities: None,
key_internal_priority: 5,
key_global_priority_by_format: Some(serde_json::json!({"gemini:generate_content": 1})),
model_id: "model-openai-chat-gemini-direct-sync-1".to_string(),
global_model_id: "global-model-openai-chat-gemini-direct-sync-1".to_string(),
global_model_name: "gpt-5".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:generate_content".to_string()]),
endpoint_ids: None,
operations: None,
}]),
model_supports_streaming: Some(true),
model_is_active: true,
model_is_available: true,
}
}
fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
"provider-openai-chat-gemini-direct-sync-1".to_string(),
"gemini".to_string(),
Some("https://example.com".to_string()),
"custom".to_string(),
)
.expect("provider should build")
.with_transport_fields(
true,
false,
true,
None,
Some(2),
None,
Some(20.0),
None,
None,
)
}
fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
"endpoint-openai-chat-gemini-direct-sync-1".to_string(),
"provider-openai-chat-gemini-direct-sync-1".to_string(),
"gemini:generate_content".to_string(),
Some("gemini".to_string()),
Some("chat".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
"https://generativelanguage.googleapis.com".to_string(),
None,
None,
Some(2),
Some("/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-openai-chat-gemini-direct-sync-1".to_string(),
"provider-openai-chat-gemini-direct-sync-1".to_string(),
"prod".to_string(),
"api_key".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(serde_json::json!(["gemini:generate_content"])),
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
"sk-upstream-openai-chat-gemini-direct-sync",
)
.expect("api key should encrypt"),
None,
None,
Some(serde_json::json!({"gemini:generate_content": 1})),
None,
None,
None,
None,
)
.expect("key transport should build")
}
let report_hits = Arc::new(Mutex::new(0usize));
let report_hits_clone = Arc::clone(&report_hits);
let finalize_hits = Arc::new(Mutex::new(0usize));
let finalize_hits_clone = Arc::clone(&finalize_hits);
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 usage_repository = Arc::new(InMemoryUsageReadRepository::default());
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/finalize-sync",
any(move |_request: Request| {
let finalize_hits_inner = Arc::clone(&finalize_hits_clone);
async move {
*finalize_hits_inner.lock().expect("mutex should lock") += 1;
(
StatusCode::IM_A_TEAPOT,
Body::from("finalize-sync-should-not-be-hit"),
)
}
}),
)
.route(
"/api/internal/gateway/report-sync",
any(move |_request: Request| {
let report_hits_inner = Arc::clone(&report_hits_clone);
async move {
*report_hits_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/sync",
any(|_request: Request| async move {
Json(json!({
"request_id": "trace-openai-chat-gemini-direct-sync-123",
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"responseId": "resp_gemini_direct_sync_123",
"candidates": [{
"content": {
"parts": [{"text": "Hello Gemini direct"}],
"role": "model"
},
"finishReason": "STOP",
"index": 0
}],
"modelVersion": "gemini-2.5-pro-upstream",
"usageMetadata": {
"promptTokenCount": 1,
"candidatesTokenCount": 2,
"totalTokenCount": 3
}
}
},
"telemetry": {
"elapsed_ms": 25
}
}))
}),
);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-client-openai-chat-gemini-direct-sync")),
sample_auth_snapshot(
"api-key-openai-chat-gemini-direct-sync-1",
"user-openai-chat-gemini-direct-sync-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(
GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_and_usage_for_tests(
auth_repository,
candidate_selection_repository,
provider_catalog_repository,
Arc::clone(&request_candidate_repository),
Arc::clone(&usage_repository),
DEVELOPMENT_ENCRYPTION_KEY,
),
);
let gateway = build_router_with_state(gateway_state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let started_at = std::time::Instant::now();
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-chat-gemini-direct-sync",
)
.header(TRACE_ID_HEADER, "trace-openai-chat-gemini-direct-sync-123")
.body("{\"model\":\"gpt-5\",\"messages\":[]}")
.send()
.await
.expect("request should succeed");
let elapsed = started_at.elapsed();
assert_eq!(response.status(), StatusCode::OK);
let response_json: serde_json::Value = response.json().await.expect("body should parse");
assert_eq!(
response_json,
json!({
"id": "resp_gemini_direct_sync_123",
"object": "chat.completion",
"model": "gemini-2.5-pro-upstream",
"choices": [{
"index": 0,
"message": {
"role": "assistant",
"content": "Hello Gemini direct"
},
"finish_reason": "stop"
}],
"usage": {
"prompt_tokens": 1,
"completion_tokens": 2,
"total_tokens": 3
}
})
);
assert!(
elapsed < std::time::Duration::from_millis(3_500),
"response took unexpectedly long for local finalize path"
);
let mut stored_candidates = Vec::new();
for _ in 0..50 {
stored_candidates = request_candidate_repository
.list_by_request_id("trace-openai-chat-gemini-direct-sync-123")
.await
.expect("request candidate trace should read");
if stored_candidates.len() == 1
&& stored_candidates[0].status == RequestCandidateStatus::Success
{
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
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_eq!(*report_hits.lock().expect("mutex should lock"), 0);
assert_eq!(*finalize_hits.lock().expect("mutex should lock"), 0);
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();
}
Reference in New Issue View Git Blame Copy Permalink