Files
Aether/apps/aether-gateway/src/tests/control/admin/provider_query.rs

5234 lines
197 KiB
Rust
Raw Normal View History

use std::sync::{Arc, Mutex};
use aether_contracts::ExecutionPlan;
use aether_crypto::DEVELOPMENT_ENCRYPTION_KEY;
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
use aether_data::repository::global_models::InMemoryGlobalModelReadRepository;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data_contracts::repository::candidates::{
RequestCandidateReadRepository, RequestCandidateStatus,
};
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogEndpoint;
use axum::body::Body;
use axum::routing::any;
use axum::{extract::Request, Json, Router};
use base64::Engine as _;
use http::StatusCode;
use serde_json::json;
use super::super::{
build_router_with_state, build_state_with_execution_runtime_override,
sample_admin_provider_model, sample_endpoint, sample_key, sample_provider, start_server,
AppState,
};
use crate::constants::{
GATEWAY_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER, TRUSTED_ADMIN_USER_ID_HEADER,
TRUSTED_ADMIN_USER_ROLE_HEADER,
};
use crate::data::GatewayDataState;
fn crc32(data: &[u8]) -> u32 {
let mut crc = 0xffff_ffffu32;
for &byte in data {
crc ^= byte as u32;
for _ in 0..8 {
let mask = if crc & 1 == 1 { 0xedb8_8320 } else { 0 };
crc = (crc >> 1) ^ mask;
}
}
!crc
}
fn encode_string_header(name: &str, value: &str) -> Vec<u8> {
let mut out = Vec::new();
out.push(name.len() as u8);
out.extend_from_slice(name.as_bytes());
out.push(7);
out.extend_from_slice(&(value.len() as u16).to_be_bytes());
out.extend_from_slice(value.as_bytes());
out
}
fn encode_frame(headers: Vec<u8>, payload: Vec<u8>) -> Vec<u8> {
let total_len = 12 + headers.len() + payload.len() + 4;
let header_len = headers.len();
let mut out = Vec::with_capacity(total_len);
out.extend_from_slice(&(total_len as u32).to_be_bytes());
out.extend_from_slice(&(header_len as u32).to_be_bytes());
let prelude_crc = crc32(&out[..8]);
out.extend_from_slice(&prelude_crc.to_be_bytes());
out.extend_from_slice(&headers);
out.extend_from_slice(&payload);
let message_crc = crc32(&out);
out.extend_from_slice(&message_crc.to_be_bytes());
out
}
fn encode_kiro_event_frame(event_type: &str, payload: serde_json::Value) -> Vec<u8> {
let mut headers = encode_string_header(":message-type", "event");
headers.extend_from_slice(&encode_string_header(":event-type", event_type));
let payload = serde_json::to_vec(&payload).expect("payload should encode");
encode_frame(headers, payload)
}
fn encode_kiro_exception_frame(exception_type: &str) -> Vec<u8> {
let mut headers = encode_string_header(":message-type", "exception");
headers.extend_from_slice(&encode_string_header(":exception-type", exception_type));
encode_frame(headers, Vec::new())
}
async fn assert_admin_provider_query_route(
path: &str,
request_payload: serde_json::Value,
expected_status: StatusCode,
expected_payload_assertions: impl FnOnce(&serde_json::Value),
) {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
path,
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}{path}"))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&request_payload)
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), expected_status);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
expected_payload_assertions(&payload);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_provider_query_models_fetches_upstream_for_selected_key() {
let execution_runtime_hits = Arc::new(Mutex::new(0usize));
let execution_runtime_hits_clone = Arc::clone(&execution_runtime_hits);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| {
let execution_runtime_hits_inner = Arc::clone(&execution_runtime_hits_clone);
async move {
*execution_runtime_hits_inner
.lock()
.expect("mutex should lock") += 1;
assert_eq!(plan.url, "https://api.openai.example/v1/models");
assert_eq!(
plan.headers.get("authorization").map(String::as_str),
Some("Bearer sk-test")
);
Json(json!({
"request_id": "req-provider-query-selected",
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"data": [{
"id": "LLM-Research/Llama-4-Maverick-17B-128E-Instruct",
"object": "",
"owned_by": "system",
"created": 1732517497u64
}]
}
}
}))
}
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-openai", "OpenAI", 10);
provider.provider_type = "openai".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![StoredProviderCatalogEndpoint::new(
"endpoint-openai-chat".to_string(),
"provider-openai".to_string(),
"openai:chat".to_string(),
Some("chat".to_string()),
Some("primary".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
"https://api.openai.example".to_string(),
None,
None,
None,
None,
None,
None,
None,
)
.expect("endpoint transport should build")],
vec![sample_key(
"key-openai-selected",
"provider-openai",
"openai:chat",
"sk-test",
)],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/models"))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-openai",
"api_key_id": "key-openai-selected"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(payload["provider"]["id"], "provider-openai");
assert_eq!(payload["provider"]["name"], "OpenAI");
assert_eq!(payload["provider"]["display_name"], "OpenAI");
assert_eq!(payload["data"]["error"], serde_json::Value::Null);
assert_eq!(payload["data"]["from_cache"], json!(false));
assert_eq!(payload["data"]["keys_total"], serde_json::Value::Null);
let models = payload["data"]["models"]
.as_array()
.expect("models should be an array");
assert_eq!(models.len(), 1);
assert_eq!(
models[0]["id"],
json!("LLM-Research/Llama-4-Maverick-17B-128E-Instruct")
);
assert_eq!(models[0]["owned_by"], json!("system"));
assert_eq!(models[0]["api_formats"], json!(["openai:chat"]));
assert_eq!(
*execution_runtime_hits.lock().expect("mutex should lock"),
1
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_provider_query_models_with_openai_responses_endpoint() {
let execution_runtime_hits = Arc::new(Mutex::new(0usize));
let execution_runtime_hits_clone = Arc::clone(&execution_runtime_hits);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| {
let execution_runtime_hits_inner = Arc::clone(&execution_runtime_hits_clone);
async move {
*execution_runtime_hits_inner
.lock()
.expect("mutex should lock") += 1;
assert_eq!(plan.endpoint_id, "endpoint-openai-responses");
assert_eq!(plan.provider_api_format, "openai:responses");
Json(json!({
"request_id": "req-provider-query-responses",
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"data": [{
"id": "gpt-4.1",
"object": "model",
"owned_by": "system",
"created": 1732517497u64
}]
}
}
}))
}
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-openai", "OpenAI", 10);
provider.provider_type = "openai".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![StoredProviderCatalogEndpoint::new(
"endpoint-openai-responses".to_string(),
"provider-openai".to_string(),
"openai:responses".to_string(),
Some("responses".to_string()),
Some("primary".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
"https://api.openai.example".to_string(),
None,
None,
None,
None,
None,
None,
None,
)
.expect("endpoint transport should build")],
vec![sample_key(
"key-openai-responses",
"provider-openai",
"openai:responses",
"sk-test-responses",
)],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/models"))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-openai",
"api_key_id": "key-openai-responses"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(payload["data"]["error"], serde_json::Value::Null);
assert_eq!(payload["data"]["from_cache"], json!(false));
assert_eq!(payload["data"]["models"][0]["id"], json!("gpt-4.1"));
assert_eq!(
payload["data"]["models"][0]["api_formats"],
json!(["openai:responses"])
);
assert_eq!(
*execution_runtime_hits.lock().expect("mutex should lock"),
1
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
2026-05-04 01:16:56 +08:00
#[tokio::test]
async fn gateway_handles_admin_provider_query_models_falls_back_to_codex_preset_when_token_invalidated(
) {
let execution_runtime_hits = Arc::new(Mutex::new(0usize));
let execution_runtime_hits_clone = Arc::clone(&execution_runtime_hits);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| {
let execution_runtime_hits_inner = Arc::clone(&execution_runtime_hits_clone);
async move {
*execution_runtime_hits_inner
.lock()
.expect("mutex should lock") += 1;
assert_eq!(
plan.url,
"https://chatgpt.com/backend-api/codex/models?client_version=0.128.0-alpha.1"
);
Json(json!({
"request_id": "req-provider-query-codex-invalidated",
"status_code": 403,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"error": {
"message": "Your authentication token has been invalidated. Please sign in again."
}
}
}
}))
}
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-codex", "Codex", 10);
provider.provider_type = "codex".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-codex-responses",
"provider-codex",
"openai:responses",
"https://chatgpt.com/backend-api/codex",
)],
vec![sample_key(
"key-codex-invalidated",
"provider-codex",
"openai:responses",
"invalidated-token",
)],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/models"))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-codex",
"api_key_id": "key-codex-invalidated"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(payload["data"]["error"], serde_json::Value::Null);
let model_ids = payload["data"]["models"]
.as_array()
.expect("models should be an array")
.iter()
.map(|model| model["id"].as_str().expect("model id"))
.collect::<Vec<_>>();
assert_eq!(
model_ids,
vec![
"gpt-5.3-codex",
"gpt-5.3-codex-spark",
"gpt-5.4",
"gpt-5.4-mini",
"gpt-5.5",
]
);
assert_eq!(
*execution_runtime_hits.lock().expect("mutex should lock"),
1
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_provider_query_models_respecting_key_api_formats() {
let execution_runtime_hits = Arc::new(Mutex::new(0usize));
let execution_runtime_hits_clone = Arc::clone(&execution_runtime_hits);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| {
let execution_runtime_hits_inner = Arc::clone(&execution_runtime_hits_clone);
async move {
*execution_runtime_hits_inner
.lock()
.expect("mutex should lock") += 1;
assert_eq!(plan.endpoint_id, "endpoint-openai-cli");
assert_eq!(plan.provider_api_format, "openai:responses");
Json(json!({
"request_id": "req-provider-query-cli",
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"data": [{
"id": "gpt-5-cli",
"object": "model",
"owned_by": "system",
"created": 1732517497u64
}]
}
}
}))
}
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-openai", "OpenAI", 10);
provider.provider_type = "openai".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![
StoredProviderCatalogEndpoint::new(
"endpoint-openai-chat".to_string(),
"provider-openai".to_string(),
"openai:chat".to_string(),
Some("chat".to_string()),
Some("primary".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
"https://api.openai.example".to_string(),
None,
None,
None,
None,
None,
None,
None,
)
.expect("endpoint transport should build"),
StoredProviderCatalogEndpoint::new(
"endpoint-openai-cli".to_string(),
"provider-openai".to_string(),
"openai:responses".to_string(),
Some("cli".to_string()),
Some("secondary".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
"https://api.openai.example".to_string(),
None,
None,
None,
None,
None,
None,
None,
)
.expect("endpoint transport should build"),
],
vec![sample_key(
"key-openai-cli",
"provider-openai",
"openai:responses",
"sk-test-cli",
)],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/models"))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-openai",
"api_key_id": "key-openai-cli"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(payload["data"]["error"], serde_json::Value::Null);
assert_eq!(payload["data"]["from_cache"], json!(false));
assert_eq!(
payload["data"]["models"][0]["api_formats"],
json!(["openai:responses"])
);
assert_eq!(
*execution_runtime_hits.lock().expect("mutex should lock"),
1
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_provider_query_models_aggregating_active_keys() {
let execution_runtime_hits = Arc::new(Mutex::new(0usize));
let execution_runtime_hits_clone = Arc::clone(&execution_runtime_hits);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| {
let execution_runtime_hits_inner = Arc::clone(&execution_runtime_hits_clone);
async move {
*execution_runtime_hits_inner
.lock()
.expect("mutex should lock") += 1;
assert_eq!(plan.url, "https://api.openai.example/v1/models");
let auth = plan
.headers
.get("authorization")
.map(String::as_str)
.unwrap_or_default()
.to_string();
let body = if auth == "Bearer sk-test-1" {
json!({
"data": [{
"id": "gpt-5",
"api_formats": ["openai:chat"],
"object": "model",
"owned_by": "system",
"created": 1732517497u64
}]
})
} else {
json!({
"data": [{
"id": "gpt-4.1",
"api_formats": ["openai:chat"],
"object": "model",
"owned_by": "system",
"created": 1732517498u64
}]
})
};
Json(json!({
"request_id": format!("req-provider-query-{auth}"),
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": body
}
}))
}
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-openai", "OpenAI", 10);
provider.provider_type = "openai".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![StoredProviderCatalogEndpoint::new(
"endpoint-openai-chat".to_string(),
"provider-openai".to_string(),
"openai:chat".to_string(),
Some("chat".to_string()),
Some("primary".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
"https://api.openai.example".to_string(),
None,
None,
None,
None,
None,
None,
None,
)
.expect("endpoint transport should build")],
vec![
sample_key(
"key-openai-1",
"provider-openai",
"openai:chat",
"sk-test-1",
),
sample_key(
"key-openai-2",
"provider-openai",
"openai:chat",
"sk-test-2",
),
],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/models"))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-openai"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(payload["data"]["from_cache"], json!(false));
assert_eq!(payload["data"]["keys_total"], json!(2));
assert_eq!(payload["data"]["keys_cached"], json!(0));
assert_eq!(payload["data"]["keys_fetched"], json!(2));
let models = payload["data"]["models"]
.as_array()
.expect("models should be an array");
assert_eq!(models.len(), 2);
let model_ids = models
.iter()
.map(|model| model["id"].as_str().expect("id should exist"))
.collect::<Vec<_>>();
feat(model-fetch): 对齐 Rust 上游模型抓取行为到 Python 语义 将 Rust 版上游可用模型抓取逻辑收敛到 Python 版行为,统一后台自动抓模 与管理员 provider-query 的模型发现路径,消除标准 /models、固定模型目录、 Antigravity、Vertex AI 等 provider 在两端实现上的分叉。 核心变更: - 在 aether-model-fetch 中引入统一抓模策略层 - 覆盖标准 /models、Vertex API Key、Vertex Service Account、 Antigravity fetchAvailableModels、固定模型目录五类抓模路径 - 将 provider-query 与后台自动抓模都切换到共享抓模入口,避免重复拼接 URL、headers 和 provider 特判逻辑 标准模型抓取对齐: - 按 Python 语义调整抓模优先级: openai:chat > openai:cli > openai:compact claude:chat > claude:cli gemini:chat > gemini:cli - 从抓模候选中移除 openai:responses - 为 openai:cli/openai:compact、claude:cli、gemini:* 补齐 Python 同款 User-Agent / 浏览器指纹请求头 - Claude 抓模保留 after_id 分页语义 - Gemini 抓模统一为 v1beta/models?key=... 语义 provider-query 对齐: - 返回结果改为按 model id 聚合,并合并/排序 api_formats - 最终模型列表按 model id 排序,行为与 Python 保持一致 - 固定目录 provider(codex/kiro/claude_code/gemini_cli)不再依赖活跃 endpoint,即使无 endpoint 也能返回预设模型目录 - Antigravity 多 key 查询改为按账户可用性 + tier 排序,首个成功结果即 停止,并接入 provider 级缓存 - 仅配置 openai:responses 的 provider 不再被视为抓模成功路径 自动抓模对齐: - 自动抓模成功时写入 allowed_models、upstream_models cache,并同步 upstream_metadata - upstream_metadata 合并逻辑对齐 Python,对 quota_by_model 做模型级合并, 并保留已有 reset_time - 自动抓模失败时不覆盖已有 allowed_models - 固定目录 provider 在无 endpoint 场景下也可成功更新 allowed_models Antigravity 对齐: - 使用 POST /v1internal:fetchAvailableModels 抓取可用模型 - 按 Python 规则处理 URL fallback 和 429/404/408/5xx fallback 状态 - 强制要求 auth_config.project_id - 过滤 Python 黑名单模型 - 解析并持久化 upstream_metadata.antigravity.quota_by_model Vertex AI 对齐: - API Key 模式仅抓取 publishers/google/models - Service Account 模式新增 JWT token exchange,并按 Python region 顺序 抓取 google + anthropic publishers - 模型 owned_by / display_name / api_format 推断与 Python 对齐 - 软 404 处理行为与 Python 收敛 Gemini CLI / 固定目录对齐: - Gemini CLI 改为返回 Python 预设模型目录 - 在可用时通过 loadCodeAssist 补充 plan_type/project_id 元数据 - Codex/Kiro/Claude Code 改为共享固定模型目录实现 测试: - 扩展 aether-model-fetch 单元测试,覆盖格式优先级、openai:responses 排除、 请求头、Claude 分页、Gemini query auth、固定目录与 metadata 合并 - 调整 provider-query 控制面测试到 Python 语义 - 新增自动抓模运行时测试,覆盖固定目录成功、Antigravity metadata 合并、 失败保留旧 allowed_models 验证: - cargo nextest run -p aether-model-fetch --lib - cargo nextest run -p aether-gateway control::admin::provider_query model_fetch::runtime::tests
2026-04-12 10:14:29 +08:00
assert_eq!(model_ids, vec!["gpt-4.1", "gpt-5"]);
assert_eq!(
*execution_runtime_hits.lock().expect("mutex should lock"),
2
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
feat(model-fetch): 对齐 Rust 上游模型抓取行为到 Python 语义 将 Rust 版上游可用模型抓取逻辑收敛到 Python 版行为,统一后台自动抓模 与管理员 provider-query 的模型发现路径,消除标准 /models、固定模型目录、 Antigravity、Vertex AI 等 provider 在两端实现上的分叉。 核心变更: - 在 aether-model-fetch 中引入统一抓模策略层 - 覆盖标准 /models、Vertex API Key、Vertex Service Account、 Antigravity fetchAvailableModels、固定模型目录五类抓模路径 - 将 provider-query 与后台自动抓模都切换到共享抓模入口,避免重复拼接 URL、headers 和 provider 特判逻辑 标准模型抓取对齐: - 按 Python 语义调整抓模优先级: openai:chat > openai:cli > openai:compact claude:chat > claude:cli gemini:chat > gemini:cli - 从抓模候选中移除 openai:responses - 为 openai:cli/openai:compact、claude:cli、gemini:* 补齐 Python 同款 User-Agent / 浏览器指纹请求头 - Claude 抓模保留 after_id 分页语义 - Gemini 抓模统一为 v1beta/models?key=... 语义 provider-query 对齐: - 返回结果改为按 model id 聚合,并合并/排序 api_formats - 最终模型列表按 model id 排序,行为与 Python 保持一致 - 固定目录 provider(codex/kiro/claude_code/gemini_cli)不再依赖活跃 endpoint,即使无 endpoint 也能返回预设模型目录 - Antigravity 多 key 查询改为按账户可用性 + tier 排序,首个成功结果即 停止,并接入 provider 级缓存 - 仅配置 openai:responses 的 provider 不再被视为抓模成功路径 自动抓模对齐: - 自动抓模成功时写入 allowed_models、upstream_models cache,并同步 upstream_metadata - upstream_metadata 合并逻辑对齐 Python,对 quota_by_model 做模型级合并, 并保留已有 reset_time - 自动抓模失败时不覆盖已有 allowed_models - 固定目录 provider 在无 endpoint 场景下也可成功更新 allowed_models Antigravity 对齐: - 使用 POST /v1internal:fetchAvailableModels 抓取可用模型 - 按 Python 规则处理 URL fallback 和 429/404/408/5xx fallback 状态 - 强制要求 auth_config.project_id - 过滤 Python 黑名单模型 - 解析并持久化 upstream_metadata.antigravity.quota_by_model Vertex AI 对齐: - API Key 模式仅抓取 publishers/google/models - Service Account 模式新增 JWT token exchange,并按 Python region 顺序 抓取 google + anthropic publishers - 模型 owned_by / display_name / api_format 推断与 Python 对齐 - 软 404 处理行为与 Python 收敛 Gemini CLI / 固定目录对齐: - Gemini CLI 改为返回 Python 预设模型目录 - 在可用时通过 loadCodeAssist 补充 plan_type/project_id 元数据 - Codex/Kiro/Claude Code 改为共享固定模型目录实现 测试: - 扩展 aether-model-fetch 单元测试,覆盖格式优先级、openai:responses 排除、 请求头、Claude 分页、Gemini query auth、固定目录与 metadata 合并 - 调整 provider-query 控制面测试到 Python 语义 - 新增自动抓模运行时测试,覆盖固定目录成功、Antigravity metadata 合并、 失败保留旧 allowed_models 验证: - cargo nextest run -p aether-model-fetch --lib - cargo nextest run -p aether-gateway control::admin::provider_query model_fetch::runtime::tests
2026-04-12 10:14:29 +08:00
#[tokio::test]
async fn gateway_handles_admin_provider_query_models_for_fixed_provider_without_endpoint() {
let execution_runtime_hits = Arc::new(Mutex::new(0usize));
let execution_runtime_hits_clone = Arc::clone(&execution_runtime_hits);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |_request: Request| {
let execution_runtime_hits_inner = Arc::clone(&execution_runtime_hits_clone);
async move {
*execution_runtime_hits_inner
.lock()
.expect("mutex should lock") += 1;
Json(json!({
"request_id": "unexpected",
"status_code": 500
}))
}
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-codex", "Codex", 10);
provider.provider_type = "codex".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![],
vec![sample_key(
"key-codex-oauth",
"provider-codex",
"openai:responses",
feat(model-fetch): 对齐 Rust 上游模型抓取行为到 Python 语义 将 Rust 版上游可用模型抓取逻辑收敛到 Python 版行为,统一后台自动抓模 与管理员 provider-query 的模型发现路径,消除标准 /models、固定模型目录、 Antigravity、Vertex AI 等 provider 在两端实现上的分叉。 核心变更: - 在 aether-model-fetch 中引入统一抓模策略层 - 覆盖标准 /models、Vertex API Key、Vertex Service Account、 Antigravity fetchAvailableModels、固定模型目录五类抓模路径 - 将 provider-query 与后台自动抓模都切换到共享抓模入口,避免重复拼接 URL、headers 和 provider 特判逻辑 标准模型抓取对齐: - 按 Python 语义调整抓模优先级: openai:chat > openai:cli > openai:compact claude:chat > claude:cli gemini:chat > gemini:cli - 从抓模候选中移除 openai:responses - 为 openai:cli/openai:compact、claude:cli、gemini:* 补齐 Python 同款 User-Agent / 浏览器指纹请求头 - Claude 抓模保留 after_id 分页语义 - Gemini 抓模统一为 v1beta/models?key=... 语义 provider-query 对齐: - 返回结果改为按 model id 聚合,并合并/排序 api_formats - 最终模型列表按 model id 排序,行为与 Python 保持一致 - 固定目录 provider(codex/kiro/claude_code/gemini_cli)不再依赖活跃 endpoint,即使无 endpoint 也能返回预设模型目录 - Antigravity 多 key 查询改为按账户可用性 + tier 排序,首个成功结果即 停止,并接入 provider 级缓存 - 仅配置 openai:responses 的 provider 不再被视为抓模成功路径 自动抓模对齐: - 自动抓模成功时写入 allowed_models、upstream_models cache,并同步 upstream_metadata - upstream_metadata 合并逻辑对齐 Python,对 quota_by_model 做模型级合并, 并保留已有 reset_time - 自动抓模失败时不覆盖已有 allowed_models - 固定目录 provider 在无 endpoint 场景下也可成功更新 allowed_models Antigravity 对齐: - 使用 POST /v1internal:fetchAvailableModels 抓取可用模型 - 按 Python 规则处理 URL fallback 和 429/404/408/5xx fallback 状态 - 强制要求 auth_config.project_id - 过滤 Python 黑名单模型 - 解析并持久化 upstream_metadata.antigravity.quota_by_model Vertex AI 对齐: - API Key 模式仅抓取 publishers/google/models - Service Account 模式新增 JWT token exchange,并按 Python region 顺序 抓取 google + anthropic publishers - 模型 owned_by / display_name / api_format 推断与 Python 对齐 - 软 404 处理行为与 Python 收敛 Gemini CLI / 固定目录对齐: - Gemini CLI 改为返回 Python 预设模型目录 - 在可用时通过 loadCodeAssist 补充 plan_type/project_id 元数据 - Codex/Kiro/Claude Code 改为共享固定模型目录实现 测试: - 扩展 aether-model-fetch 单元测试,覆盖格式优先级、openai:responses 排除、 请求头、Claude 分页、Gemini query auth、固定目录与 metadata 合并 - 调整 provider-query 控制面测试到 Python 语义 - 新增自动抓模运行时测试,覆盖固定目录成功、Antigravity metadata 合并、 失败保留旧 allowed_models 验证: - cargo nextest run -p aether-model-fetch --lib - cargo nextest run -p aether-gateway control::admin::provider_query model_fetch::runtime::tests
2026-04-12 10:14:29 +08:00
"sk-test-codex",
)],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/models"))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-codex",
"api_key_id": "key-codex-oauth"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(payload["data"]["error"], serde_json::Value::Null);
assert_eq!(payload["data"]["from_cache"], json!(false));
let models = payload["data"]["models"]
.as_array()
.expect("models should be an array");
assert!(models.iter().any(|model| model["id"] == "gpt-5.4"));
assert_eq!(
*execution_runtime_hits.lock().expect("mutex should lock"),
0
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_provider_query_test_model_locally_with_trusted_admin_principal() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
assert_eq!(plan.provider_id, "provider-openai");
assert_eq!(plan.endpoint_id, "endpoint-openai-chat");
assert_eq!(plan.key_id, "key-openai-primary");
assert_eq!(plan.provider_api_format, "openai:chat");
assert_eq!(plan.model_name.as_deref(), Some("gpt-4.1"));
assert!(!plan.stream);
assert_eq!(
plan.headers.get("content-type").map(String::as_str),
Some("application/json")
);
assert_eq!(
plan.headers.get("authorization").map(String::as_str),
Some("Bearer sk-test-primary")
);
assert_eq!(
plan.body
.json_body
.as_ref()
.and_then(|body| body.get("model")),
Some(&json!("gpt-4.1"))
);
Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"id": "chatcmpl-test-model",
"object": "chat.completion",
"choices": [{
"message": {
"role": "assistant",
"content": "Hello from OpenAI"
}
}]
}
},
"telemetry": {
"elapsed_ms": 18
}
}))
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-openai", "OpenAI", 10);
provider.provider_type = "openai".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-openai-chat",
"provider-openai",
"openai:chat",
"https://api.openai.example",
)],
vec![sample_key(
"key-openai-primary",
"provider-openai",
"openai:chat",
"sk-test-primary",
)],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/test-model"))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-openai",
"model": "gpt-4.1",
"api_format": "openai:chat"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(payload["provider"]["id"], json!("provider-openai"));
assert_eq!(payload["model"], json!("gpt-4.1"));
assert_eq!(payload["error"], serde_json::Value::Null);
assert_eq!(
payload["data"]["response"]["choices"][0]["message"]["content"],
json!("Hello from OpenAI")
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_provider_query_embedding_model_test() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
assert_eq!(plan.provider_id, "provider-siliconflow");
assert_eq!(plan.endpoint_id, "endpoint-siliconflow-embedding");
assert_eq!(plan.key_id, "key-siliconflow-embedding");
assert_eq!(plan.client_api_format, "openai:embedding");
assert_eq!(plan.provider_api_format, "openai:embedding");
assert_eq!(plan.url, "https://api.siliconflow.example/v1/embeddings");
assert_eq!(plan.model_name.as_deref(), Some("Qwen/Qwen3-Embedding-4B"));
assert!(!plan.stream);
assert_eq!(
plan.headers.get("authorization").map(String::as_str),
Some("Bearer sk-siliconflow-embedding")
);
let body = plan.body.json_body.as_ref().expect("json body");
assert_eq!(body["model"], json!("Qwen/Qwen3-Embedding-4B"));
assert_eq!(body["input"], json!("This is a test embedding input."));
assert!(
body.get("stream").is_none(),
"embedding provider body must not carry stream"
);
Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"object": "list",
"model": "Qwen/Qwen3-Embedding-4B",
"data": [{
"object": "embedding",
"index": 0,
"embedding": [0.1, 0.2, 0.3]
}]
}
},
"telemetry": {
"elapsed_ms": 24
}
}))
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-siliconflow", "SiliconFlow", 10);
provider.provider_type = "custom".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-siliconflow-embedding",
"provider-siliconflow",
"openai:embedding",
"https://api.siliconflow.example",
)],
vec![sample_key(
"key-siliconflow-embedding",
"provider-siliconflow",
"openai:embedding",
"sk-siliconflow-embedding",
)],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/test-model"))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-siliconflow",
"model": "Qwen/Qwen3-Embedding-4B",
"api_format": "openai:embedding",
"endpoint_id": "endpoint-siliconflow-embedding",
"request_body": {
"model": "Qwen/Qwen3-Embedding-4B",
"input": "This is a test embedding input."
}
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(payload["error"], serde_json::Value::Null);
assert_eq!(payload["attempts"][0]["status"], json!("success"));
assert_eq!(
payload["attempts"][0]["response_body"]["data"][0]["object"],
json!("embedding")
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_provider_query_doubao_text_embedding_model_test() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
assert_eq!(plan.provider_id, "provider-doubao");
assert_eq!(plan.endpoint_id, "endpoint-doubao-embedding");
assert_eq!(plan.key_id, "key-doubao-embedding");
assert_eq!(plan.client_api_format, "openai:embedding");
assert_eq!(plan.provider_api_format, "doubao:embedding");
assert_eq!(plan.url, "https://ark.volces.example/api/v3/embeddings");
assert_eq!(
plan.model_name.as_deref(),
Some("doubao-embedding-text-240515")
);
assert!(!plan.stream);
assert_eq!(
plan.headers.get("authorization").map(String::as_str),
Some("Bearer sk-doubao-embedding")
);
let body = plan.body.json_body.as_ref().expect("json body");
assert_eq!(body["model"], json!("doubao-embedding-text-240515"));
assert_eq!(body["input"], json!(["This is a test embedding input."]));
assert!(
body.get("stream").is_none(),
"doubao embedding provider body must not carry stream"
);
Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"object": "list",
"model": "doubao-embedding-text-240515",
"data": [{
"object": "embedding",
"index": 0,
"embedding": [0.1, 0.2, 0.3]
}]
}
},
"telemetry": {
"elapsed_ms": 28
}
}))
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-doubao", "Doubao", 10);
provider.provider_type = "doubao".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-doubao-embedding",
"provider-doubao",
"doubao:embedding",
"https://ark.volces.example/api/v3",
)],
vec![sample_key(
"key-doubao-embedding",
"provider-doubao",
"doubao:embedding",
"sk-doubao-embedding",
)],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/test-model"))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-doubao",
"model": "doubao-embedding-text-240515",
"api_format": "doubao:embedding",
"endpoint_id": "endpoint-doubao-embedding",
"request_body": {
"model": "doubao-embedding-text-240515",
"input": "This is a test embedding input."
}
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(payload["error"], serde_json::Value::Null);
assert_eq!(payload["attempts"][0]["status"], json!("success"));
assert_eq!(
payload["attempts"][0]["request_body"]["input"],
json!(["This is a test embedding input."])
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_provider_query_gemini_embedding_model_test() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
assert_eq!(plan.provider_id, "provider-gemini");
assert_eq!(plan.endpoint_id, "endpoint-gemini-embedding");
assert_eq!(plan.key_id, "key-gemini-embedding");
assert_eq!(plan.client_api_format, "openai:embedding");
assert_eq!(plan.provider_api_format, "gemini:embedding");
assert_eq!(
plan.url,
"https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-001:embedContent"
);
assert_eq!(plan.model_name.as_deref(), Some("gemini-embedding-001"));
assert!(!plan.stream);
assert_eq!(
plan.headers.get("x-goog-api-key").map(String::as_str),
Some("sk-gemini-embedding")
);
let body = plan.body.json_body.as_ref().expect("json body");
assert_eq!(body["model"], json!("gemini-embedding-001"));
assert_eq!(
body["content"]["parts"][0]["text"],
json!("This is a test embedding input.")
);
assert!(body.get("requests").is_none());
assert!(
body.get("stream").is_none(),
"gemini embedding provider body must not carry stream"
);
Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"model": "gemini-embedding-001",
"embedding": {
"values": [0.1, 0.2, 0.3]
}
}
},
"telemetry": {
"elapsed_ms": 27
}
}))
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-gemini", "Gemini", 10);
provider.provider_type = "gemini".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-gemini-embedding",
"provider-gemini",
"gemini:embedding",
"https://generativelanguage.googleapis.com/v1beta",
)],
vec![sample_key(
"key-gemini-embedding",
"provider-gemini",
"gemini:embedding",
"sk-gemini-embedding",
)],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/test-model"))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-gemini",
"model": "gemini-embedding-001",
"api_format": "gemini:embedding",
"endpoint_id": "endpoint-gemini-embedding",
"request_body": {
"model": "gemini-embedding-001",
"input": "This is a test embedding input."
}
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(payload["error"], serde_json::Value::Null);
assert_eq!(payload["attempts"][0]["status"], json!("success"));
assert_eq!(
payload["attempts"][0]["response_body"]["embedding"]["values"][0],
json!(0.1)
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_provider_query_jina_embedding_model_test() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
assert_eq!(plan.provider_id, "provider-jina-embedding");
assert_eq!(plan.endpoint_id, "endpoint-jina-embedding");
assert_eq!(plan.key_id, "key-jina-embedding");
assert_eq!(plan.client_api_format, "openai:embedding");
assert_eq!(plan.provider_api_format, "jina:embedding");
assert_eq!(plan.url, "https://api.jina.example/v1/embeddings");
assert_eq!(plan.model_name.as_deref(), Some("jina-embeddings-v3"));
assert!(!plan.stream);
assert_eq!(
plan.headers.get("authorization").map(String::as_str),
Some("Bearer sk-jina-embedding")
);
let body = plan.body.json_body.as_ref().expect("json body");
assert_eq!(body["model"], json!("jina-embeddings-v3"));
assert_eq!(body["task"], json!("text-matching"));
assert_eq!(body["input"], json!("This is a test embedding input."));
assert!(
body.get("stream").is_none(),
"jina embedding provider body must not carry stream"
);
Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"model": "jina-embeddings-v3",
"data": [{
"object": "embedding",
"index": 0,
"embedding": [0.1, 0.2, 0.3]
}]
}
},
"telemetry": {
"elapsed_ms": 29
}
}))
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-jina-embedding", "Jina", 10);
provider.provider_type = "jina".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-jina-embedding",
"provider-jina-embedding",
"jina:embedding",
"https://api.jina.example",
)],
vec![sample_key(
"key-jina-embedding",
"provider-jina-embedding",
"jina:embedding",
"sk-jina-embedding",
)],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/test-model"))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-jina-embedding",
"model": "jina-embeddings-v3",
"api_format": "jina:embedding",
"endpoint_id": "endpoint-jina-embedding",
"request_body": {
"model": "jina-embeddings-v3",
"input": "This is a test embedding input."
}
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(payload["error"], serde_json::Value::Null);
assert_eq!(payload["attempts"][0]["status"], json!("success"));
assert_eq!(
payload["attempts"][0]["request_body"]["task"],
json!("text-matching")
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_provider_query_openai_rerank_model_test() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
assert_eq!(plan.provider_id, "provider-openai-rerank");
assert_eq!(plan.endpoint_id, "endpoint-openai-rerank");
assert_eq!(plan.key_id, "key-openai-rerank");
assert_eq!(plan.client_api_format, "openai:rerank");
assert_eq!(plan.provider_api_format, "openai:rerank");
assert_eq!(plan.url, "https://api.openai.example/v1/rerank");
assert_eq!(plan.model_name.as_deref(), Some("bge-reranker-base"));
assert!(!plan.stream);
assert_eq!(
plan.headers.get("authorization").map(String::as_str),
Some("Bearer sk-openai-rerank")
);
let body = plan.body.json_body.as_ref().expect("json body");
assert_eq!(body["model"], json!("bge-reranker-base"));
assert_eq!(body["query"], json!("Apple"));
assert_eq!(
body["documents"],
json!(["apple", "banana", "fruit", "vegetable"])
);
assert_eq!(body["return_documents"], json!(true));
assert_eq!(body["top_n"], json!(4));
assert!(
body.get("stream").is_none(),
"openai rerank provider body must not carry stream"
);
Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"model": "bge-reranker-base",
"results": [{
"index": 0,
"relevance_score": 0.91
}]
}
},
"telemetry": {
"elapsed_ms": 32
}
}))
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-openai-rerank", "OpenAI Rerank", 10);
provider.provider_type = "openai".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-openai-rerank",
"provider-openai-rerank",
"openai:rerank",
"https://api.openai.example/v1",
)],
vec![sample_key(
"key-openai-rerank",
"provider-openai-rerank",
"openai:rerank",
"sk-openai-rerank",
)],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/test-model"))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-openai-rerank",
"model": "bge-reranker-base",
"api_format": "openai:rerank",
"endpoint_id": "endpoint-openai-rerank",
"request_body": {
"model": "bge-reranker-base",
"query": "Apple",
"documents": ["apple", "banana", "fruit", "vegetable"],
"return_documents": true,
"top_n": 4
}
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(payload["error"], serde_json::Value::Null);
assert_eq!(payload["attempts"][0]["status"], json!("success"));
assert_eq!(
payload["attempts"][0]["response_body"]["results"][0]["relevance_score"],
json!(0.91)
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_provider_query_rerank_model_test() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
assert_eq!(plan.provider_id, "provider-jina");
assert_eq!(plan.endpoint_id, "endpoint-jina-rerank");
assert_eq!(plan.key_id, "key-jina-rerank");
assert_eq!(plan.client_api_format, "openai:rerank");
assert_eq!(plan.provider_api_format, "jina:rerank");
assert_eq!(plan.url, "https://api.jina.example/v1/rerank");
assert_eq!(
plan.model_name.as_deref(),
Some("jina-reranker-v2-base-multilingual")
);
assert!(!plan.stream);
assert_eq!(
plan.headers.get("authorization").map(String::as_str),
Some("Bearer sk-jina-rerank")
);
let body = plan.body.json_body.as_ref().expect("json body");
assert_eq!(body["model"], json!("jina-reranker-v2-base-multilingual"));
assert_eq!(body["query"], json!("Apple"));
assert_eq!(
body["documents"],
json!(["apple", "banana", "fruit", "vegetable"])
);
assert_eq!(body["return_documents"], json!(true));
assert_eq!(body["top_n"], json!(4));
assert!(
body.get("stream").is_none(),
"rerank provider body must not carry stream"
);
Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"model": "jina-reranker-v2-base-multilingual",
"results": [{
"index": 0,
"relevance_score": 0.93
}]
}
},
"telemetry": {
"elapsed_ms": 31
}
}))
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-jina", "Jina", 10);
provider.provider_type = "jina".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-jina-rerank",
"provider-jina",
"jina:rerank",
"https://api.jina.example",
)],
vec![sample_key(
"key-jina-rerank",
"provider-jina",
"jina:rerank",
"sk-jina-rerank",
)],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/test-model"))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-jina",
"model": "jina-reranker-v2-base-multilingual",
"api_format": "jina:rerank",
"endpoint_id": "endpoint-jina-rerank",
"request_body": {
"model": "jina-reranker-v2-base-multilingual",
"query": "Apple",
"documents": ["apple", "banana", "fruit", "vegetable"],
"return_documents": true,
"top_n": 4
}
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(payload["error"], serde_json::Value::Null);
assert_eq!(payload["attempts"][0]["status"], json!("success"));
assert_eq!(
payload["attempts"][0]["response_body"]["results"][0]["relevance_score"],
json!(0.93)
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_maps_admin_provider_model_before_model_list_test_request() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
assert_eq!(plan.provider_id, "provider-minimax");
assert_eq!(plan.endpoint_id, "endpoint-minimax-chat");
assert_eq!(plan.key_id, "key-minimax-primary");
assert_eq!(plan.provider_api_format, "openai:chat");
let requested_model = plan
.model_name
.as_deref()
.expect("test plan should carry model name");
assert!(
matches!(
requested_model,
"MiniMax-M2.7-highspeed" | "MiniMax-M2.7-balanced" | "claude-opus-4-6"
),
"unexpected requested model: {requested_model}"
);
assert_eq!(
plan.body
.json_body
.as_ref()
.and_then(|body| body.get("model")),
Some(&json!(requested_model))
);
let response_model = if requested_model.starts_with("MiniMax-M2.7") {
"MiniMax-M2.7"
} else {
requested_model
};
Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"id": "chatcmpl-minimax-mapped-model",
"object": "chat.completion",
"model": response_model,
"choices": [{
"message": {
"role": "assistant",
"content": "Hello from MiniMax"
}
}]
}
},
"telemetry": {
"elapsed_ms": 22
}
}))
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-minimax", "MiniMax", 10);
provider.provider_type = "custom".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-minimax-chat",
"provider-minimax",
"openai:chat",
"https://api.minimax.example",
)],
vec![sample_key(
"key-minimax-primary",
"provider-minimax",
"openai:chat",
"sk-minimax-primary",
)],
));
let mut provider_model = sample_admin_provider_model(
"model-minimax-claude-opus",
"provider-minimax",
"global-claude-opus-4-6",
"claude-opus-4-6",
);
provider_model.global_model_name = Some("claude-opus-4-6".to_string());
provider_model.global_model_display_name = Some("Claude Opus 4.6".to_string());
provider_model.provider_model_mappings = Some(json!([
{
"name": "ignored-anthropic-model",
"priority": 1,
"api_formats": ["anthropic:messages"],
},
{
"name": "MiniMax-M2.7-highspeed",
"priority": 2,
"api_formats": ["openai:chat"],
"endpoint_ids": ["endpoint-minimax-chat"],
},
{
"name": "MiniMax-M2.7-balanced",
"priority": 3,
"api_formats": ["openai:chat"],
"endpoint_ids": ["endpoint-minimax-chat"],
}
]));
let global_model_repository = Arc::new(
InMemoryGlobalModelReadRepository::seed(Vec::new())
.with_admin_provider_models(vec![provider_model]),
);
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)
.with_global_model_repository_for_tests(global_model_repository),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/test-model"))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-minimax",
"model": "claude-opus-4-6",
"api_format": "openai:chat",
"endpoint_id": "endpoint-minimax-chat"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(payload["model"], json!("claude-opus-4-6"));
assert_eq!(
payload["data"]["response"]["choices"][0]["message"]["content"],
json!("Hello from MiniMax")
);
assert_eq!(
payload["attempts"][0]["effective_model"],
json!("MiniMax-M2.7-highspeed")
);
assert_eq!(
payload["attempts"][0]["request_body"]["model"],
json!("MiniMax-M2.7-highspeed")
);
assert_eq!(
payload["attempts"][0]["response_body"]["model"],
json!("MiniMax-M2.7")
);
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/test-model"))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-minimax",
"model": "claude-opus-4-6",
"api_format": "openai:chat",
"endpoint_id": "endpoint-minimax-chat",
"mapped_model_name": "MiniMax-M2.7-balanced",
"request_body": {
"model": "stale-model-from-ui",
"messages": [{
"role": "user",
"content": "custom prompt"
}]
}
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(
payload["attempts"][0]["effective_model"],
json!("MiniMax-M2.7-balanced")
);
assert_eq!(
payload["attempts"][0]["request_body"]["model"],
json!("MiniMax-M2.7-balanced")
);
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/test-model"))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-minimax",
"model": "claude-opus-4-6",
"api_format": "openai:chat",
"endpoint_id": "endpoint-minimax-chat",
"mapped_model_name": "ignored-anthropic-model"
}))
.send()
.await
.expect("request should fail");
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-query/test-model-failover"
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-minimax",
"mode": "global",
"model": "claude-opus-4-6",
"failover_models": ["claude-opus-4-6"],
"api_format": "openai:chat",
"endpoint_id": "endpoint-minimax-chat",
"apply_model_mapping": false
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(payload["model"], json!("claude-opus-4-6"));
assert_eq!(
payload["attempts"][0]["effective_model"],
json!("claude-opus-4-6")
);
assert_eq!(
payload["attempts"][0]["request_body"]["model"],
json!("claude-opus-4-6")
);
assert_eq!(
payload["attempts"][0]["response_body"]["model"],
json!("claude-opus-4-6")
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_streams_codex_openai_responses_upstream_for_admin_pool_model_test() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
assert_eq!(plan.provider_id, "provider-codex");
assert_eq!(plan.endpoint_id, "endpoint-codex-responses");
assert_eq!(plan.key_id, "key-codex-primary");
assert_eq!(plan.provider_api_format, "openai:responses");
assert_eq!(plan.model_name.as_deref(), Some("gpt-5.3-codex-spark"));
assert!(plan.stream, "Codex openai:responses must stream upstream");
assert_eq!(
plan.body
.json_body
.as_ref()
.and_then(|body| body.get("stream")),
Some(&json!(true))
);
Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"id": "resp-codex-model-test",
"model": "gpt-5.3-codex-spark",
"output_text": "ok"
}
},
"telemetry": {
"elapsed_ms": 18
}
}))
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-codex", "Codex", 10);
provider.provider_type = "codex".to_string();
let mut endpoint = sample_endpoint(
"endpoint-codex-responses",
"provider-codex",
"openai:responses",
"https://chatgpt.com/backend-api/codex",
);
endpoint.config = Some(json!({"upstream_stream_policy": "force_stream"}));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![sample_key(
"key-codex-primary",
"provider-codex",
"openai:responses",
"sk-codex-primary",
)],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-query/test-model-failover"
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-codex",
"mode": "pool",
"model": "gpt-5.3-codex-spark",
"failover_models": ["gpt-5.3-codex-spark"],
"api_format": "openai:responses",
"endpoint_id": "endpoint-codex-responses",
"request_body": {
"model": "gpt-5.3-codex-spark",
"input": "hello",
"stream": true
}
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(
payload["attempts"][0]["request_body"]["stream"],
json!(true)
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_uses_pool_scheduler_order_for_admin_pool_model_test() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
assert_eq!(plan.key_id, "key-codex-plus");
Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"id": "resp-codex-plus",
"model": "gpt-5.4-mini",
"output_text": "ok"
}
},
"telemetry": {
"elapsed_ms": 21
}
}))
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-codex", "Codex", 10);
provider.provider_type = "codex".to_string();
provider.config = Some(json!({
"pool_advanced": {
"scheduling_presets": [
{"preset": "plus_first", "enabled": true}
]
}
}));
let mut free_key = sample_key(
"key-codex-free",
"provider-codex",
"openai:responses",
"sk-codex-free",
);
free_key.internal_priority = 0;
free_key.status_snapshot = Some(json!({
"quota": {
"provider_type": "codex",
"code": "limited",
"plan_type": "free",
"usage_ratio": 0.1
}
}));
let mut plus_key = sample_key(
"key-codex-plus",
"provider-codex",
"openai:responses",
"sk-codex-plus",
);
plus_key.internal_priority = 100;
plus_key.status_snapshot = Some(json!({
"quota": {
"provider_type": "codex",
"code": "limited",
"plan_type": "plus",
"usage_ratio": 0.1
}
}));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-codex-responses",
"provider-codex",
"openai:responses",
"https://chatgpt.com/backend-api/codex",
)],
vec![free_key, plus_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_request_candidate_repository_for_tests(
request_candidate_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY.to_string())
.attach_provider_catalog_repository_for_tests(provider_catalog_repository),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-query/test-model-failover"
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-codex",
"mode": "pool",
"model": "gpt-5.4-mini",
"failover_models": ["gpt-5.4-mini"],
"api_format": "openai:responses",
"endpoint_id": "endpoint-codex-responses",
"request_id": "provider-query-model-test-trace-123",
"request_body": {
"model": "gpt-5.4-mini",
"input": "hello",
"stream": true
}
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(payload["attempts"][0]["key_id"], json!("key-codex-plus"));
assert_eq!(
payload["candidate_summary"]["winning_key_id"],
json!("key-codex-plus")
);
assert_eq!(payload["candidate_summary"]["total_candidates"], json!(2));
assert_eq!(payload["candidate_summary"]["attempted"], json!(1));
assert_eq!(payload["candidate_summary"]["unused"], json!(1));
let stored_candidates = request_candidate_repository
.list_by_request_id("provider-query-model-test-trace-123")
.await
.expect("model-test trace candidates should read");
assert_eq!(stored_candidates.len(), 2);
assert_eq!(
stored_candidates[0].key_id.as_deref(),
Some("key-codex-plus")
);
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success);
assert_eq!(
stored_candidates[1].key_id.as_deref(),
Some("key-codex-free")
);
assert_eq!(stored_candidates[1].status, RequestCandidateStatus::Unused);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_provider_query_test_model_failover_locally_with_trusted_admin_principal(
) {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
let auth = plan
.headers
.get("authorization")
.map(String::as_str)
.unwrap_or_default()
.to_string();
assert_eq!(plan.provider_id, "provider-openai");
assert_eq!(plan.endpoint_id, "endpoint-openai-chat");
assert_eq!(plan.provider_api_format, "openai:chat");
assert_eq!(plan.model_name.as_deref(), Some("gpt-4.1"));
assert_eq!(
plan.headers.get("x-test-header").map(String::as_str),
Some("from-admin")
);
assert_eq!(
plan.body
.json_body
.as_ref()
.and_then(|body| body.get("messages"))
.and_then(|messages| messages.as_array())
.and_then(|messages| messages.first())
.and_then(|message| message.get("content")),
Some(&json!("custom prompt"))
);
let payload = if auth == "Bearer sk-test-first" {
json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 429,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"error": {
"message": "too many requests"
}
}
},
"telemetry": {
"elapsed_ms": 11
}
})
} else {
json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"id": "chatcmpl-failover",
"object": "chat.completion",
"choices": [{
"message": {
"role": "assistant",
"content": "Recovered from OpenAI failover"
}
}]
}
},
"telemetry": {
"elapsed_ms": 27
}
})
};
Json(payload)
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-openai", "OpenAI", 10);
provider.provider_type = "openai".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-openai-chat",
"provider-openai",
"openai:chat",
"https://api.openai.example",
)],
vec![
sample_key(
"key-openai-first",
"provider-openai",
"openai:chat",
"sk-test-first",
),
sample_key(
"key-openai-second",
"provider-openai",
"openai:chat",
"sk-test-second",
),
],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-query/test-model-failover"
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-openai",
"mode": "direct",
"model_name": "gpt-4.1",
"failover_models": ["gpt-4.1"],
"api_format": "openai:chat",
"request_headers": {
"x-test-header": "from-admin"
},
"request_body": {
"model": "ignored-model",
"messages": [{
"role": "user",
"content": "custom prompt"
}]
},
"request_id": "provider-test-openai"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(payload["total_candidates"], json!(2));
assert_eq!(payload["total_attempts"], json!(2));
assert_eq!(payload["candidate_summary"]["total_candidates"], json!(2));
assert_eq!(payload["candidate_summary"]["attempted"], json!(2));
assert_eq!(payload["candidate_summary"]["failed"], json!(1));
assert_eq!(payload["candidate_summary"]["success"], json!(1));
assert_eq!(
payload["candidate_summary"]["stop_reason"],
json!("first_success")
);
let attempts = payload["attempts"]
.as_array()
.expect("attempts should be an array");
assert_eq!(attempts.len(), 2);
assert_eq!(attempts[0]["status"], json!("failed"));
assert_eq!(attempts[0]["status_code"], json!(429));
assert_eq!(attempts[1]["status"], json!("success"));
assert_eq!(attempts[1]["key_id"], json!("key-openai-second"));
assert_eq!(attempts[1]["request_body"]["model"], json!("gpt-4.1"));
assert_eq!(payload["data"]["stream"], json!(false));
assert_eq!(
payload["data"]["response"]["choices"][0]["message"]["content"],
json!("Recovered from OpenAI failover")
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_provider_query_test_model_for_kiro_locally() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
assert_eq!(plan.provider_id, "provider-kiro");
assert_eq!(plan.endpoint_id, "endpoint-kiro-cli");
assert_eq!(plan.key_id, "key-kiro-primary");
assert_eq!(plan.provider_api_format, "claude:messages");
assert_eq!(plan.model_name.as_deref(), Some("claude-sonnet-4-upstream"));
Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/vnd.amazon.eventstream"
},
"body": {
"body_bytes_b64": base64::engine::general_purpose::STANDARD.encode(
[
encode_kiro_event_frame("assistantResponseEvent", json!({"content": "Hello from Kiro"})),
encode_kiro_exception_frame("ContentLengthExceededException"),
]
.concat()
)
},
"telemetry": {
"elapsed_ms": 42
}
}))
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-kiro", "Kiro", 10);
provider.provider_type = "kiro".to_string();
let mut key = sample_key(
"key-kiro-primary",
"provider-kiro",
"claude:messages",
"__placeholder__",
);
key.auth_type = "oauth".to_string();
key.encrypted_auth_config = Some(
aether_crypto::encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{
"provider_type":"kiro",
"auth_method":"idc",
"access_token":"cached-kiro-token",
"expires_at":4102444800,
"refresh_token":"rrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrr",
"machine_id":"123e4567-e89b-12d3-a456-426614174000",
"api_region":"us-east-1",
"client_id":"client-id",
"client_secret":"client-secret"
}"#,
)
.expect("auth config should encrypt"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![StoredProviderCatalogEndpoint::new(
"endpoint-kiro-cli".to_string(),
"provider-kiro".to_string(),
"claude:messages".to_string(),
Some("claude".to_string()),
Some("messages".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
"https://q.{region}.amazonaws.com".to_string(),
None,
None,
None,
None,
None,
None,
None,
)
.expect("endpoint transport should build")],
vec![key],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/test-model"))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-kiro",
"model_name": "claude-sonnet-4-upstream",
"api_format": "claude:messages"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(payload["provider"]["id"], json!("provider-kiro"));
assert_eq!(payload["model"], json!("claude-sonnet-4-upstream"));
assert_eq!(
payload["data"]["response"]["content"][0]["text"],
json!("Hello from Kiro")
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_uses_kiro_mapped_model_name_for_explicit_model_mapping_test() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
assert_eq!(plan.provider_id, "provider-kiro");
assert_eq!(plan.endpoint_id, "endpoint-kiro-cli");
assert_eq!(plan.key_id, "key-kiro-primary");
assert_eq!(plan.provider_api_format, "claude:messages");
assert_eq!(plan.model_name.as_deref(), Some("claude-haiku-4.5"));
let body = plan.body.json_body.as_ref().expect("json body");
assert_eq!(
body.pointer("/conversationState/currentMessage/userInputMessage/modelId"),
Some(&json!("claude-haiku-4.5"))
);
Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/vnd.amazon.eventstream"
},
"body": {
"body_bytes_b64": base64::engine::general_purpose::STANDARD.encode(
[
encode_kiro_event_frame("assistantResponseEvent", json!({"content": "Hello from Kiro mapped model"})),
encode_kiro_exception_frame("ContentLengthExceededException"),
]
.concat()
)
},
"telemetry": {
"elapsed_ms": 42
}
}))
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-kiro", "Kiro", 10);
provider.provider_type = "kiro".to_string();
let mut key = sample_key(
"key-kiro-primary",
"provider-kiro",
"claude:messages",
"__placeholder__",
);
key.auth_type = "oauth".to_string();
key.encrypted_auth_config = Some(
aether_crypto::encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{
"provider_type":"kiro",
"auth_method":"idc",
"access_token":"cached-kiro-token",
"expires_at":4102444800,
"refresh_token":"rrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrr",
"machine_id":"123e4567-e89b-12d3-a456-426614174000",
"api_region":"us-east-1",
"client_id":"client-id",
"client_secret":"client-secret"
}"#,
)
.expect("auth config should encrypt"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![StoredProviderCatalogEndpoint::new(
"endpoint-kiro-cli".to_string(),
"provider-kiro".to_string(),
"claude:messages".to_string(),
Some("claude".to_string()),
Some("messages".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
"https://q.{region}.amazonaws.com".to_string(),
None,
None,
None,
None,
None,
None,
None,
)
.expect("endpoint transport should build")],
vec![key],
));
let mut provider_model = sample_admin_provider_model(
"model-kiro-haiku",
"provider-kiro",
"global-kiro-haiku",
"claude-haiku-4-5-20251001",
);
provider_model.provider_model_mappings = Some(json!([{
"name": "claude-haiku-4.5",
"priority": 1,
"api_formats": ["claude:messages"],
"endpoint_ids": ["endpoint-kiro-cli"]
}]));
let global_model_repository = Arc::new(
InMemoryGlobalModelReadRepository::seed(Vec::new())
.with_admin_provider_models(vec![provider_model]),
);
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)
.with_global_model_repository_for_tests(global_model_repository),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-query/test-model-failover"
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-kiro",
"mode": "pool",
"model_name": "claude-haiku-4-5-20251001",
"failover_models": ["claude-haiku-4-5-20251001"],
"mapped_model_name": "claude-haiku-4.5",
"api_format": "claude:messages",
"endpoint_id": "endpoint-kiro-cli",
"request_id": "provider-test-kiro-mapped"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(
payload["attempts"][0]["effective_model"],
json!("claude-haiku-4.5")
);
assert_eq!(
payload["attempts"][0]["request_body"]
.pointer("/conversationState/currentMessage/userInputMessage/modelId"),
Some(&json!("claude-haiku-4.5"))
);
assert_eq!(
payload["data"]["response"]["content"][0]["text"],
json!("Hello from Kiro mapped model")
);
let direct_response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/test-model"))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-kiro",
"mode": "direct",
"model_name": "claude-haiku-4-5-20251001",
"mapped_model_name": "claude-haiku-4.5",
"api_format": "claude:messages",
"endpoint_id": "endpoint-kiro-cli",
"request_id": "provider-test-kiro-mapped-direct"
}))
.send()
.await
.expect("direct request should succeed");
assert_eq!(direct_response.status(), StatusCode::OK);
let direct_payload: serde_json::Value = direct_response
.json()
.await
.expect("json body should parse");
assert_eq!(direct_payload["success"], json!(true));
assert_eq!(
direct_payload["data"]["response"]["content"][0]["text"],
json!("Hello from Kiro mapped model")
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_provider_query_test_model_failover_for_kiro_locally() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
let payload = if plan.key_id == "key-kiro-first" {
json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 429,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"message": "too many requests"
}
},
"telemetry": {
"elapsed_ms": 11
}
})
} else {
json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/vnd.amazon.eventstream"
},
"body": {
"body_bytes_b64": base64::engine::general_purpose::STANDARD.encode(
[
encode_kiro_event_frame("assistantResponseEvent", json!({"content": "Recovered from failover"})),
encode_kiro_exception_frame("ContentLengthExceededException"),
]
.concat()
)
},
"telemetry": {
"elapsed_ms": 27
}
})
};
Json(payload)
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-kiro", "Kiro", 10);
provider.provider_type = "kiro".to_string();
let build_key = |id: &str| {
2026-04-29 09:25:19 +08:00
let mut key = sample_key(id, "provider-kiro", "claude:messages", "__placeholder__");
key.auth_type = "oauth".to_string();
key.encrypted_auth_config = Some(
aether_crypto::encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{
"provider_type":"kiro",
"auth_method":"idc",
"access_token":"cached-kiro-token",
"expires_at":4102444800,
"refresh_token":"rrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrr",
"machine_id":"123e4567-e89b-12d3-a456-426614174000",
"api_region":"us-east-1",
"client_id":"client-id",
"client_secret":"client-secret"
}"#,
)
.expect("auth config should encrypt"),
);
key
};
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![StoredProviderCatalogEndpoint::new(
"endpoint-kiro-cli".to_string(),
"provider-kiro".to_string(),
2026-04-29 09:25:19 +08:00
"claude:messages".to_string(),
Some("claude".to_string()),
2026-04-29 09:25:19 +08:00
Some("messages".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
"https://q.{region}.amazonaws.com".to_string(),
None,
None,
None,
None,
None,
None,
None,
)
.expect("endpoint transport should build")],
vec![build_key("key-kiro-first"), build_key("key-kiro-second")],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-query/test-model-failover"
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-kiro",
"mode": "direct",
"model_name": "claude-sonnet-4-upstream",
"failover_models": ["claude-sonnet-4-upstream"],
2026-04-29 09:25:19 +08:00
"api_format": "claude:messages",
"request_id": "provider-test-kiro"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(payload["total_candidates"], json!(2));
assert_eq!(payload["total_attempts"], json!(2));
assert_eq!(payload["candidate_summary"]["total_candidates"], json!(2));
assert_eq!(payload["candidate_summary"]["attempted"], json!(2));
assert_eq!(payload["candidate_summary"]["failed"], json!(1));
assert_eq!(payload["candidate_summary"]["success"], json!(1));
assert_eq!(
payload["candidate_summary"]["stop_reason"],
json!("first_success")
);
let attempts = payload["attempts"]
.as_array()
.expect("attempts should be an array");
assert_eq!(attempts.len(), 2);
assert_eq!(attempts[0]["status"], json!("failed"));
assert_eq!(attempts[0]["status_code"], json!(429));
assert_eq!(attempts[1]["status"], json!("success"));
assert_eq!(
payload["data"]["response"]["content"][0]["text"],
json!("Recovered from failover")
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_retries_kiro_failover_after_http_error_without_message() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
let payload = if plan.key_id == "key-kiro-first" {
json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 500,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {}
},
"telemetry": {
"elapsed_ms": 9
}
})
} else {
json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/vnd.amazon.eventstream"
},
"body": {
"body_bytes_b64": base64::engine::general_purpose::STANDARD.encode(
[
encode_kiro_event_frame("assistantResponseEvent", json!({"content": "Recovered from Kiro empty error"})),
encode_kiro_exception_frame("ContentLengthExceededException"),
]
.concat()
)
},
"telemetry": {
"elapsed_ms": 21
}
})
};
Json(payload)
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-kiro", "Kiro", 10);
provider.provider_type = "kiro".to_string();
let build_key = |id: &str| {
2026-04-29 09:25:19 +08:00
let mut key = sample_key(id, "provider-kiro", "claude:messages", "__placeholder__");
key.auth_type = "oauth".to_string();
key.encrypted_auth_config = Some(
aether_crypto::encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{
"provider_type":"kiro",
"auth_method":"idc",
"access_token":"cached-kiro-token",
"expires_at":4102444800,
"refresh_token":"rrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrr",
"machine_id":"123e4567-e89b-12d3-a456-426614174000",
"api_region":"us-east-1",
"client_id":"client-id",
"client_secret":"client-secret"
}"#,
)
.expect("auth config should encrypt"),
);
key
};
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![StoredProviderCatalogEndpoint::new(
"endpoint-kiro-cli".to_string(),
"provider-kiro".to_string(),
2026-04-29 09:25:19 +08:00
"claude:messages".to_string(),
Some("claude".to_string()),
2026-04-29 09:25:19 +08:00
Some("messages".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
"https://q.{region}.amazonaws.com".to_string(),
None,
None,
None,
None,
None,
None,
None,
)
.expect("endpoint transport should build")],
vec![build_key("key-kiro-first"), build_key("key-kiro-second")],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-query/test-model-failover"
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-kiro",
"mode": "direct",
"model_name": "claude-sonnet-4-upstream",
"failover_models": ["claude-sonnet-4-upstream"],
2026-04-29 09:25:19 +08:00
"api_format": "claude:messages",
"request_id": "provider-test-kiro-empty-error"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(payload["total_attempts"], json!(2));
let attempts = payload["attempts"]
.as_array()
.expect("attempts should be an array");
assert_eq!(attempts[0]["status"], json!("failed"));
assert_eq!(attempts[0]["status_code"], json!(500));
assert_eq!(attempts[1]["status"], json!("success"));
assert_eq!(
payload["data"]["response"]["content"][0]["text"],
json!("Recovered from Kiro empty error")
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_handles_non_kiro_multi_model_failover_locally() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
let payload = if plan.model_name.as_deref() == Some("gpt-4.1") {
json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 500,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"error": {
"message": "primary model failed"
}
}
},
"telemetry": {
"elapsed_ms": 8
}
})
} else {
json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"id": "chatcmpl-multi-model",
"choices": [{
"message": {
"role": "assistant",
"content": "Recovered with fallback model"
}
}]
}
},
"telemetry": {
"elapsed_ms": 12
}
})
};
Json(payload)
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-openai", "OpenAI", 10);
provider.provider_type = "openai".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-openai-chat",
"provider-openai",
"openai:chat",
"https://api.openai.example",
)],
vec![sample_key(
"key-openai-primary",
"provider-openai",
"openai:chat",
"sk-test-primary",
)],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-query/test-model-failover"
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-openai",
"failover_models": ["gpt-4.1", "gpt-4o-mini"],
"api_format": "openai:chat"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(payload["total_attempts"], json!(2));
let attempts = payload["attempts"]
.as_array()
.expect("attempts should be an array");
assert_eq!(attempts[0]["effective_model"], json!("gpt-4.1"));
assert_eq!(attempts[1]["effective_model"], json!("gpt-4o-mini"));
assert_eq!(
payload["data"]["response"]["choices"][0]["message"]["content"],
json!("Recovered with fallback model")
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_handles_openai_responses_test_model_locally() {
let prompt = "Tell me whether the CLI request preserved this prompt.";
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
assert_eq!(plan.provider_id, "provider-openai");
assert_eq!(plan.endpoint_id, "endpoint-openai-cli");
assert_eq!(plan.key_id, "key-openai-cli");
assert_eq!(plan.provider_api_format, "openai:responses");
assert_eq!(plan.url, "https://tiger.bookapi.cc/codex/responses");
assert!(plan.stream);
assert_eq!(
plan.headers.get("authorization").map(String::as_str),
Some("Bearer sk-test-cli")
);
assert_eq!(
plan.headers.get("x-stainless-runtime").map(String::as_str),
Some("node")
);
assert_eq!(
plan.body
.json_body
.as_ref()
.and_then(|body| body.get("model")),
Some(&json!("gpt-5.4-mini"))
);
assert_eq!(
plan.body
.json_body
.as_ref()
.and_then(|body| body.get("stream")),
Some(&json!(true))
);
assert!(plan
.body
.json_body
.as_ref()
.and_then(|body| body.get("input"))
.is_some());
assert_eq!(
plan.body
.json_body
.as_ref()
.and_then(|body| body.get("input"))
.and_then(|input| input.as_array())
.and_then(|items| items.first())
.and_then(|item| item.get("type"))
.and_then(|value| value.as_str()),
Some("message")
);
assert_eq!(
plan.body
.json_body
.as_ref()
.and_then(|body| body.get("input"))
.and_then(|input| input.as_array())
.and_then(|items| items.first())
.and_then(|item| item.get("content"))
.and_then(|content| content.as_array())
.and_then(|parts| parts.first())
.and_then(|part| part.get("text"))
.and_then(|value| value.as_str()),
Some(prompt)
);
assert_eq!(
plan.body
.json_body
.as_ref()
.and_then(|body| body.get("instructions")),
Some(&json!(""))
);
assert_eq!(
plan.body
.json_body
.as_ref()
.and_then(|body| body.get("store")),
Some(&json!(false))
);
assert!(plan
.body
.json_body
.as_ref()
.and_then(|body| body.get("prompt_cache_key"))
.is_some());
Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"id": "chatcmpl-openai-cli-test-model",
"object": "chat.completion",
"choices": [{
"message": {
"role": "assistant",
"content": "Hello from OpenAI Responses"
}
}]
}
},
"telemetry": {
"elapsed_ms": 17
}
}))
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-openai", "OpenAI", 10);
provider.provider_type = "codex".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-openai-cli",
"provider-openai",
"openai:responses",
"https://tiger.bookapi.cc/codex",
)],
vec![sample_key(
"key-openai-cli",
"provider-openai",
"openai:responses",
"sk-test-cli",
)],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/test-model"))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-openai",
"model": "gpt-5.4-mini",
"api_format": "openai:responses",
"message": prompt,
"request_headers": {
"x-stainless-runtime": "node"
}
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(
payload["data"]["response"]["choices"][0]["message"]["content"],
json!("Hello from OpenAI Responses")
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_handles_openai_image_test_model_locally() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
assert_eq!(plan.provider_id, "provider-openai");
assert_eq!(plan.endpoint_id, "endpoint-openai-image");
assert_eq!(plan.key_id, "key-openai-image");
assert_eq!(plan.client_api_format, "openai:image");
assert_eq!(plan.provider_api_format, "openai:image");
assert_eq!(plan.model_name.as_deref(), Some("gpt-image-1"));
assert_eq!(plan.url, "https://api.openai.example/v1/responses");
assert!(plan.stream);
assert_eq!(
plan.headers.get("authorization").map(String::as_str),
Some("Bearer sk-test-image")
);
assert_eq!(
plan.body
.json_body
.as_ref()
.and_then(|body| body.get("model")),
Some(&json!(crate::ai_serving::CODEX_OPENAI_IMAGE_INTERNAL_MODEL))
);
assert_eq!(
plan.body
.json_body
.as_ref()
.and_then(|body| body.get("input"))
.and_then(|input| input.as_array())
.and_then(|items| items.first())
.and_then(|item| item.get("content"))
.and_then(|value| value.as_str()),
Some("Draw a small blue square")
);
Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "text/event-stream"
},
"body": {
"body_bytes_b64": base64::engine::general_purpose::STANDARD.encode(
concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"created_at\":1776839946}}\n\n",
"event: response.output_item.done\n",
"data: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"type\":\"image_generation_call\",\"output_format\":\"png\",\"revised_prompt\":\"revised prompt\",\"result\":\"aGVsbG8=\"}}\n\n",
"event: response.completed\n",
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_img_123\",\"model\":\"gpt-image-1\",\"status\":\"completed\",\"tool_usage\":{\"image_gen\":{\"input_tokens\":171,\"output_tokens\":1372,\"total_tokens\":1543}}}}\n\n"
)
.as_bytes()
)
},
"telemetry": {
"elapsed_ms": 19
}
}))
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-openai", "OpenAI", 10);
provider.provider_type = "codex".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-openai-image",
"provider-openai",
"openai:image",
"https://api.openai.example",
)],
vec![sample_key(
"key-openai-image",
"provider-openai",
"openai:image",
"sk-test-image",
)],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/test-model"))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-openai",
"model": "gpt-image-1",
"api_format": "openai:image",
"message": "Draw a small blue square"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(
payload["data"]["response"]["data"][0]["b64_json"],
json!("aGVsbG8=")
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_reports_transport_unsupported_reason_for_non_kiro_provider() {
let mut provider = sample_provider("provider-antigravity", "Antigravity", 10);
provider.provider_type = "antigravity".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-antigravity-gemini",
"provider-antigravity",
2026-04-29 09:25:19 +08:00
"gemini:generate_content",
"https://cloudcode-pa.googleapis.com",
)],
vec![sample_key(
"key-antigravity-gemini",
"provider-antigravity",
2026-04-29 09:25:19 +08:00
"gemini:generate_content",
"sk-test-antigravity",
)],
));
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/test-model"))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-antigravity",
"model": "gemini-2.5-pro",
2026-04-29 09:25:19 +08:00
"api_format": "gemini:generate_content"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(false));
assert_eq!(
payload["error"],
json!(
"Rust local provider-query model test cannot execute endpoint format gemini:generate_content (transport_antigravity_auth_unsupported)"
)
);
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_handles_antigravity_endpoint_test_model_locally() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
assert_eq!(plan.provider_id, "provider-antigravity");
assert_eq!(plan.endpoint_id, "endpoint-antigravity-gemini");
assert_eq!(plan.key_id, "key-antigravity-gemini");
assert_eq!(plan.provider_api_format, "gemini:generate_content");
assert_eq!(
plan.url,
"https://antigravity.googleapis.com/v1internal:generateContent"
);
assert_eq!(plan.stream, false);
assert_eq!(
plan.body
.json_body
.as_ref()
.and_then(|body| body.get("project"))
.and_then(|value| value.as_str()),
Some("project-ant-123")
);
assert_eq!(
plan.body
.json_body
.as_ref()
.and_then(|body| body.get("requestType"))
.and_then(|value| value.as_str()),
Some("endpoint_test")
);
assert_eq!(
plan.body
.json_body
.as_ref()
.and_then(|body| body.get("model"))
.and_then(|value| value.as_str()),
Some("gemini-2.5-pro")
);
assert_eq!(
plan.body
.json_body
.as_ref()
.and_then(|body| body.get("request"))
.and_then(|request| request.get("contents"))
.and_then(|contents| contents.as_array())
.and_then(|items| items.first())
.and_then(|item| item.get("parts"))
.and_then(|parts| parts.as_array())
.and_then(|parts| parts.first())
.and_then(|part| part.get("text"))
.and_then(|value| value.as_str()),
Some("Say hello")
);
assert_eq!(
plan.headers.get("x-client-name").map(String::as_str),
Some("antigravity")
);
assert!(plan.headers.contains_key("x-goog-api-client"));
Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"response": {
"candidates": [{
"content": {
"parts": [
{"text": "Hello from Antigravity EndpointTest"}
],
"role": "model"
},
"finishReason": "STOP",
"index": 0
}],
"modelVersion": "claude-sonnet-4-5",
"usageMetadata": {
"promptTokenCount": 2,
"candidatesTokenCount": 3,
"totalTokenCount": 5
}
},
"responseId": "resp-antigravity-test-123"
}
},
"telemetry": {
"elapsed_ms": 23
}
}))
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-antigravity", "Antigravity", 10);
provider.provider_type = "antigravity".to_string();
let mut key = sample_key(
"key-antigravity-gemini",
"provider-antigravity",
"gemini:generate_content",
"cached-antigravity-token",
);
key.auth_type = "oauth".to_string();
key.encrypted_auth_config = Some(
aether_crypto::encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{
"provider_type":"antigravity",
"project_id":"project-ant-123",
"client_version":"1.2.3",
"session_id":"sess-ant-123",
"refresh_token":"rt-ant-123"
}"#,
)
.expect("auth config should encrypt"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-antigravity-gemini",
"provider-antigravity",
"gemini:generate_content",
"https://antigravity.googleapis.com",
)],
vec![key],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/test-model"))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-antigravity",
"model": "gemini-2.5-pro",
"api_format": "gemini:generate_content",
"message": "Say hello"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(
payload["data"]["response"]["choices"][0]["message"]["content"],
json!("Hello from Antigravity EndpointTest")
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_prefers_supported_non_kiro_endpoint_when_api_format_is_omitted() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
assert_eq!(plan.endpoint_id, "endpoint-openai-chat");
assert_eq!(plan.provider_api_format, "openai:chat");
Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"id": "chatcmpl-preferred-endpoint",
"choices": [{
"message": {
"role": "assistant",
"content": "Selected supported endpoint"
}
}]
}
},
"telemetry": {
"elapsed_ms": 10
}
}))
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-openai", "OpenAI", 10);
provider.provider_type = "openai".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![
sample_endpoint(
"endpoint-openai-cli",
"provider-openai",
"openai:responses",
"https://api.openai.example",
),
sample_endpoint(
"endpoint-openai-chat",
"provider-openai",
"openai:chat",
"https://api.openai.example",
),
],
vec![
sample_key(
"key-openai-cli",
"provider-openai",
"openai:responses",
"sk-test-cli",
),
sample_key(
"key-openai-chat",
"provider-openai",
"openai:chat",
"sk-test-chat",
),
],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/test-model"))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-openai",
"model": "gpt-5.4-mini"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(
payload["data"]["response"]["choices"][0]["message"]["content"],
json!("Selected supported endpoint")
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_prefers_transport_supported_non_kiro_endpoint_when_api_format_is_omitted() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
assert_eq!(plan.endpoint_id, "endpoint-openai-chat-supported");
assert_eq!(plan.provider_api_format, "openai:chat");
assert_eq!(
plan.headers.get("content-type").map(String::as_str),
Some("application/json")
);
Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"id": "chatcmpl-supported-transport",
"choices": [{
"message": {
"role": "assistant",
"content": "Selected locally supported endpoint"
}
}]
}
},
"telemetry": {
"elapsed_ms": 12
}
}))
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-openai", "OpenAI", 10);
provider.provider_type = "openai".to_string();
let mut unsupported_endpoint = sample_endpoint(
"endpoint-openai-chat-unsupported",
"provider-openai",
"openai:chat",
"https://api.openai.example",
);
unsupported_endpoint.header_rules = Some(json!({"invalid": true}));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![
unsupported_endpoint,
sample_endpoint(
"endpoint-openai-chat-supported",
"provider-openai",
"openai:chat",
"https://api.openai.example",
),
],
vec![sample_key(
"key-openai-chat",
"provider-openai",
"openai:chat",
"sk-test-chat",
)],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/test-model"))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-openai",
"model": "gpt-5.4-mini"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(
payload["data"]["response"]["choices"][0]["message"]["content"],
json!("Selected locally supported endpoint")
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_prefers_supported_non_kiro_endpoint_with_compatible_key_when_api_format_is_omitted(
) {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
assert_eq!(plan.endpoint_id, "endpoint-openai-chat");
assert_eq!(plan.provider_api_format, "openai:chat");
assert_eq!(plan.key_id, "key-openai-chat");
Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"id": "chatcmpl-key-compatible-endpoint",
"choices": [{
"message": {
"role": "assistant",
"content": "Selected endpoint with compatible key"
}
}]
}
},
"telemetry": {
"elapsed_ms": 11
}
}))
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-openai", "OpenAI", 10);
provider.provider_type = "openai".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![
sample_endpoint(
"endpoint-gemini-chat",
"provider-openai",
2026-04-29 09:25:19 +08:00
"gemini:generate_content",
"https://api.gemini.example",
),
sample_endpoint(
"endpoint-openai-chat",
"provider-openai",
"openai:chat",
"https://api.openai.example",
),
],
vec![sample_key(
"key-openai-chat",
"provider-openai",
"openai:chat",
"sk-test-chat",
)],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/test-model"))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-openai",
"model": "gpt-5.4-mini"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(
payload["data"]["response"]["choices"][0]["message"]["content"],
json!("Selected endpoint with compatible key")
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_uses_compatible_cli_endpoint_when_api_format_is_omitted() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
assert_eq!(plan.endpoint_id, "endpoint-openai-cli");
assert_eq!(plan.provider_api_format, "openai:responses");
assert_eq!(plan.key_id, "key-openai-cli");
Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"id": "chatcmpl-cli-only-endpoint",
"choices": [{
"message": {
"role": "assistant",
"content": "Selected compatible CLI endpoint"
}
}]
}
},
"telemetry": {
"elapsed_ms": 13
}
}))
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-openai", "OpenAI", 10);
provider.provider_type = "openai".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![
sample_endpoint(
"endpoint-openai-chat",
"provider-openai",
"openai:chat",
"https://api.openai.example",
),
sample_endpoint(
"endpoint-openai-cli",
"provider-openai",
"openai:responses",
"https://api.openai.example",
),
],
vec![sample_key(
"key-openai-cli",
"provider-openai",
"openai:responses",
"sk-test-cli",
)],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/test-model"))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-openai",
"model": "gpt-5.4-mini"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(
payload["data"]["response"]["choices"][0]["message"]["content"],
json!("Selected compatible CLI endpoint")
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_uses_runnable_cli_endpoint_after_chat_preference_when_api_format_is_omitted() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
assert_eq!(plan.endpoint_id, "endpoint-openai-cli-runnable");
assert_eq!(plan.provider_api_format, "openai:responses");
assert_eq!(plan.key_id, "key-openai-shared");
Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"id": "chatcmpl-cli-runnable-after-chat-preference",
"choices": [{
"message": {
"role": "assistant",
"content": "Selected runnable CLI endpoint after unsupported chat"
}
}]
}
},
"telemetry": {
"elapsed_ms": 18
}
}))
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-openai", "OpenAI", 10);
provider.provider_type = "openai".to_string();
let mut unsupported_chat_endpoint = sample_endpoint(
"endpoint-openai-chat-unsupported",
"provider-openai",
"openai:chat",
"https://api.openai.example",
);
unsupported_chat_endpoint.header_rules = Some(json!({"invalid": true}));
let cli_endpoint = sample_endpoint(
"endpoint-openai-cli-runnable",
"provider-openai",
"openai:responses",
"https://api.openai.example",
);
let mut shared_key = sample_key(
"key-openai-shared",
"provider-openai",
"openai:chat",
"sk-test-shared",
);
shared_key.api_formats = Some(json!(["openai:chat", "openai:responses"]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![unsupported_chat_endpoint, cli_endpoint],
vec![shared_key],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/test-model"))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-openai",
"model": "gpt-5.4-mini"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(
payload["data"]["response"]["choices"][0]["message"]["content"],
json!("Selected runnable CLI endpoint after unsupported chat")
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_handles_openai_responses_test_model_failover_locally() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
assert_eq!(plan.provider_id, "provider-openai");
assert_eq!(plan.endpoint_id, "endpoint-openai-cli");
assert_eq!(plan.key_id, "key-openai-cli");
assert_eq!(plan.provider_api_format, "openai:responses");
assert_eq!(plan.model_name.as_deref(), Some("gpt-5.4-mini"));
Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"id": "chatcmpl-openai-cli-failover",
"choices": [{
"message": {
"role": "assistant",
"content": "OpenAI Responses failover path succeeded"
}
}]
}
},
"telemetry": {
"elapsed_ms": 15
}
}))
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-openai", "OpenAI", 10);
provider.provider_type = "openai".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-openai-cli",
"provider-openai",
"openai:responses",
"https://api.openai.example",
)],
vec![sample_key(
"key-openai-cli",
"provider-openai",
"openai:responses",
"sk-test-cli",
)],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-query/test-model-failover"
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-openai",
"failover_models": ["gpt-5.4-mini"],
"api_format": "openai:responses"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(payload["total_attempts"], json!(1));
assert_eq!(
payload["data"]["response"]["choices"][0]["message"]["content"],
json!("OpenAI Responses failover path succeeded")
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_handles_claude_cli_test_model_locally() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
assert_eq!(plan.provider_id, "provider-claude");
assert_eq!(plan.endpoint_id, "endpoint-claude-cli");
assert_eq!(plan.key_id, "key-claude-cli");
2026-04-29 09:25:19 +08:00
assert_eq!(plan.provider_api_format, "claude:messages");
assert_eq!(plan.url, "https://api.anthropic.example/v1/messages");
assert_eq!(plan.model_name.as_deref(), Some("claude-sonnet-4-5"));
Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"id": "chatcmpl-claude-cli-test-model",
"choices": [{
"message": {
"role": "assistant",
"content": "Hello from Claude CLI"
}
}]
}
},
"telemetry": {
"elapsed_ms": 14
}
}))
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-claude", "Claude", 10);
provider.provider_type = "anthropic".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-claude-cli",
"provider-claude",
2026-04-29 09:25:19 +08:00
"claude:messages",
"https://api.anthropic.example",
)],
vec![sample_key(
"key-claude-cli",
"provider-claude",
2026-04-29 09:25:19 +08:00
"claude:messages",
"sk-test-claude-cli",
)],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/test-model"))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-claude",
"model": "claude-sonnet-4-5",
2026-04-29 09:25:19 +08:00
"api_format": "claude:messages"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(
payload["data"]["response"]["choices"][0]["message"]["content"],
json!("Hello from Claude CLI")
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_uses_compatible_claude_cli_endpoint_when_api_format_is_omitted() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
assert_eq!(plan.endpoint_id, "endpoint-claude-cli");
2026-04-29 09:25:19 +08:00
assert_eq!(plan.provider_api_format, "claude:messages");
assert_eq!(plan.key_id, "key-claude-cli");
Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"id": "chatcmpl-claude-cli-only-endpoint",
"choices": [{
"message": {
"role": "assistant",
"content": "Selected compatible Claude CLI endpoint"
}
}]
}
},
"telemetry": {
"elapsed_ms": 12
}
}))
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-claude", "Claude", 10);
provider.provider_type = "anthropic".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-claude-cli",
"provider-claude",
2026-04-29 09:25:19 +08:00
"claude:messages",
"https://api.anthropic.example",
)],
vec![sample_key(
"key-claude-cli",
"provider-claude",
2026-04-29 09:25:19 +08:00
"claude:messages",
"sk-test-claude-cli",
)],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/test-model"))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-claude",
"model": "claude-sonnet-4-5"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(
payload["data"]["response"]["choices"][0]["message"]["content"],
json!("Selected compatible Claude CLI endpoint")
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_handles_claude_cli_test_model_failover_locally() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
assert_eq!(plan.provider_id, "provider-claude");
assert_eq!(plan.endpoint_id, "endpoint-claude-cli");
assert_eq!(plan.key_id, "key-claude-cli");
2026-04-29 09:25:19 +08:00
assert_eq!(plan.provider_api_format, "claude:messages");
assert_eq!(plan.model_name.as_deref(), Some("claude-sonnet-4-5"));
Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"id": "chatcmpl-claude-cli-failover",
"choices": [{
"message": {
"role": "assistant",
"content": "Claude CLI failover path succeeded"
}
}]
}
},
"telemetry": {
"elapsed_ms": 16
}
}))
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-claude", "Claude", 10);
provider.provider_type = "anthropic".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-claude-cli",
"provider-claude",
2026-04-29 09:25:19 +08:00
"claude:messages",
"https://api.anthropic.example",
)],
vec![sample_key(
"key-claude-cli",
"provider-claude",
2026-04-29 09:25:19 +08:00
"claude:messages",
"sk-test-claude-cli",
)],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-query/test-model-failover"
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-claude",
"failover_models": ["claude-sonnet-4-5"],
2026-04-29 09:25:19 +08:00
"api_format": "claude:messages"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(payload["total_attempts"], json!(1));
assert_eq!(
payload["data"]["response"]["choices"][0]["message"]["content"],
json!("Claude CLI failover path succeeded")
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_handles_gemini_cli_test_model_locally() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
assert_eq!(plan.provider_id, "provider-gemini");
assert_eq!(plan.endpoint_id, "endpoint-gemini-cli");
assert_eq!(plan.key_id, "key-gemini-cli");
2026-04-29 09:25:19 +08:00
assert_eq!(plan.provider_api_format, "gemini:generate_content");
assert_eq!(
plan.url,
"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-pro:generateContent"
);
assert_eq!(plan.model_name.as_deref(), Some("gemini-2.5-pro"));
Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"id": "chatcmpl-gemini-cli-test-model",
"choices": [{
"message": {
"role": "assistant",
"content": "Hello from Gemini CLI"
}
}]
}
},
"telemetry": {
"elapsed_ms": 19
}
}))
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-gemini", "Gemini", 10);
provider.provider_type = "google".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-gemini-cli",
"provider-gemini",
2026-04-29 09:25:19 +08:00
"gemini:generate_content",
"https://generativelanguage.googleapis.com",
)],
vec![sample_key(
"key-gemini-cli",
"provider-gemini",
2026-04-29 09:25:19 +08:00
"gemini:generate_content",
"sk-test-gemini-cli",
)],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/test-model"))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-gemini",
"model": "gemini-2.5-pro",
2026-04-29 09:25:19 +08:00
"api_format": "gemini:generate_content"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(
payload["data"]["response"]["choices"][0]["message"]["content"],
json!("Hello from Gemini CLI")
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_handles_gemini_cli_test_model_with_oauth_header_fallback() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
assert_eq!(plan.provider_id, "provider-gemini");
assert_eq!(plan.endpoint_id, "endpoint-gemini-cli");
assert_eq!(plan.key_id, "key-gemini-cli");
2026-04-29 09:25:19 +08:00
assert_eq!(plan.provider_api_format, "gemini:generate_content");
assert_eq!(
plan.url,
"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-pro:generateContent"
);
assert_eq!(
plan.headers.get("authorization").map(String::as_str),
Some("Bearer cached-gemini-cli-token")
);
assert!(!plan.headers.contains_key("x-goog-api-key"));
Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"id": "chatcmpl-gemini-cli-test-model",
"choices": [{
"message": {
"role": "assistant",
"content": "Hello from Gemini CLI"
}
}]
}
},
"telemetry": {
"elapsed_ms": 19
}
}))
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-gemini", "Gemini", 10);
provider.provider_type = "gemini_cli".to_string();
let mut key = sample_key(
"key-gemini-cli",
"provider-gemini",
2026-04-29 09:25:19 +08:00
"gemini:generate_content",
"cached-gemini-cli-token",
);
key.auth_type = "oauth".to_string();
key.encrypted_auth_config = Some(
aether_crypto::encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"gemini_cli"}"#,
)
.expect("auth config should encrypt"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-gemini-cli",
"provider-gemini",
2026-04-29 09:25:19 +08:00
"gemini:generate_content",
"https://generativelanguage.googleapis.com",
)],
vec![key],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/test-model"))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-gemini",
"model": "gemini-2.5-pro",
2026-04-29 09:25:19 +08:00
"api_format": "gemini:generate_content"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(
payload["data"]["response"]["choices"][0]["message"]["content"],
json!("Hello from Gemini CLI")
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_uses_compatible_gemini_cli_endpoint_when_api_format_is_omitted() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
assert_eq!(plan.endpoint_id, "endpoint-gemini-cli");
2026-04-29 09:25:19 +08:00
assert_eq!(plan.provider_api_format, "gemini:generate_content");
assert_eq!(plan.key_id, "key-gemini-cli");
Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"id": "chatcmpl-gemini-cli-only-endpoint",
"choices": [{
"message": {
"role": "assistant",
"content": "Selected compatible Gemini CLI endpoint"
}
}]
}
},
"telemetry": {
"elapsed_ms": 21
}
}))
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-gemini", "Gemini", 10);
provider.provider_type = "google".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-gemini-cli",
"provider-gemini",
2026-04-29 09:25:19 +08:00
"gemini:generate_content",
"https://generativelanguage.googleapis.com",
)],
vec![sample_key(
"key-gemini-cli",
"provider-gemini",
2026-04-29 09:25:19 +08:00
"gemini:generate_content",
"sk-test-gemini-cli",
)],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/test-model"))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-gemini",
"model": "gemini-2.5-pro"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(
payload["data"]["response"]["choices"][0]["message"]["content"],
json!("Selected compatible Gemini CLI endpoint")
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_handles_gemini_cli_test_model_failover_locally() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
assert_eq!(plan.provider_id, "provider-gemini");
assert_eq!(plan.endpoint_id, "endpoint-gemini-cli");
assert_eq!(plan.key_id, "key-gemini-cli");
2026-04-29 09:25:19 +08:00
assert_eq!(plan.provider_api_format, "gemini:generate_content");
assert_eq!(plan.model_name.as_deref(), Some("gemini-2.5-pro"));
Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"id": "chatcmpl-gemini-cli-failover",
"choices": [{
"message": {
"role": "assistant",
"content": "Gemini CLI failover path succeeded"
}
}]
}
},
"telemetry": {
"elapsed_ms": 23
}
}))
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-gemini", "Gemini", 10);
provider.provider_type = "google".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-gemini-cli",
"provider-gemini",
2026-04-29 09:25:19 +08:00
"gemini:generate_content",
"https://generativelanguage.googleapis.com",
)],
vec![sample_key(
"key-gemini-cli",
"provider-gemini",
2026-04-29 09:25:19 +08:00
"gemini:generate_content",
"sk-test-gemini-cli",
)],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-query/test-model-failover"
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-gemini",
"failover_models": ["gemini-2.5-pro"],
2026-04-29 09:25:19 +08:00
"api_format": "gemini:generate_content"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(payload["total_attempts"], json!(1));
assert_eq!(
payload["data"]["response"]["choices"][0]["message"]["content"],
json!("Gemini CLI failover path succeeded")
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_provider_query_test_model_failover_with_single_model_name_alias() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
assert_eq!(plan.model_name.as_deref(), Some("gpt-4.1"));
Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"id": "chatcmpl-alias",
"choices": [{
"message": {
"role": "assistant",
"content": "Alias path succeeded"
}
}]
}
},
"telemetry": {
"elapsed_ms": 9
}
}))
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-openai", "OpenAI", 10);
provider.provider_type = "openai".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-openai-chat",
"provider-openai",
"openai:chat",
"https://api.openai.example",
)],
vec![sample_key(
"key-openai-alias",
"provider-openai",
"openai:chat",
"sk-test-alias",
)],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-query/test-model-failover"
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-openai",
"failover_models": ["gpt-4.1"],
"api_format": "openai:chat"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(payload["model"], json!("gpt-4.1"));
assert_eq!(payload["total_attempts"], json!(1));
assert_eq!(
payload["data"]["response"]["choices"][0]["message"]["content"],
json!("Alias path succeeded")
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_retries_non_kiro_failover_after_http_error_without_message() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
let auth = plan
.headers
.get("authorization")
.map(String::as_str)
.unwrap_or_default()
.to_string();
let payload = if auth == "Bearer sk-test-first" {
json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 500,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {}
},
"telemetry": {
"elapsed_ms": 7
}
})
} else {
json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"id": "chatcmpl-retry",
"choices": [{
"message": {
"role": "assistant",
"content": "Recovered after empty error"
}
}]
}
},
"telemetry": {
"elapsed_ms": 13
}
})
};
Json(payload)
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-openai", "OpenAI", 10);
provider.provider_type = "openai".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-openai-chat",
"provider-openai",
"openai:chat",
"https://api.openai.example",
)],
vec![
sample_key(
"key-openai-first",
"provider-openai",
"openai:chat",
"sk-test-first",
),
sample_key(
"key-openai-second",
"provider-openai",
"openai:chat",
"sk-test-second",
),
],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-query/test-model-failover"
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-openai",
"failover_models": ["gpt-4.1"],
"api_format": "openai:chat"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(payload["total_attempts"], json!(2));
let attempts = payload["attempts"]
.as_array()
.expect("attempts should be an array");
assert_eq!(attempts[0]["status"], json!("failed"));
assert_eq!(attempts[0]["status_code"], json!(500));
assert_eq!(attempts[1]["status"], json!("success"));
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_retries_non_kiro_failover_after_success_status_without_body() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
let auth = plan
.headers
.get("authorization")
.map(String::as_str)
.unwrap_or_default()
.to_string();
let payload = if auth == "Bearer sk-test-first" {
json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {},
"telemetry": {
"elapsed_ms": 6
}
})
} else {
json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"id": "chatcmpl-retry-empty-body",
"choices": [{
"message": {
"role": "assistant",
"content": "Recovered after empty success body"
}
}]
}
},
"telemetry": {
"elapsed_ms": 14
}
})
};
Json(payload)
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-openai", "OpenAI", 10);
provider.provider_type = "openai".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-openai-chat",
"provider-openai",
"openai:chat",
"https://api.openai.example",
)],
vec![
sample_key(
"key-openai-first",
"provider-openai",
"openai:chat",
"sk-test-first",
),
sample_key(
"key-openai-second",
"provider-openai",
"openai:chat",
"sk-test-second",
),
],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-query/test-model-failover"
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-openai",
"failover_models": ["gpt-4.1"],
"api_format": "openai:chat"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(payload["total_attempts"], json!(2));
assert_eq!(
payload["data"]["response"]["choices"][0]["message"]["content"],
json!("Recovered after empty success body")
);
let attempts = payload["attempts"]
.as_array()
.expect("attempts should be an array");
assert_eq!(attempts[0]["status"], json!("failed"));
assert_eq!(attempts[0]["status_code"], json!(200));
assert_eq!(
attempts[0]["error_message"],
json!("Provider returned HTTP 200 without a model-test response body")
);
assert_eq!(attempts[1]["status"], json!("success"));
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_preserves_non_success_status_for_test_model_local_wrapper() {
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/test-model"))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-does-not-exist",
"model": "gpt-4.1"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::NOT_FOUND);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["detail"], json!("Provider not found"));
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_rejects_admin_provider_query_invalid_json_body() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/provider-query/models",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/models"))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.header(reqwest::header::CONTENT_TYPE, "application/json")
.body("{")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["detail"], json!("Invalid JSON request body"));
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_rejects_admin_provider_query_test_model_without_provider_id() {
assert_admin_provider_query_route(
"/api/admin/provider-query/test-model",
json!({ "model": "gpt-4.1" }),
StatusCode::BAD_REQUEST,
|payload| {
assert_eq!(payload["detail"], json!("provider_id is required"));
},
)
.await;
}
#[tokio::test]
async fn gateway_rejects_admin_provider_query_test_model_without_model() {
assert_admin_provider_query_route(
"/api/admin/provider-query/test-model",
json!({ "provider_id": "provider-openai" }),
StatusCode::BAD_REQUEST,
|payload| {
assert_eq!(payload["detail"], json!("model is required"));
},
)
.await;
}
#[tokio::test]
async fn gateway_rejects_admin_provider_query_test_model_failover_without_provider_id() {
assert_admin_provider_query_route(
"/api/admin/provider-query/test-model-failover",
json!({ "failover_models": ["gpt-4.1"] }),
StatusCode::BAD_REQUEST,
|payload| {
assert_eq!(payload["detail"], json!("provider_id is required"));
},
)
.await;
}
#[tokio::test]
async fn gateway_rejects_admin_provider_query_test_model_failover_without_models() {
assert_admin_provider_query_route(
"/api/admin/provider-query/test-model-failover",
json!({ "provider_id": "provider-openai", "failover_models": [] }),
StatusCode::BAD_REQUEST,
|payload| {
assert_eq!(
payload["detail"],
json!("failover_models should not be empty")
);
},
)
.await;
}