i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 09:20:22 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(kiro,pool,model): 对齐 Kiro 管理链路并修复全局模型删除行为 (#305)

Browse Source 
* feat(pool): 号池支持跳过额度耗尽账号

- 新增 pool_advanced.skip_exhausted_accounts 开关及高级设置 UI, 默认关闭并兼容旧配置
- 为 Codex/Kiro 增加额度耗尽判定, 接入请求侧候选跳过并新增 account_quota_exhausted skip reason
- 号池列表将额度耗尽账号标记为 blocked/额度耗尽, 并补充前后端相关测试

* fix(kiro): 对齐账号管理与 provider-query 的 Rust 行为

- 修复 Kiro 单条导入误走 import-refresh-token 的前端分流, 并为误用路径返回明确错误提示
- 为 Kiro 导入与本地请求链补齐 bearer 兼容, 同步放开账号启停等 Key 更新操作的 auth_type 校验
- 实现 Kiro provider-query 本地模型测试与 failover 执行链, 并修复结果弹窗在无 trace 时无法展示 attempts/响应体的问题

* fix(model): 删除全局模型时级联清理关联提供商模型

- 对齐 Python 版本删除逻辑, GlobalModel 删除前先在事务内清理关联的 Provider Model 记录
- 修复已绑定 Provider 的模型在 Rust SQL 仓库下会被外键约束拦住、无法正常删除的问题
- 增加管理端回归测试, 覆盖绑定 Provider Model 的 GlobalModel 删除场景

* fix(kiro,ci): 恢复 Kiro OAuth 持久化并修复 Rust CI

* Fix oauth-managed provider key semantics

---------

Co-authored-by: fawney19 <elky0401@gmail.com>
This commit is contained in:
Entropy.Xu
2026-04-17 12:57:06 +08:00
committed by GitHub
parent 96a25d058b
commit ac1a126756
43 changed files with 2909 additions and 168 deletions
Download Patch File Download Diff File Expand all files Collapse all files

76
apps/aether-gateway/src/tests/control/admin/models/global.rs
View File

@@ -1027,6 +1027,82 @@ async fn gateway_batch_deletes_admin_global_models_locally_with_trusted_admin_pr
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_deletes_admin_global_model_with_bound_provider_models_locally() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/models/global/global-gpt-5",
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 global_model_repository = Arc::new(
InMemoryGlobalModelReadRepository::seed(Vec::new())
.with_admin_global_models(vec![sample_admin_global_model(
"global-gpt-5",
"gpt-5",
"GPT 5",
)])
.with_admin_provider_models(vec![sample_admin_provider_model(
"model-openai-gpt5",
"provider-openai",
"global-gpt-5",
"gpt-5-upstream",
)]),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::disabled()
.with_global_model_repository_for_tests(global_model_repository.clone()),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.delete(format!(
"{gateway_url}/api/admin/models/global/global-gpt-5"
))
.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")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::NO_CONTENT);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
let deleted = global_model_repository
.get_admin_global_model_by_id("global-gpt-5")
.await
.expect("model lookup should succeed");
assert!(deleted.is_none());
let provider_models = global_model_repository
.list_admin_provider_models(&AdminProviderModelListQuery {
provider_id: "provider-openai".to_string(),
is_active: None,
offset: 0,
limit: 20,
})
.await
.expect("provider models should read");
assert!(provider_models.is_empty());
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_assigns_admin_global_model_to_providers_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));

94
apps/aether-gateway/src/tests/control/admin/oauth.rs
View File

@@ -1792,6 +1792,51 @@ async fn gateway_imports_admin_provider_oauth_refresh_token_locally_with_trusted
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_rejects_kiro_single_refresh_token_import_with_clear_error() {
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![{
let mut provider = sample_provider("provider-kiro", "kiro", 10);
provider.provider_type = "kiro".to_string();
provider
}],
vec![sample_endpoint(
"endpoint-kiro-chat",
"provider-kiro",
"kiro:generateAssistantResponse",
"https://service.kiro.dev",
)],
Vec::new(),
));
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
));
let response = local_admin_provider_oauth_response(
&state,
http::Method::POST,
"/api/admin/provider-oauth/providers/provider-kiro/import-refresh-token",
Some(json!({
"refresh_token": "kiro-refresh-token"
})),
)
.await;
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
let payload: serde_json::Value = serde_json::from_slice(
&to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should read"),
)
.expect("json body should parse");
assert_eq!(
payload["detail"],
json!("Kiro 不支持单条 Refresh Token 导入，请使用批量导入或设备授权。")
);
}
#[tokio::test]
async fn gateway_imports_admin_provider_oauth_refresh_token_via_execution_runtime_proxy_node() {
let execution_plans = Arc::new(Mutex::new(Vec::<ExecutionPlan>::new()));
@@ -2382,20 +2427,37 @@ async fn gateway_batch_imports_admin_provider_oauth_kiro_via_execution_runtime_p
.lock()
.expect("mutex should lock")
.push(plan.clone());
assert_eq!(plan.request_id, "kiro_batch_refresh:social");
assert_eq!(plan.url, "https://oauth.example/refreshToken");
assert_eq!(
plan.proxy
.as_ref()
.and_then(|proxy| proxy.node_id.as_deref()),
Some("proxy-node-kiro-batch-runtime")
);
assert_eq!(
plan.headers
.get(EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER)
.map(String::as_str),
Some("true")
);
if plan.request_id == "kiro_batch_refresh:social" {
assert_eq!(plan.url, "https://oauth.example/refreshToken");
assert_eq!(
plan.headers
.get(EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER)
.map(String::as_str),
Some("true")
);
return Json(json!({
"request_id": plan.request_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"accessToken": sample_kiro_device_access_token("kiro-runtime@example.com"),
"refreshToken": "kiro-runtime-refresh-token-new",
"expiresIn": 1800,
}
}
}));
}
assert_eq!(plan.request_id, "kiro-quota:key-kiro-batch-runtime");
Json(json!({
"request_id": plan.request_id,
"status_code": 200,
@@ -2404,9 +2466,17 @@ async fn gateway_batch_imports_admin_provider_oauth_kiro_via_execution_runtime_p
},
"body": {
"json_body": {
"accessToken": sample_kiro_device_access_token("kiro-runtime@example.com"),
"refreshToken": "kiro-runtime-refresh-token-new",
"expiresIn": 1800,
"subscriptionInfo": {
"subscriptionTitle": "KIRO PRO+"
},
"usageBreakdownList": [{
"currentUsageWithPrecision": 5.0,
"usageLimitWithPrecision": 20.0,
"nextDateReset": 1_900_000_000u64
}],
"desktopUserInfo": {
"email": "kiro-runtime@example.com"
}
}
}
}))
@@ -2496,7 +2566,7 @@ async fn gateway_batch_imports_admin_provider_oauth_kiro_via_execution_runtime_p
{
let plans = execution_plans.lock().expect("mutex should lock");
assert_eq!(plans.len(), 1);
assert_eq!(plans.len(), 2);
}
let stored_key = provider_catalog_repository

215
apps/aether-gateway/src/tests/control/admin/pool.rs
View File

@@ -1109,6 +1109,162 @@ async fn gateway_formats_codex_quota_countdown_from_reset_after_seconds() {
);
}
#[tokio::test]
async fn gateway_marks_exhausted_codex_pool_key_as_blocked_when_flag_enabled() {
let mut provider = sample_provider("provider-codex", "codex", 10).with_transport_fields(
true,
false,
true,
None,
None,
None,
None,
None,
Some(json!({
"pool_advanced": {
"enabled": true,
"skip_exhausted_accounts": true
}
})),
);
provider.provider_type = "codex".to_string();
let mut key = sample_key(
"key-codex-exhausted",
"provider-codex",
"openai:cli",
"oauth-placeholder",
);
key.name = "codex exhausted".to_string();
key.auth_type = "oauth".to_string();
key.upstream_metadata = Some(json!({
"codex": {
"secondary_used_percent": 100.0,
"plan_type": "plus"
}
}));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
Vec::new(),
vec![key],
));
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
));
let response = local_admin_pool_response(
&state,
http::Method::GET,
"/api/admin/pool/provider-codex/keys?page=1&page_size=50&status=all",
None,
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = serde_json::from_slice(
&to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should read"),
)
.expect("json body should parse");
let keys = payload["keys"].as_array().expect("keys should be array");
assert_eq!(keys[0]["scheduling_status"], json!("blocked"));
assert_eq!(
keys[0]["scheduling_reason"],
json!("account_quota_exhausted")
);
assert_eq!(keys[0]["scheduling_label"], json!("额度耗尽"));
assert_eq!(
keys[0]["scheduling_reasons"][0],
json!({
"code": "account_quota_exhausted",
"label": "额度耗尽",
"blocking": true,
"source": "quota",
"ttl_seconds": serde_json::Value::Null,
"detail": serde_json::Value::Null,
})
);
assert_eq!(keys[0]["account_quota"], json!("5H剩余 0.0%"));
}
#[tokio::test]
async fn gateway_marks_exhausted_kiro_pool_key_as_blocked_when_flag_enabled() {
let mut provider = sample_provider("provider-kiro", "kiro", 10).with_transport_fields(
true,
false,
true,
None,
None,
None,
None,
None,
Some(json!({
"pool_advanced": {
"enabled": true,
"skip_exhausted_accounts": true
}
})),
);
provider.provider_type = "kiro".to_string();
let mut key = sample_key(
"key-kiro-exhausted",
"provider-kiro",
"claude:cli",
"oauth-placeholder",
);
key.name = "kiro exhausted".to_string();
key.auth_type = "oauth".to_string();
key.upstream_metadata = Some(json!({
"kiro": {
"remaining": 0.0,
"usage_limit": 100.0,
"current_usage": 100.0
}
}));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
Vec::new(),
vec![key],
));
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
));
let response = local_admin_pool_response(
&state,
http::Method::GET,
"/api/admin/pool/provider-kiro/keys?page=1&page_size=50&status=all",
None,
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = serde_json::from_slice(
&to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should read"),
)
.expect("json body should parse");
let keys = payload["keys"].as_array().expect("keys should be array");
assert_eq!(keys[0]["scheduling_status"], json!("blocked"));
assert_eq!(
keys[0]["scheduling_reason"],
json!("account_quota_exhausted")
);
assert_eq!(keys[0]["scheduling_label"], json!("额度耗尽"));
assert_eq!(keys[0]["account_quota"], json!("剩余 0/100"));
}
#[tokio::test]
async fn gateway_codex_quota_resets_to_full_after_countdown_elapsed() {
let mut provider = sample_provider("provider-codex", "codex", 10).with_transport_fields(
@@ -1427,6 +1583,65 @@ async fn gateway_handles_admin_pool_resolve_selection_locally_with_trusted_admin
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_resolve_selection_marks_legacy_kiro_bearer_keys_as_oauth_managed() {
let mut provider = sample_provider("provider-kiro", "kiro", 10);
provider.provider_type = "kiro".to_string();
let mut key = sample_key(
"key-kiro-legacy",
"provider-kiro",
"kiro:generateAssistantResponse",
"kiro-access-token",
);
key.auth_type = "bearer".to_string();
key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"kiro","email":"legacy-kiro@example.com","refresh_token":"legacy-kiro-refresh-token"}"#,
)
.expect("auth config ciphertext should build"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
Vec::new(),
vec![key],
));
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_reader_for_tests(provider_catalog_repository)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
);
let response = local_admin_pool_response(
&state,
http::Method::POST,
"/api/admin/pool/provider-kiro/keys/resolve-selection",
Some(json!({})),
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = serde_json::from_slice(
&to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should read"),
)
.expect("json body should parse");
assert_eq!(payload["total"], json!(1));
let items = payload["items"].as_array().expect("items should be array");
assert_eq!(items.len(), 1);
assert_eq!(items[0]["key_id"], json!("key-kiro-legacy"));
assert_eq!(items[0]["auth_type"], json!("bearer"));
assert_eq!(items[0]["credential_kind"], json!("oauth_session"));
assert_eq!(items[0]["runtime_auth_kind"], json!("bearer"));
assert_eq!(items[0]["oauth_managed"], json!(true));
assert_eq!(items[0]["can_refresh_oauth"], json!(true));
assert_eq!(items[0]["can_export_oauth"], json!(true));
assert_eq!(items[0]["can_edit_oauth"], json!(true));
}
#[tokio::test]
async fn gateway_handles_admin_pool_batch_action_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));

322
apps/aether-gateway/src/tests/control/admin/provider_query.rs
View File

@@ -7,6 +7,7 @@ use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogEn
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;
@@ -20,6 +21,56 @@ use crate::constants::{
};
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,
@@ -676,6 +727,277 @@ async fn gateway_handles_admin_provider_query_test_model_failover_locally_with_t
.await;
}
#[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:cli");
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:cli",
"__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",
"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:cli".to_string(),
Some("claude".to_string()),
Some("cli".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: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["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_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| {
let mut key = sample_key(id, "provider-kiro", "claude:cli", "__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",
"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(),
"claude:cli".to_string(),
Some("claude".to_string()),
Some("cli".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"],
"api_format": "claude:cli",
"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));
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_handles_admin_provider_query_test_model_failover_with_single_model_name_alias() {
assert_admin_provider_query_route(