i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-03 01:40:21 +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

13
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/complete/key.rs
View File

@@ -10,6 +10,7 @@ use super::shared::{
};
use crate::handlers::admin::provider::shared::paths::admin_provider_oauth_complete_key_id;
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::provider_key_auth::provider_key_is_oauth_managed;
use crate::GatewayError;
use axum::{
body::{Body, Bytes},
@@ -76,12 +77,6 @@ pub(super) async fn handle_admin_provider_oauth_complete_key(
"Key 不存在",
));
};
if !key.auth_type.eq_ignore_ascii_case("oauth") {
return Ok(build_internal_control_error_response(
http::StatusCode::BAD_REQUEST,
"该 Key 不是 oauth 认证类型",
));
}
if !state_data.provider_id.trim().is_empty() && state_data.provider_id != key.provider_id {
return Ok(build_internal_control_error_response(
http::StatusCode::BAD_REQUEST,
@@ -102,6 +97,12 @@ pub(super) async fn handle_admin_provider_oauth_complete_key(
));
};
let provider_type = provider.provider_type.trim().to_ascii_lowercase();
if !provider_key_is_oauth_managed(&key, provider_type.as_str()) {
return Ok(build_internal_control_error_response(
http::StatusCode::BAD_REQUEST,
"该 Key 不是 OAuth 管理账号",
));
}
if !is_fixed_provider_type_for_provider_oauth(&provider_type) {
return Ok(build_internal_control_error_response(
http::StatusCode::BAD_REQUEST,

6
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/import.rs
View File

@@ -94,6 +94,12 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
"该 Provider 不是固定类型，无法使用 provider-oauth",
));
}
if provider_type == "kiro" {
return Ok(build_internal_control_error_response(
http::StatusCode::BAD_REQUEST,
"Kiro 不支持单条 Refresh Token 导入，请使用批量导入或设备授权。",
));
}
let Some(template) = admin_provider_oauth_template(&provider_type) else {
return Ok(build_admin_provider_oauth_backend_unavailable_response());
};

14
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/refresh/request.rs
View File

@@ -4,6 +4,7 @@ use super::helpers::{self, RefreshDispatch, RefreshRequestContext};
use super::response;
use crate::handlers::admin::provider::shared::paths::admin_provider_oauth_refresh_key_id;
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::provider_key_auth::provider_key_is_oauth_managed;
use crate::GatewayError;
use axum::http;
@@ -28,13 +29,6 @@ pub(super) async fn parse_admin_provider_oauth_refresh_request(
"Key 不存在",
)));
};
if !key.auth_type.eq_ignore_ascii_case("oauth") {
return Ok(RefreshDispatch::Respond(response::control_error_response(
http::StatusCode::BAD_REQUEST,
"该 Key 不是 oauth 认证类型",
)));
}
let Some(encrypted_auth_config) = key.encrypted_auth_config.as_deref() else {
return Ok(RefreshDispatch::Respond(response::control_error_response(
http::StatusCode::BAD_REQUEST,
@@ -69,6 +63,12 @@ pub(super) async fn parse_admin_provider_oauth_refresh_request(
)));
};
let provider_type = provider.provider_type.trim().to_ascii_lowercase();
if !provider_key_is_oauth_managed(&key, provider_type.as_str()) {
return Ok(RefreshDispatch::Respond(response::control_error_response(
http::StatusCode::BAD_REQUEST,
"该 Key 不是 OAuth 管理账号",
)));
}
if !is_fixed_provider_type_for_provider_oauth(&provider_type) {
return Ok(RefreshDispatch::Respond(response::control_error_response(
http::StatusCode::BAD_REQUEST,

14
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/start.rs
View File

@@ -8,6 +8,7 @@ use crate::handlers::admin::provider::shared::paths::{
admin_provider_oauth_start_key_id, admin_provider_oauth_start_provider_id,
};
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::provider_key_auth::provider_key_is_oauth_managed;
use crate::GatewayError;
use axum::{
body::Body,
@@ -37,13 +38,6 @@ pub(super) async fn handle_admin_provider_oauth_start_key(
"Key 不存在",
));
};
if !key.auth_type.eq_ignore_ascii_case("oauth") {
return Ok(build_internal_control_error_response(
http::StatusCode::BAD_REQUEST,
"该 Key 不是 oauth 认证类型",
));
}
let provider_id = key.provider_id.clone();
let provider = state
.read_provider_catalog_providers_by_ids(std::slice::from_ref(&provider_id))
@@ -57,6 +51,12 @@ pub(super) async fn handle_admin_provider_oauth_start_key(
));
};
let provider_type = provider.provider_type.trim().to_ascii_lowercase();
if !provider_key_is_oauth_managed(&key, provider_type.as_str()) {
return Ok(build_internal_control_error_response(
http::StatusCode::BAD_REQUEST,
"该 Key 不是 OAuth 管理账号",
));
}
if !is_fixed_provider_type_for_provider_oauth(&provider_type) {
return Ok(build_internal_control_error_response(
http::StatusCode::BAD_REQUEST,

8
apps/aether-gateway/src/handlers/admin/provider/oauth/duplicates.rs
View File

@@ -1,4 +1,5 @@
use crate::handlers::admin::request::AdminAppState;
use crate::provider_key_auth::provider_key_is_oauth_managed;
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
fn normalize_codex_plan_group_for_provider_oauth(
@@ -141,8 +142,13 @@ pub(crate) async fn find_duplicate_provider_oauth_key(
.await
.map_err(|err| format!("{err:?}"))?;
let provider_type = auth_config
.get("provider_type")
.and_then(serde_json::Value::as_str)
.unwrap_or_default()
.to_string();
for existing_key in existing_keys.into_iter().filter(|key| {
key.auth_type.trim().eq_ignore_ascii_case("oauth")
provider_key_is_oauth_managed(key, provider_type.as_str())
&& exclude_key_id.is_none_or(|exclude| key.id != exclude)
}) {
let Some(existing_auth_config) = state.parse_catalog_auth_config_json(&existing_key) else {

12
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/codex/mod.rs
View File

@@ -18,6 +18,7 @@ use super::shared::{
should_auto_remove_structured_reason, ProviderQuotaExecutionOutcome,
};
use crate::handlers::admin::request::AdminAppState;
use crate::provider_key_auth::provider_key_is_oauth_managed;
use crate::GatewayError;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
@@ -55,11 +56,12 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
}
};
let resolved_oauth_auth = if key.auth_type.trim().eq_ignore_ascii_case("oauth") {
state.resolve_local_oauth_header_auth(&transport).await?
} else {
None
};
let resolved_oauth_auth =
if provider_key_is_oauth_managed(&key, provider.provider_type.as_str()) {
state.resolve_local_oauth_header_auth(&transport).await?
} else {
None
};
let headers = match build_codex_refresh_headers(&transport, resolved_oauth_auth) {
Ok(headers) => headers,

4
apps/aether-gateway/src/handlers/admin/provider/oauth/runtime.rs
View File

@@ -2,6 +2,7 @@ use super::quota::antigravity::refresh_antigravity_provider_quota_locally;
use super::quota::codex::refresh_codex_provider_quota_locally;
use super::quota::kiro::refresh_kiro_provider_quota_locally;
use crate::handlers::admin::request::AdminAppState;
use crate::provider_key_auth::provider_key_is_oauth_managed;
use crate::GatewayError;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogProvider,
@@ -85,6 +86,9 @@ pub(crate) async fn refresh_provider_oauth_account_state_after_update(
else {
return Ok((false, None));
};
if provider_type == "kiro" && !provider_key_is_oauth_managed(&key, provider_type.as_str()) {
return Ok((false, None));
}
let payload = match provider_type.as_str() {
"codex" => {

59
apps/aether-gateway/src/handlers/admin/provider/pool/config.rs
View File

@@ -53,6 +53,7 @@ pub(crate) fn admin_provider_pool_config(
let Some(pool_advanced) = raw_pool_advanced.as_object() else {
return Some(AdminProviderPoolConfig {
lru_enabled: false,
skip_exhausted_accounts: false,
cost_window_seconds: 18_000,
cost_limit_per_key_tokens: None,
});
@@ -60,6 +61,10 @@ pub(crate) fn admin_provider_pool_config(
Some(AdminProviderPoolConfig {
lru_enabled: admin_provider_pool_lru_enabled(pool_advanced),
skip_exhausted_accounts: pool_advanced
.get("skip_exhausted_accounts")
.and_then(Value::as_bool)
.unwrap_or(false),
cost_window_seconds: pool_advanced
.get("cost_window_seconds")
.and_then(json_u64)
@@ -70,3 +75,57 @@ pub(crate) fn admin_provider_pool_config(
.and_then(json_u64),
})
}
#[cfg(test)]
mod tests {
use super::admin_provider_pool_config;
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogProvider;
use serde_json::json;
fn sample_provider(config: serde_json::Value) -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
"provider-1".to_string(),
"provider-1".to_string(),
Some("https://example.com".to_string()),
"codex".to_string(),
)
.expect("provider should build")
.with_transport_fields(
true,
false,
false,
None,
None,
None,
None,
None,
Some(config),
)
}
#[test]
fn defaults_skip_exhausted_accounts_to_false() {
let provider = sample_provider(json!({ "pool_advanced": {} }));
let config = admin_provider_pool_config(&provider).expect("pool config should exist");
assert!(!config.skip_exhausted_accounts);
}
#[test]
fn parses_skip_exhausted_accounts_from_pool_advanced() {
let provider = sample_provider(json!({
"pool_advanced": {
"skip_exhausted_accounts": true,
"lru_enabled": true,
"cost_window_seconds": 7200,
"cost_limit_per_key_tokens": 12000
}
}));
let config = admin_provider_pool_config(&provider).expect("pool config should exist");
assert!(config.skip_exhausted_accounts);
assert!(config.lru_enabled);
assert_eq!(config.cost_window_seconds, 7200);
assert_eq!(config.cost_limit_per_key_tokens, Some(12_000));
}
}

89
apps/aether-gateway/src/handlers/admin/provider/pool_admin/payloads.rs
View File

@@ -3,6 +3,8 @@ use crate::handlers::admin::provider::shared::support::{
};
use crate::handlers::admin::request::AdminAppState;
use crate::handlers::admin::shared::{provider_key_status_snapshot_payload, unix_secs_to_rfc3339};
use crate::provider_key_auth::provider_key_auth_semantics;
use aether_admin::provider::pool as admin_provider_pool_pure;
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use serde_json::json;
@@ -122,10 +124,11 @@ fn admin_pool_normalize_oauth_plan_type(value: &str, provider_type: &str) -> Opt
}
fn admin_pool_derive_oauth_expires_at(
provider_type: &str,
key: &StoredProviderCatalogKey,
auth_config: Option<&serde_json::Map<String, serde_json::Value>>,
) -> Option<u64> {
if !key.auth_type.trim().eq_ignore_ascii_case("oauth") {
if !provider_key_auth_semantics(key, provider_type).oauth_managed() {
return None;
}
@@ -144,7 +147,7 @@ fn admin_pool_derive_oauth_plan_type(
provider_type: &str,
auth_config: Option<&serde_json::Map<String, serde_json::Value>>,
) -> Option<String> {
if !key.auth_type.trim().eq_ignore_ascii_case("oauth") {
if !provider_key_auth_semantics(key, provider_type).oauth_managed() {
return None;
}
@@ -551,6 +554,7 @@ fn admin_pool_scheduling_payload(
cooldown_ttl_seconds: Option<u64>,
health_score: f64,
circuit_breaker_open: bool,
account_quota_exhausted: bool,
) -> (String, String, String, Vec<serde_json::Value>) {
if !key.is_active {
return (
@@ -567,6 +571,21 @@ fn admin_pool_scheduling_payload(
})],
);
}
if account_quota_exhausted {
return (
"blocked".to_string(),
"account_quota_exhausted".to_string(),
"额度耗尽".to_string(),
vec![json!({
"code": "account_quota_exhausted",
"label": "额度耗尽",
"blocking": true,
"source": "quota",
"ttl_seconds": serde_json::Value::Null,
"detail": serde_json::Value::Null,
})],
);
}
if let Some(reason) = cooldown_reason {
return (
"degraded".to_string(),
@@ -632,6 +651,9 @@ pub(super) fn build_admin_pool_key_payload(
.and_then(|_| runtime.cooldown_ttl_by_key.get(&key.id).copied());
let health_score = admin_pool_health_score(key);
let circuit_breaker_open = admin_pool_circuit_breaker_open(key);
let auth_semantics = provider_key_auth_semantics(key, provider_type);
let account_quota_exhausted = pool_config.is_some_and(|config| config.skip_exhausted_accounts)
&& admin_provider_pool_pure::admin_pool_key_account_quota_exhausted(key, provider_type);
let (scheduling_status, scheduling_reason, scheduling_label, scheduling_reasons) =
admin_pool_scheduling_payload(
key,
@@ -639,9 +661,11 @@ pub(super) fn build_admin_pool_key_payload(
cooldown_ttl_seconds,
health_score,
circuit_breaker_open,
account_quota_exhausted,
);
let auth_config = state.parse_catalog_auth_config_json(key);
let oauth_expires_at = admin_pool_derive_oauth_expires_at(key, auth_config.as_ref());
let oauth_expires_at =
admin_pool_derive_oauth_expires_at(provider_type, key, auth_config.as_ref());
let oauth_plan_type =
admin_pool_derive_oauth_plan_type(key, provider_type, auth_config.as_ref());
let status_snapshot = provider_key_status_snapshot_payload(key);
@@ -656,15 +680,29 @@ pub(super) fn build_admin_pool_key_payload(
.and_then(serde_json::Value::as_object);
let quota_updated_at =
admin_pool_json_to_u64(quota_snapshot.and_then(|item| item.get("updated_at")));
let oauth_invalid_at =
let oauth_invalid_at = if auth_semantics.can_show_oauth_metadata() {
admin_pool_json_to_u64(oauth_snapshot.and_then(|item| item.get("invalid_at")))
.or(key.oauth_invalid_at_unix_secs);
let oauth_account_id = admin_pool_trimmed_string_from_map(auth_config.as_ref(), "account_id");
let oauth_account_name =
admin_pool_trimmed_string_from_map(auth_config.as_ref(), "account_name");
let oauth_account_user_id =
admin_pool_trimmed_string_from_map(auth_config.as_ref(), "account_user_id");
let oauth_organizations = admin_pool_oauth_organizations(auth_config.as_ref());
.or(key.oauth_invalid_at_unix_secs)
} else {
None
};
let oauth_account_id = auth_semantics
.can_show_oauth_metadata()
.then(|| admin_pool_trimmed_string_from_map(auth_config.as_ref(), "account_id"))
.flatten();
let oauth_account_name = auth_semantics
.can_show_oauth_metadata()
.then(|| admin_pool_trimmed_string_from_map(auth_config.as_ref(), "account_name"))
.flatten();
let oauth_account_user_id = auth_semantics
.can_show_oauth_metadata()
.then(|| admin_pool_trimmed_string_from_map(auth_config.as_ref(), "account_user_id"))
.flatten();
let oauth_organizations = if auth_semantics.can_show_oauth_metadata() {
admin_pool_oauth_organizations(auth_config.as_ref())
} else {
Vec::new()
};
let account_status_code = admin_pool_trimmed_string_from_map(account_snapshot, "code");
let account_status_label =
admin_pool_trimmed_string(account_snapshot.and_then(|item| item.get("label")));
@@ -686,11 +724,38 @@ pub(super) fn build_admin_pool_key_payload(
payload.insert("key_name".to_string(), json!(key.name));
payload.insert("is_active".to_string(), json!(key.is_active));
payload.insert("auth_type".to_string(), json!(key.auth_type));
payload.insert(
"credential_kind".to_string(),
json!(auth_semantics.credential_kind().as_str()),
);
payload.insert(
"runtime_auth_kind".to_string(),
json!(auth_semantics.runtime_auth_kind().as_str()),
);
payload.insert(
"oauth_managed".to_string(),
json!(auth_semantics.oauth_managed()),
);
payload.insert(
"can_refresh_oauth".to_string(),
json!(auth_semantics.can_refresh_oauth()),
);
payload.insert(
"can_export_oauth".to_string(),
json!(auth_semantics.can_export_oauth()),
);
payload.insert(
"can_edit_oauth".to_string(),
json!(auth_semantics.can_edit_oauth()),
);
payload.insert("oauth_expires_at".to_string(), json!(oauth_expires_at));
payload.insert("oauth_invalid_at".to_string(), json!(oauth_invalid_at));
payload.insert(
"oauth_invalid_reason".to_string(),
json!(key.oauth_invalid_reason),
json!(auth_semantics
.can_show_oauth_metadata()
.then_some(key.oauth_invalid_reason.clone())
.flatten()),
);
payload.insert("oauth_plan_type".to_string(), json!(oauth_plan_type));
payload.insert("oauth_account_id".to_string(), json!(oauth_account_id));

26
apps/aether-gateway/src/handlers/admin/provider/pool_admin/read_routes/resolve_selection.rs
View File

@@ -3,6 +3,7 @@ use super::{
AdminPoolResolveSelectionRequest, ADMIN_POOL_PROVIDER_CATALOG_READER_UNAVAILABLE_DETAIL,
};
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::provider_key_auth::provider_key_auth_semantics;
use crate::GatewayError;
use aether_admin::provider::pool as admin_provider_pool_pure;
use axum::{
@@ -11,6 +12,7 @@ use axum::{
response::{IntoResponse, Response},
Json,
};
use serde_json::json;
pub(super) async fn build_admin_pool_resolve_selection_response(
state: &AdminAppState<'_>,
@@ -88,5 +90,27 @@ pub(super) async fn build_admin_pool_resolve_selection_response(
.then_with(|| left.name.cmp(&right.name))
});
Ok(Json(admin_provider_pool_pure::build_admin_pool_selection_payload(&keys)).into_response())
let items = keys
.iter()
.map(|key| {
let auth_semantics = provider_key_auth_semantics(key, &provider_type);
json!({
"key_id": key.id,
"key_name": key.name,
"auth_type": key.auth_type,
"credential_kind": auth_semantics.credential_kind().as_str(),
"runtime_auth_kind": auth_semantics.runtime_auth_kind().as_str(),
"oauth_managed": auth_semantics.oauth_managed(),
"can_refresh_oauth": auth_semantics.can_refresh_oauth(),
"can_export_oauth": auth_semantics.can_export_oauth(),
"can_edit_oauth": auth_semantics.can_edit_oauth(),
})
})
.collect::<Vec<_>>();
Ok(Json(json!({
"total": items.len(),
"items": items,
}))
.into_response())
}

3
apps/aether-gateway/src/handlers/admin/provider/pool_admin/selection.rs
View File

@@ -1,4 +1,5 @@
use crate::handlers::admin::request::AdminAppState;
use crate::provider_key_auth::provider_key_is_oauth_managed;
use aether_admin::provider::pool as admin_provider_pool_pure;
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
@@ -44,7 +45,7 @@ fn admin_pool_derive_oauth_plan_type(
}
};
if key.auth_type.trim() != "oauth" {
if !provider_key_is_oauth_managed(key, provider_type) {
return None;
}

818
apps/aether-gateway/src/handlers/admin/provider/query/models.rs
View File

@@ -1,17 +1,29 @@
use super::payload::{
provider_query_extract_api_key_id, provider_query_extract_force_refresh,
provider_query_extract_provider_id,
provider_query_extract_model, provider_query_extract_provider_id,
provider_query_extract_request_id,
};
use super::response::{
build_admin_provider_query_bad_request_response, build_admin_provider_query_not_found_response,
ADMIN_PROVIDER_QUERY_API_KEY_NOT_FOUND_DETAIL, ADMIN_PROVIDER_QUERY_NO_ACTIVE_API_KEY_DETAIL,
ADMIN_PROVIDER_QUERY_API_KEY_NOT_FOUND_DETAIL, ADMIN_PROVIDER_QUERY_MODEL_REQUIRED_DETAIL,
ADMIN_PROVIDER_QUERY_NO_ACTIVE_API_KEY_DETAIL, ADMIN_PROVIDER_QUERY_NO_LOCAL_MODELS_DETAIL,
ADMIN_PROVIDER_QUERY_PROVIDER_ID_REQUIRED_DETAIL,
ADMIN_PROVIDER_QUERY_PROVIDER_NOT_FOUND_DETAIL,
};
use crate::ai_pipeline::{maybe_build_sync_finalize_outcome, GatewayControlDecision};
use crate::execution_runtime;
use crate::handlers::admin::request::AdminAppState;
use crate::model_fetch::ModelFetchRuntimeState;
use crate::provider_transport::kiro::{
build_kiro_generate_assistant_response_url, build_kiro_provider_headers,
build_kiro_provider_request_body, supports_local_kiro_request_transport_with_network,
KiroProviderHeadersInput, KIRO_ENVELOPE_NAME,
};
use crate::usage::GatewaySyncReportRequest;
use crate::{AppState, GatewayError};
use aether_admin::provider::pool as admin_provider_pool_pure;
use aether_contracts::{ExecutionPlan, RequestBody};
use aether_data_contracts::repository::global_models::AdminProviderModelListQuery;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
@@ -19,9 +31,16 @@ use aether_model_fetch::{
aggregate_models_for_cache, fetch_models_from_transports, json_string_list,
preset_models_for_provider, selected_models_fetch_endpoints,
};
use axum::{body::Body, http::Response, response::IntoResponse, Json};
use axum::{
body::{to_bytes, Body},
http::{HeaderMap, HeaderName, HeaderValue},
response::{IntoResponse, Response},
Json,
};
use base64::Engine as _;
use serde_json::{json, Value};
use std::collections::BTreeSet;
use std::collections::{BTreeMap, BTreeSet};
use uuid::Uuid;
pub(crate) const ADMIN_PROVIDER_QUERY_LOCAL_TEST_MODEL_MESSAGE: &str =
"Rust local provider-query model test is not configured";
@@ -32,7 +51,10 @@ const ADMIN_PROVIDER_QUERY_NO_ACTIVE_ENDPOINT_DETAIL: &str =
const ADMIN_PROVIDER_QUERY_NO_MODELS_FROM_ENDPOINT_DETAIL: &str =
"No models returned from any endpoint";
const ADMIN_PROVIDER_QUERY_NO_MODELS_FROM_KEY_DETAIL: &str = "No models returned from any key";
const ADMIN_PROVIDER_QUERY_NO_ACTIVE_TEST_CANDIDATE_DETAIL: &str =
"No active endpoint or API key found";
const ANTIGRAVITY_PROVIDER_CACHE_KEY_PREFIX: &str = "upstream_models_provider:";
const DEFAULT_PROVIDER_QUERY_TEST_MESSAGE: &str = "Hello! This is a test message.";
#[derive(Debug)]
struct ProviderQueryKeyFetchResult {
@@ -42,6 +64,47 @@ struct ProviderQueryKeyFetchResult {
has_success: bool,
}
#[derive(Debug, Clone)]
struct ProviderQueryTestCandidate {
endpoint: StoredProviderCatalogEndpoint,
key: StoredProviderCatalogKey,
effective_model: String,
}
#[derive(Debug, Clone)]
struct ProviderQueryTestAttempt {
candidate_index: usize,
endpoint_api_format: String,
endpoint_base_url: String,
key_name: String,
key_id: String,
auth_type: String,
effective_model: String,
status: &'static str,
skip_reason: Option<String>,
error_message: Option<String>,
status_code: Option<u16>,
latency_ms: Option<u64>,
request_url: Option<String>,
request_headers: Option<BTreeMap<String, String>>,
request_body: Option<Value>,
response_headers: Option<BTreeMap<String, String>>,
response_body: Option<Value>,
}
#[derive(Debug, Clone)]
struct ProviderQueryExecutionOutcome {
status: &'static str,
error_message: Option<String>,
status_code: Option<u16>,
latency_ms: Option<u64>,
request_url: String,
request_headers: BTreeMap<String, String>,
request_body: Value,
response_headers: BTreeMap<String, String>,
response_body: Option<Value>,
}
fn provider_query_provider_payload(provider: &StoredProviderCatalogProvider) -> Value {
json!({
"id": provider.id.clone(),
@@ -50,6 +113,753 @@ fn provider_query_provider_payload(provider: &StoredProviderCatalogProvider) ->
})
}
fn provider_query_test_mode(payload: &Value) -> &str {
payload
.get("mode")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("global")
}
fn provider_query_extract_endpoint_id(payload: &Value) -> Option<String> {
payload
.get("endpoint_id")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
fn provider_query_extract_api_format(payload: &Value) -> Option<String> {
payload
.get("api_format")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
fn provider_query_extract_message(payload: &Value) -> Option<String> {
payload
.get("message")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
fn provider_query_extract_request_body(payload: &Value) -> Option<Value> {
payload
.get("request_body")
.filter(|value| value.is_object())
.cloned()
}
fn provider_query_extract_request_headers(payload: &Value) -> HeaderMap {
let mut headers = HeaderMap::new();
let Some(values) = payload.get("request_headers").and_then(Value::as_object) else {
return headers;
};
for (key, value) in values {
let key = key.trim();
if key.is_empty() {
continue;
}
let Some(value) = (match value {
Value::String(value) => Some(value.trim().to_string()),
Value::Bool(value) => Some(value.to_string()),
Value::Number(value) => Some(value.to_string()),
other => serde_json::to_string(other).ok(),
}) else {
continue;
};
if value.is_empty() {
continue;
}
let Ok(name) = HeaderName::from_bytes(key.as_bytes()) else {
continue;
};
let Ok(value) = HeaderValue::from_str(&value) else {
continue;
};
headers.insert(name, value);
}
headers
}
fn provider_query_build_test_request_body(payload: &Value, model: &str) -> Value {
if let Some(mut body) = provider_query_extract_request_body(payload) {
if let Some(object) = body.as_object_mut() {
object.insert("model".to_string(), Value::String(model.to_string()));
}
return body;
}
json!({
"model": model,
"messages": [{
"role": "user",
"content": provider_query_extract_message(payload)
.unwrap_or_else(|| DEFAULT_PROVIDER_QUERY_TEST_MESSAGE.to_string())
}],
"max_tokens": 30,
"temperature": 0.7,
"stream": true,
})
}
fn provider_query_select_kiro_endpoint<'a>(
endpoints: &'a [StoredProviderCatalogEndpoint],
endpoint_id: Option<&str>,
api_format: Option<&str>,
) -> Result<Option<&'a StoredProviderCatalogEndpoint>, &'static str> {
if let Some(endpoint_id) = endpoint_id {
let endpoint = endpoints.iter().find(|endpoint| endpoint.id == endpoint_id);
return endpoint
.ok_or("Endpoint not found")
.map(|endpoint| Some(endpoint));
}
if let Some(api_format) = api_format {
let endpoint = endpoints.iter().find(|endpoint| {
endpoint.is_active && endpoint.api_format.trim().eq_ignore_ascii_case(api_format)
});
return Ok(endpoint);
}
Ok(endpoints.iter().find(|endpoint| endpoint.is_active))
}
fn provider_query_key_supports_endpoint(
key: &StoredProviderCatalogKey,
endpoint_api_format: &str,
) -> bool {
let formats = json_string_list(key.api_formats.as_ref());
formats.is_empty()
|| formats
.iter()
.any(|value| value.eq_ignore_ascii_case(endpoint_api_format))
}
fn provider_query_test_key_sort_key(
provider_type: &str,
key: &StoredProviderCatalogKey,
endpoint_api_format: &str,
) -> (u8, u8, i32, u64, i32) {
let quota_exhausted =
admin_provider_pool_pure::admin_pool_key_account_quota_exhausted(key, provider_type);
let circuit_open = key
.circuit_breaker_by_format
.as_ref()
.and_then(Value::as_object)
.and_then(|value| value.get(endpoint_api_format))
.and_then(Value::as_object)
.and_then(|value| value.get("open"))
.and_then(Value::as_bool)
.unwrap_or(false);
let health_score = key
.health_by_format
.as_ref()
.and_then(Value::as_object)
.and_then(|value| value.get(endpoint_api_format))
.and_then(Value::as_object)
.and_then(|value| value.get("health_score"))
.and_then(Value::as_f64)
.unwrap_or(1.0);
let consecutive_failures = key
.health_by_format
.as_ref()
.and_then(Value::as_object)
.and_then(|value| value.get(endpoint_api_format))
.and_then(Value::as_object)
.and_then(|value| value.get("consecutive_failures"))
.and_then(Value::as_u64)
.unwrap_or(0);
let normalized_health = (health_score.clamp(0.0, 1.0) * 1000.0).round() as i32;
(
if quota_exhausted { 1 } else { 0 },
if circuit_open { 1 } else { 0 },
-normalized_health,
consecutive_failures,
key.internal_priority,
)
}
async fn provider_query_resolve_global_effective_model(
state: &AdminAppState<'_>,
provider_id: &str,
requested_model: &str,
) -> Result<String, GatewayError> {
let models = state
.list_admin_provider_models(&AdminProviderModelListQuery {
provider_id: provider_id.to_string(),
is_active: Some(true),
offset: 0,
limit: 1024,
})
.await
.unwrap_or_default();
Ok(models
.into_iter()
.find(|model| {
model.is_available
&& model
.global_model_name
.as_deref()
.is_some_and(|value| value.eq_ignore_ascii_case(requested_model))
})
.map(|model| model.provider_model_name)
.unwrap_or_else(|| requested_model.to_string()))
}
async fn provider_query_build_kiro_test_candidates(
state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider,
payload: &Value,
) -> Result<Vec<ProviderQueryTestCandidate>, Response<Body>> {
let provider_ids = vec![provider.id.clone()];
let endpoints = state
.app()
.list_provider_catalog_endpoints_by_provider_ids(&provider_ids)
.await
.map_err(|_| {
build_admin_provider_query_bad_request_response(
ADMIN_PROVIDER_QUERY_NO_ACTIVE_API_KEY_DETAIL,
)
})?;
let endpoint = match provider_query_select_kiro_endpoint(
&endpoints,
provider_query_extract_endpoint_id(payload).as_deref(),
provider_query_extract_api_format(payload).as_deref(),
) {
Ok(Some(endpoint)) => endpoint.clone(),
Ok(None) => {
return Err(build_admin_provider_query_not_found_response(
ADMIN_PROVIDER_QUERY_NO_ACTIVE_API_KEY_DETAIL,
));
}
Err("Endpoint not found") => {
return Err(build_admin_provider_query_not_found_response(
ADMIN_PROVIDER_QUERY_API_KEY_NOT_FOUND_DETAIL,
));
}
Err(_) => {
return Err(build_admin_provider_query_not_found_response(
ADMIN_PROVIDER_QUERY_NO_ACTIVE_API_KEY_DETAIL,
));
}
};
let all_keys = state
.app()
.list_provider_catalog_keys_by_provider_ids(&provider_ids)
.await
.map_err(|_| {
build_admin_provider_query_bad_request_response(
ADMIN_PROVIDER_QUERY_NO_ACTIVE_API_KEY_DETAIL,
)
})?;
let selected_key_id = provider_query_extract_api_key_id(payload);
if let Some(api_key_id) = selected_key_id.as_deref() {
let Some(key) = all_keys.iter().find(|key| key.id == api_key_id) else {
return Err(build_admin_provider_query_not_found_response(
ADMIN_PROVIDER_QUERY_API_KEY_NOT_FOUND_DETAIL,
));
};
if !key.is_active || !provider_query_key_supports_endpoint(key, &endpoint.api_format) {
return Err(build_admin_provider_query_not_found_response(
ADMIN_PROVIDER_QUERY_NO_ACTIVE_TEST_CANDIDATE_DETAIL,
));
}
}
let requested_model = provider_query_extract_model(payload).ok_or_else(|| {
build_admin_provider_query_bad_request_response(ADMIN_PROVIDER_QUERY_MODEL_REQUIRED_DETAIL)
})?;
let effective_model = if provider_query_test_mode(payload).eq_ignore_ascii_case("direct") {
requested_model.clone()
} else {
provider_query_resolve_global_effective_model(state, &provider.id, &requested_model)
.await
.unwrap_or(requested_model.clone())
};
let mut keys = all_keys
.into_iter()
.filter(|key| key.is_active)
.filter(|key| {
selected_key_id
.as_deref()
.is_none_or(|value| value == key.id.as_str())
})
.filter(|key| provider_query_key_supports_endpoint(key, &endpoint.api_format))
.collect::<Vec<_>>();
keys.sort_by_key(|key| {
provider_query_test_key_sort_key(provider.provider_type.as_str(), key, &endpoint.api_format)
});
let candidates = keys
.into_iter()
.map(|key| ProviderQueryTestCandidate {
endpoint: endpoint.clone(),
key,
effective_model: effective_model.clone(),
})
.collect::<Vec<_>>();
if candidates.is_empty() {
return Err(build_admin_provider_query_not_found_response(
ADMIN_PROVIDER_QUERY_NO_ACTIVE_TEST_CANDIDATE_DETAIL,
));
}
Ok(candidates)
}
fn provider_query_decode_execution_body(
result: &aether_contracts::ExecutionResult,
) -> Option<Vec<u8>> {
result
.body
.as_ref()
.and_then(|body| body.body_bytes_b64.as_deref())
.and_then(|value| base64::engine::general_purpose::STANDARD.decode(value).ok())
}
fn provider_query_extract_error_message(
result: &aether_contracts::ExecutionResult,
) -> Option<String> {
result
.body
.as_ref()
.and_then(|body| body.json_body.as_ref())
.and_then(Value::as_object)
.and_then(|value| {
value
.get("error")
.and_then(Value::as_object)
.and_then(|error| error.get("message"))
.and_then(Value::as_str)
.or_else(|| value.get("message").and_then(Value::as_str))
})
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.or_else(|| {
provider_query_decode_execution_body(result)
.and_then(|bytes| String::from_utf8(bytes).ok())
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
})
.or_else(|| {
result
.error
.as_ref()
.map(|error| error.message.trim().to_string())
.filter(|value| !value.is_empty())
})
}
async fn provider_query_finalize_kiro_result(
route_path: &str,
trace_id: &str,
requested_model: &str,
endpoint_api_format: &str,
effective_model: &str,
original_request_body: &Value,
result: &aether_contracts::ExecutionResult,
) -> Result<Option<Value>, GatewayError> {
let decision = GatewayControlDecision::synthetic(
route_path,
Some("admin_proxy".to_string()),
Some("provider_query_manage".to_string()),
Some("test_model_failover".to_string()),
Some(endpoint_api_format.to_string()),
);
let payload = GatewaySyncReportRequest {
trace_id: trace_id.to_string(),
report_kind: "claude_cli_sync_finalize".to_string(),
report_context: Some(json!({
"client_api_format": endpoint_api_format,
"provider_api_format": endpoint_api_format,
"model": requested_model,
"mapped_model": effective_model,
"needs_conversion": false,
"has_envelope": true,
"envelope_name": KIRO_ENVELOPE_NAME,
"original_request_body": original_request_body,
})),
status_code: result.status_code,
headers: result.headers.clone(),
body_json: result.body.as_ref().and_then(|body| body.json_body.clone()),
client_body_json: None,
body_base64: result
.body
.as_ref()
.and_then(|body| body.body_bytes_b64.clone()),
telemetry: result.telemetry.clone(),
};
let Some(outcome) = maybe_build_sync_finalize_outcome(trace_id, &decision, &payload)? else {
return Ok(None);
};
let bytes = to_bytes(outcome.response.into_body(), usize::MAX)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
serde_json::from_slice::<Value>(&bytes)
.map(Some)
.map_err(|err| GatewayError::Internal(err.to_string()))
}
async fn provider_query_execute_kiro_test_candidate(
state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider,
candidate: &ProviderQueryTestCandidate,
payload: &Value,
route_path: &str,
trace_id: &str,
requested_model: &str,
) -> Result<ProviderQueryExecutionOutcome, GatewayError> {
let Some(transport) = state
.read_provider_transport_snapshot(&provider.id, &candidate.endpoint.id, &candidate.key.id)
.await?
else {
return Ok(ProviderQueryExecutionOutcome {
status: "skipped",
error_message: None,
status_code: None,
latency_ms: None,
request_url: String::new(),
request_headers: BTreeMap::new(),
request_body: Value::Null,
response_headers: BTreeMap::new(),
response_body: None,
});
};
if !supports_local_kiro_request_transport_with_network(&transport) {
return Ok(ProviderQueryExecutionOutcome {
status: "skipped",
error_message: None,
status_code: None,
latency_ms: None,
request_url: String::new(),
request_headers: BTreeMap::new(),
request_body: Value::Null,
response_headers: BTreeMap::new(),
response_body: None,
});
}
let Some(kiro_auth) = state
.resolve_local_oauth_kiro_request_auth(&transport)
.await?
else {
return Ok(ProviderQueryExecutionOutcome {
status: "failed",
error_message: Some("oauth auth failed".to_string()),
status_code: None,
latency_ms: None,
request_url: String::new(),
request_headers: BTreeMap::new(),
request_body: Value::Null,
response_headers: BTreeMap::new(),
response_body: None,
});
};
let request_body = provider_query_build_test_request_body(payload, &candidate.effective_model);
let provider_request_body = match build_kiro_provider_request_body(
&request_body,
&candidate.effective_model,
&kiro_auth.auth_config,
transport.endpoint.body_rules.as_ref(),
) {
Some(body) => body,
None => {
return Ok(ProviderQueryExecutionOutcome {
status: "failed",
error_message: Some("provider request body build failed".to_string()),
status_code: None,
latency_ms: None,
request_url: String::new(),
request_headers: BTreeMap::new(),
request_body,
response_headers: BTreeMap::new(),
response_body: None,
});
}
};
let mut synthetic_request = http::Request::builder()
.uri(route_path)
.body(())
.map_err(|err| GatewayError::Internal(err.to_string()))?;
*synthetic_request.headers_mut() = provider_query_extract_request_headers(payload);
let (parts, _) = synthetic_request.into_parts();
let request_url = build_kiro_generate_assistant_response_url(
&transport.endpoint.base_url,
parts.uri.query(),
Some(kiro_auth.auth_config.effective_api_region()),
)
.ok_or_else(|| GatewayError::Internal("kiro request url is unavailable".to_string()))?;
let request_headers = build_kiro_provider_headers(KiroProviderHeadersInput {
headers: &parts.headers,
provider_request_body: &provider_request_body,
original_request_body: &request_body,
header_rules: transport.endpoint.header_rules.as_ref(),
auth_header: kiro_auth.name,
auth_value: &kiro_auth.value,
auth_config: &kiro_auth.auth_config,
machine_id: kiro_auth.machine_id.as_str(),
})
.ok_or_else(|| GatewayError::Internal("kiro request headers are unavailable".to_string()))?;
let plan = ExecutionPlan {
request_id: trace_id.to_string(),
candidate_id: Some(format!("provider-query-{}", candidate.key.id)),
provider_name: Some(provider.name.clone()),
provider_id: provider.id.clone(),
endpoint_id: candidate.endpoint.id.clone(),
key_id: candidate.key.id.clone(),
method: "POST".to_string(),
url: request_url.clone(),
headers: request_headers.clone(),
content_type: Some("application/json".to_string()),
content_encoding: None,
body: RequestBody::from_json(provider_request_body.clone()),
stream: true,
client_api_format: candidate.endpoint.api_format.clone(),
provider_api_format: candidate.endpoint.api_format.clone(),
model_name: Some(candidate.effective_model.clone()),
proxy: state
.resolve_transport_proxy_snapshot_with_tunnel_affinity(&transport)
.await,
tls_profile: state.resolve_transport_tls_profile(&transport),
timeouts: state.resolve_transport_execution_timeouts(&transport),
};
let result = state
.execute_execution_runtime_sync_plan(Some(trace_id), &plan)
.await?;
let response_body = if result.status_code < 400 {
provider_query_finalize_kiro_result(
route_path,
trace_id,
requested_model,
candidate.endpoint.api_format.as_str(),
candidate.effective_model.as_str(),
&request_body,
&result,
)
.await?
} else {
result.body.as_ref().and_then(|body| body.json_body.clone())
};
let error_message = if result.status_code >= 400 {
provider_query_extract_error_message(&result)
} else if response_body.is_none()
&& provider_query_decode_execution_body(&result)
.is_some_and(|body| crate::ai_pipeline::stream_body_contains_error_event(&body))
{
Some("Kiro upstream returned embedded stream error".to_string())
} else {
None
};
Ok(ProviderQueryExecutionOutcome {
status: if error_message.is_some() {
"failed"
} else {
"success"
},
error_message,
status_code: Some(result.status_code),
latency_ms: result.telemetry.as_ref().and_then(|value| value.elapsed_ms),
request_url,
request_headers,
request_body: provider_request_body,
response_headers: result.headers,
response_body,
})
}
fn provider_query_test_attempt_payload(
candidate_index: usize,
candidate: &ProviderQueryTestCandidate,
execution: &ProviderQueryExecutionOutcome,
) -> Value {
json!({
"candidate_index": candidate_index,
"retry_index": 0,
"endpoint_api_format": candidate.endpoint.api_format,
"endpoint_base_url": candidate.endpoint.base_url,
"key_name": provider_query_key_display_name(&candidate.key),
"key_id": candidate.key.id,
"auth_type": candidate.key.auth_type,
"effective_model": candidate.effective_model,
"status": execution.status,
"skip_reason": Value::Null,
"error_message": execution.error_message,
"status_code": execution.status_code,
"latency_ms": execution.latency_ms,
"request_url": execution.request_url,
"request_headers": execution.request_headers,
"request_body": execution.request_body,
"response_headers": execution.response_headers,
"response_body": execution.response_body,
})
}
async fn build_admin_provider_query_kiro_failover_response(
state: &AdminAppState<'_>,
payload: &Value,
route_path: &str,
) -> Result<Response<Body>, GatewayError> {
let Some(provider_id) = provider_query_extract_provider_id(payload) else {
return Ok(build_admin_provider_query_bad_request_response(
ADMIN_PROVIDER_QUERY_PROVIDER_ID_REQUIRED_DETAIL,
));
};
let requested_model = match provider_query_extract_model(payload) {
Some(model) => model,
None => {
return Ok(build_admin_provider_query_bad_request_response(
ADMIN_PROVIDER_QUERY_MODEL_REQUIRED_DETAIL,
));
}
};
let Some(provider) = state
.app()
.read_provider_catalog_providers_by_ids(std::slice::from_ref(&provider_id))
.await?
.into_iter()
.find(|item| item.id == provider_id)
else {
return Ok(build_admin_provider_query_not_found_response(
ADMIN_PROVIDER_QUERY_PROVIDER_NOT_FOUND_DETAIL,
));
};
if !provider.provider_type.trim().eq_ignore_ascii_case("kiro") {
return Ok(build_admin_provider_query_test_model_failover_response(
provider_id,
super::payload::provider_query_extract_failover_models(payload),
));
}
let candidates =
match provider_query_build_kiro_test_candidates(state, &provider, payload).await {
Ok(candidates) => candidates,
Err(response) => return Ok(response),
};
let trace_id = provider_query_extract_request_id(payload)
.unwrap_or_else(|| format!("provider-query-test-{}", Uuid::new_v4().simple()));
let mut attempts = Vec::new();
let mut total_attempts = 0usize;
let mut success_body = None;
for (candidate_index, candidate) in candidates.iter().enumerate() {
let execution = provider_query_execute_kiro_test_candidate(
state,
&provider,
candidate,
payload,
route_path,
&trace_id,
&requested_model,
)
.await?;
if execution.status != "skipped" {
total_attempts += 1;
}
let is_success = execution.status == "success";
let response_body = execution.response_body.clone();
attempts.push(provider_query_test_attempt_payload(
candidate_index,
candidate,
&execution,
));
if is_success {
success_body = response_body;
break;
}
}
let success = success_body.is_some();
let error = if success {
Value::Null
} else {
attempts
.iter()
.rev()
.find_map(|attempt| {
attempt
.get("error_message")
.cloned()
.filter(|value| !value.is_null())
})
.unwrap_or_else(|| json!(ADMIN_PROVIDER_QUERY_NO_LOCAL_MODELS_DETAIL))
};
Ok(Json(json!({
"success": success,
"model": requested_model,
"provider": provider_query_provider_payload(&provider),
"attempts": attempts,
"total_candidates": candidates.len(),
"total_attempts": total_attempts,
"data": success_body.as_ref().map(|body| json!({
"stream": true,
"response": body,
})),
"error": error,
}))
.into_response())
}
pub(crate) async fn build_admin_provider_query_test_model_local_response(
state: &AdminAppState<'_>,
payload: &Value,
) -> Result<Response<Body>, GatewayError> {
let response = build_admin_provider_query_kiro_failover_response(
state,
payload,
"/api/admin/provider-query/test-model",
)
.await?;
let body = to_bytes(response.into_body(), usize::MAX)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let parsed: Value =
serde_json::from_slice(&body).map_err(|err| GatewayError::Internal(err.to_string()))?;
Ok(Json(json!({
"success": parsed.get("success").cloned().unwrap_or(Value::Bool(false)),
"error": parsed.get("error").cloned().unwrap_or(Value::Null),
"data": parsed.get("data").cloned().unwrap_or(Value::Null),
"provider": parsed.get("provider").cloned().unwrap_or(Value::Null),
"model": parsed.get("model").cloned().unwrap_or(Value::Null),
}))
.into_response())
}
pub(crate) async fn build_admin_provider_query_test_model_failover_local_response(
state: &AdminAppState<'_>,
payload: &Value,
) -> Result<Response<Body>, GatewayError> {
build_admin_provider_query_kiro_failover_response(
state,
payload,
"/api/admin/provider-query/test-model-failover",
)
.await
}
fn provider_query_key_display_name(key: &StoredProviderCatalogKey) -> String {
let trimmed = key.name.trim();
if trimmed.is_empty() {

1
apps/aether-gateway/src/handlers/admin/provider/shared/support.rs
View File

@@ -13,6 +13,7 @@ pub(crate) const ADMIN_PROVIDER_OAUTH_DATA_UNAVAILABLE_DETAIL: &str =
#[derive(Debug, Clone, Copy)]
pub(crate) struct AdminProviderPoolConfig {
pub(crate) lru_enabled: bool,
pub(crate) skip_exhausted_accounts: bool,
pub(crate) cost_window_seconds: u64,
pub(crate) cost_limit_per_key_tokens: Option<u64>,
}

14
apps/aether-gateway/src/handlers/admin/provider/write/normalize.rs
View File

@@ -13,8 +13,8 @@ pub(crate) fn normalize_provider_type_input(value: &str) -> Result<String, Strin
pub(crate) fn normalize_auth_type(value: Option<&str>) -> Result<String, String> {
let auth_type = value.unwrap_or("api_key").trim().to_ascii_lowercase();
match auth_type.as_str() {
"api_key" | "service_account" | "oauth" => Ok(auth_type),
_ => Err("auth_type 仅支持 api_key / service_account / oauth".to_string()),
"api_key" | "service_account" | "oauth" | "bearer" => Ok(auth_type),
_ => Err("auth_type 仅支持 api_key / service_account / oauth / bearer".to_string()),
}
}
@@ -63,7 +63,7 @@ pub(crate) fn validate_vertex_api_formats(
#[cfg(test)]
mod tests {
use super::normalize_pool_advanced_config;
use super::{normalize_auth_type, normalize_pool_advanced_config};
use serde_json::json;
#[test]
@@ -85,4 +85,12 @@ mod tests {
"pool_advanced 必须是 JSON 对象"
);
}
#[test]
fn normalize_auth_type_supports_bearer() {
assert_eq!(
normalize_auth_type(Some("bearer")).expect("bearer should normalize"),
"bearer"
);
}
}

32
apps/aether-gateway/src/handlers/admin/provider/write/reveal.rs
View File

@@ -8,17 +8,37 @@ fn normalize_reveal_auth_type(value: &str) -> &str {
use crate::handlers::admin::request::AdminAppState;
use crate::handlers::admin::shared::parse_catalog_auth_config_json;
use crate::provider_key_auth::provider_key_auth_semantics;
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use chrono::{SecondsFormat, Utc};
use serde_json::json;
fn reveal_provider_type_from_auth_config(
auth_config: Option<&serde_json::Map<String, serde_json::Value>>,
) -> String {
auth_config
.and_then(|value| value.get("provider_type"))
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or_default()
.to_string()
}
pub(crate) fn build_admin_reveal_key_payload(
state: &AdminAppState<'_>,
key: &StoredProviderCatalogKey,
) -> Result<serde_json::Value, String> {
let auth_type = normalize_reveal_auth_type(&key.auth_type);
let parsed_auth_config = state.parse_catalog_auth_config_json(key);
let provider_type = reveal_provider_type_from_auth_config(parsed_auth_config.as_ref());
let auth_semantics = provider_key_auth_semantics(key, provider_type.as_str());
let auth_type = if auth_semantics.oauth_managed() {
"oauth"
} else {
normalize_reveal_auth_type(&key.auth_type)
};
if matches!(auth_type, "service_account") {
if let Some(auth_config) = state.parse_catalog_auth_config_json(key) {
if let Some(auth_config) = parsed_auth_config {
return Ok(json!({
"auth_type": auth_type,
"auth_config": auth_config,
@@ -92,11 +112,6 @@ pub(crate) async fn build_admin_export_key_payload(
state: &AdminAppState<'_>,
key: &StoredProviderCatalogKey,
) -> Result<serde_json::Value, String> {
let auth_type = normalize_reveal_auth_type(&key.auth_type);
if auth_type != "oauth" {
return Err("仅 OAuth 类型的 Key 支持导出".to_string());
}
let ciphertext = key
.encrypted_auth_config
.as_deref()
@@ -136,6 +151,9 @@ pub(crate) async fn build_admin_export_key_payload(
.map(|provider| provider.provider_type)
.unwrap_or_default()
};
if !provider_key_auth_semantics(key, provider_type.as_str()).can_export_oauth() {
return Err("仅 OAuth 管理账号支持导出".to_string());
}
let mut payload =
provider_oauth_export_payload(&provider_type, &auth_config, key.upstream_metadata.as_ref());

54
apps/aether-gateway/src/handlers/admin/request/provider/routes/query.rs
View File

@@ -1,7 +1,9 @@
use crate::handlers::admin::provider::query::{
models::{
build_admin_provider_query_models_response,
build_admin_provider_query_test_model_failover_local_response,
build_admin_provider_query_test_model_failover_response,
build_admin_provider_query_test_model_local_response,
build_admin_provider_query_test_model_response,
},
payload::{
@@ -78,10 +80,25 @@ impl<'a> AdminAppState<'a> {
ADMIN_PROVIDER_QUERY_MODEL_REQUIRED_DETAIL,
)));
};
Ok(Some(build_admin_provider_query_test_model_response(
provider_id,
model,
)))
let provider_type = self
.app()
.read_provider_catalog_providers_by_ids(std::slice::from_ref(&provider_id))
.await?
.into_iter()
.find(|provider| provider.id == provider_id)
.map(|provider| provider.provider_type)
.unwrap_or_default();
if provider_type.trim().eq_ignore_ascii_case("kiro") {
Ok(Some(
build_admin_provider_query_test_model_local_response(self, &payload)
.await?,
))
} else {
Ok(Some(build_admin_provider_query_test_model_response(
provider_id,
model,
)))
}
}
"test_model_failover" => {
let Some(provider_id) = provider_query_extract_provider_id(&payload) else {
@@ -107,12 +124,29 @@ impl<'a> AdminAppState<'a> {
ADMIN_PROVIDER_QUERY_FAILOVER_MODELS_REQUIRED_DETAIL,
)));
}
Ok(Some(
build_admin_provider_query_test_model_failover_response(
provider_id,
failover_models,
),
))
let provider_type = self
.app()
.read_provider_catalog_providers_by_ids(std::slice::from_ref(&provider_id))
.await?
.into_iter()
.find(|provider| provider.id == provider_id)
.map(|provider| provider.provider_type)
.unwrap_or_default();
if provider_type.trim().eq_ignore_ascii_case("kiro") {
Ok(Some(
build_admin_provider_query_test_model_failover_local_response(
self, &payload,
)
.await?,
))
} else {
Ok(Some(
build_admin_provider_query_test_model_failover_response(
provider_id,
failover_models,
),
))
}
}
_ => Ok(Some(
build_admin_provider_query_models_response(self, &payload).await?,

115
apps/aether-gateway/src/handlers/shared/catalog.rs
View File

@@ -1,4 +1,5 @@
use crate::handlers::shared::{json_string_list, unix_secs_to_rfc3339};
use crate::provider_key_auth::provider_key_auth_semantics;
use crate::AppState;
#[cfg(test)]
use aether_crypto::DEVELOPMENT_ENCRYPTION_KEY;
@@ -273,7 +274,7 @@ fn derive_catalog_oauth_plan_type(
provider_type: &str,
auth_config: Option<&serde_json::Map<String, serde_json::Value>>,
) -> Option<String> {
if !key.auth_type.trim().eq_ignore_ascii_case("oauth") {
if !provider_key_auth_semantics(key, provider_type).oauth_managed() {
return None;
}
@@ -339,13 +340,18 @@ pub(crate) fn build_admin_provider_key_response(
} else {
0.0
};
let auth_semantics = provider_key_auth_semantics(key, provider_type);
let auth_config = parse_catalog_auth_config_json(state, key);
let oauth_organizations = auth_config
.as_ref()
.and_then(|config| config.get("organizations"))
.and_then(serde_json::Value::as_array)
.cloned()
.unwrap_or_default();
let oauth_organizations = if auth_semantics.can_show_oauth_metadata() {
auth_config
.as_ref()
.and_then(|config| config.get("organizations"))
.and_then(serde_json::Value::as_array)
.cloned()
.unwrap_or_default()
} else {
Vec::new()
};
let oauth_plan_type = derive_catalog_oauth_plan_type(key, provider_type, auth_config.as_ref());
let (
health_score,
@@ -387,6 +393,30 @@ pub(crate) fn build_admin_provider_key_response(
);
payload.insert("api_key_plain".to_string(), serde_json::Value::Null);
payload.insert("auth_type".to_string(), json!(key.auth_type));
payload.insert(
"credential_kind".to_string(),
json!(auth_semantics.credential_kind().as_str()),
);
payload.insert(
"runtime_auth_kind".to_string(),
json!(auth_semantics.runtime_auth_kind().as_str()),
);
payload.insert(
"oauth_managed".to_string(),
json!(auth_semantics.oauth_managed()),
);
payload.insert(
"can_refresh_oauth".to_string(),
json!(auth_semantics.can_refresh_oauth()),
);
payload.insert(
"can_export_oauth".to_string(),
json!(auth_semantics.can_export_oauth()),
);
payload.insert(
"can_edit_oauth".to_string(),
json!(auth_semantics.can_edit_oauth()),
);
payload.insert("name".to_string(), json!(key.name));
payload.insert("rate_multipliers".to_string(), json!(key.rate_multipliers));
payload.insert(
@@ -410,40 +440,59 @@ pub(crate) fn build_admin_provider_key_response(
payload.insert("capabilities".to_string(), json!(key.capabilities));
payload.insert(
"oauth_expires_at".to_string(),
json!(key.expires_at_unix_secs),
json!(auth_semantics
.can_show_oauth_metadata()
.then_some(key.expires_at_unix_secs)
.flatten()),
);
payload.insert(
"oauth_email".to_string(),
auth_config
.as_ref()
.and_then(|config| config.get("email"))
.cloned()
.unwrap_or(serde_json::Value::Null),
if auth_semantics.can_show_oauth_metadata() {
auth_config
.as_ref()
.and_then(|config| config.get("email"))
.cloned()
.unwrap_or(serde_json::Value::Null)
} else {
serde_json::Value::Null
},
);
payload.insert("oauth_plan_type".to_string(), json!(oauth_plan_type));
payload.insert(
"oauth_account_id".to_string(),
auth_config
.as_ref()
.and_then(|config| config.get("account_id"))
.cloned()
.unwrap_or(serde_json::Value::Null),
if auth_semantics.can_show_oauth_metadata() {
auth_config
.as_ref()
.and_then(|config| config.get("account_id"))
.cloned()
.unwrap_or(serde_json::Value::Null)
} else {
serde_json::Value::Null
},
);
payload.insert(
"oauth_account_name".to_string(),
auth_config
.as_ref()
.and_then(|config| config.get("account_name"))
.cloned()
.unwrap_or(serde_json::Value::Null),
if auth_semantics.can_show_oauth_metadata() {
auth_config
.as_ref()
.and_then(|config| config.get("account_name"))
.cloned()
.unwrap_or(serde_json::Value::Null)
} else {
serde_json::Value::Null
},
);
payload.insert(
"oauth_account_user_id".to_string(),
auth_config
.as_ref()
.and_then(|config| config.get("account_user_id"))
.cloned()
.unwrap_or(serde_json::Value::Null),
if auth_semantics.can_show_oauth_metadata() {
auth_config
.as_ref()
.and_then(|config| config.get("account_user_id"))
.cloned()
.unwrap_or(serde_json::Value::Null)
} else {
serde_json::Value::Null
},
);
payload.insert(
"oauth_organizations".to_string(),
@@ -451,11 +500,17 @@ pub(crate) fn build_admin_provider_key_response(
);
payload.insert(
"oauth_invalid_at".to_string(),
json!(key.oauth_invalid_at_unix_secs),
json!(auth_semantics
.can_show_oauth_metadata()
.then_some(key.oauth_invalid_at_unix_secs)
.flatten()),
);
payload.insert(
"oauth_invalid_reason".to_string(),
json!(key.oauth_invalid_reason),
json!(auth_semantics
.can_show_oauth_metadata()
.then_some(key.oauth_invalid_reason.clone())
.flatten()),
);
payload.insert(
"status_snapshot".to_string(),