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

feat: 全栈功能增强 - 扩展 provider/pool 管理、完善调度与数据层、重构前端 Pool 页面

Browse Source 
后端:
- 扩展 pool_admin payloads 和 provider query models，增强 endpoint key 管理
- 完善 scheduler-core 候选排序与请求候选逻辑
- 增强 usage-runtime 写入、provider-transport 网络层与 OAuth 刷新
- 改进 AI pipeline 响应转换与流式处理
- 扩展 global_models/provider_catalog 数据层查询能力
- 增强 video-tasks-core 多 provider 支持
- 新增大量集成测试覆盖 pool/keys/provider_query/frontdoor

前端:
- 重构 PoolManagement 页面，拆分状态管理/对话框逻辑到独立模块
- 新增 poolAdvancedDialog/poolSchedulingDialog/poolManagementState/poolMobilePresentation 工具函数及测试
- 改进 Dialog 组件与 provider tabs 显示

部署:
- 更新 Rust CI workflow 和 Dockerfile 构建配置

Closes #275
Co-authored-by: AAEE86 <ppk0227@hotmail.com>
This commit is contained in:
fawney19
2026-04-09 13:51:50 +08:00
parent 4fc95adfb9
commit b0b40c16ff
97 changed files with 5816 additions and 1881 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
apps/aether-gateway/src/ai_pipeline/planner/decision/stream.rs
View File

@@ -106,6 +106,7 @@ async fn maybe_build_local_video_task_content_stream_decision_payload(
let crate::video_tasks::LocalVideoTaskContentAction::StreamPlan(plan) = action else {
return Ok(None);
};
let plan = *plan;
let provider_contract = plan.provider_api_format.clone();
let client_contract = plan.client_api_format.clone();
let execution_strategy = if plan.provider_api_format == plan.client_api_format {

24
apps/aether-gateway/src/ai_pipeline/planner/passthrough/provider/family/payload.rs
View File

@@ -9,7 +9,9 @@ use crate::ai_pipeline::transport::antigravity::{
};
use crate::ai_pipeline::transport::auth::build_openai_passthrough_headers;
use crate::ai_pipeline::transport::claude_code::build_claude_code_passthrough_headers;
use crate::ai_pipeline::transport::kiro::{build_kiro_provider_headers, KIRO_ENVELOPE_NAME};
use crate::ai_pipeline::transport::kiro::{
build_kiro_provider_headers, KiroProviderHeadersInput, KIRO_ENVELOPE_NAME,
};
use crate::ai_pipeline::transport::{
apply_local_header_rules, ensure_upstream_auth_header, resolve_transport_execution_timeouts,
resolve_transport_proxy_snapshot_with_tunnel_affinity, resolve_transport_tls_profile,
@@ -171,16 +173,16 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
};
let Some(provider_request_headers) = (if let Some(kiro_auth) = kiro_auth.as_ref() {
build_kiro_provider_headers(
&parts.headers,
&provider_request_body,
body_json,
transport.endpoint.header_rules.as_ref(),
auth_header.as_deref().unwrap_or_default(),
auth_value.as_deref().unwrap_or_default(),
&kiro_auth.auth_config,
kiro_auth.machine_id.as_str(),
)
build_kiro_provider_headers(KiroProviderHeadersInput {
headers: &parts.headers,
provider_request_body: &provider_request_body,
original_request_body: body_json,
header_rules: transport.endpoint.header_rules.as_ref(),
auth_header: auth_header.as_deref().unwrap_or_default(),
auth_value: auth_value.as_deref().unwrap_or_default(),
auth_config: &kiro_auth.auth_config,
machine_id: kiro_auth.machine_id.as_str(),
})
} else {
let extra_headers = antigravity_auth
.as_ref()

68
apps/aether-gateway/src/api/response.rs
View File

@@ -127,40 +127,6 @@ pub(crate) fn build_client_response_from_parts(
Ok(response)
}
#[cfg(test)]
mod tests {
use super::build_client_response_from_parts;
use axum::body::Body;
use std::collections::BTreeMap;
#[test]
fn sse_responses_disable_proxy_buffering() {
let response = build_client_response_from_parts(
200,
&BTreeMap::from([("content-type".to_string(), "text/event-stream".to_string())]),
Body::from("data: hello\n\n"),
"trace-sse-buffering-1",
None,
)
.expect("response should build");
assert_eq!(
response
.headers()
.get(http::header::CACHE_CONTROL)
.and_then(|value| value.to_str().ok()),
Some("no-cache, no-transform")
);
assert_eq!(
response
.headers()
.get("x-accel-buffering")
.and_then(|value| value.to_str().ok()),
Some("no")
);
}
}
pub(crate) fn insert_candidate_id_header_if_present(
headers: &mut http::HeaderMap,
candidate_id: Option<&str>,
@@ -357,3 +323,37 @@ pub(crate) fn build_local_overloaded_response(
control_decision,
)
}
#[cfg(test)]
mod tests {
use super::build_client_response_from_parts;
use axum::body::Body;
use std::collections::BTreeMap;
#[test]
fn sse_responses_disable_proxy_buffering() {
let response = build_client_response_from_parts(
200,
&BTreeMap::from([("content-type".to_string(), "text/event-stream".to_string())]),
Body::from("data: hello\n\n"),
"trace-sse-buffering-1",
None,
)
.expect("response should build");
assert_eq!(
response
.headers()
.get(http::header::CACHE_CONTROL)
.and_then(|value| value.to_str().ok()),
Some("no-cache, no-transform")
);
assert_eq!(
response
.headers()
.get("x-accel-buffering")
.and_then(|value| value.to_str().ok()),
Some("no")
);
}
}

4
apps/aether-gateway/src/control/route/admin.rs
View File

@@ -72,9 +72,7 @@ pub(super) fn classify_admin_route(
} else if let Some(route) = classify_admin_model_provider_family_route(method, normalized_path)
{
Some(route)
} else if let Some(route) = classify_admin_endpoints_family_route(method, normalized_path) {
Some(route)
} else {
None
classify_admin_endpoints_family_route(method, normalized_path)
}
}

135
apps/aether-gateway/src/execution_runtime/stream/execution.rs
View File

@@ -2,6 +2,7 @@ use std::io::Error as IoError;
use aether_contracts::{ExecutionPlan, ExecutionTelemetry, StreamFramePayload};
use aether_data_contracts::repository::candidates::RequestCandidateStatus;
use aether_scheduler_core::SchedulerRequestCandidateStatusUpdate;
use async_stream::stream;
use axum::body::{Body, Bytes};
use axum::http::Response;
@@ -169,16 +170,18 @@ pub(crate) async fn execute_execution_runtime_stream(
state,
&plan,
report_context.as_ref(),
RequestCandidateStatus::Failed,
Some(response.status().as_u16()),
Some("execution_runtime_http_error".to_string()),
Some(format!(
"execution runtime returned HTTP {}",
response.status()
)),
None,
Some(terminal_unix_secs),
Some(terminal_unix_secs),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Failed,
status_code: Some(response.status().as_u16()),
error_type: Some("execution_runtime_http_error".to_string()),
error_message: Some(format!(
"execution runtime returned HTTP {}",
response.status()
)),
latency_ms: None,
started_at_unix_secs: Some(terminal_unix_secs),
finished_at_unix_secs: Some(terminal_unix_secs),
},
)
.await;
return Ok(Some(attach_control_metadata_headers(
@@ -241,15 +244,17 @@ async fn execute_stream_from_frame_stream(
state,
&plan,
report_context.as_ref(),
RequestCandidateStatus::Failed,
Some(status_code),
Some("retryable_upstream_status".to_string()),
Some(format!(
"execution runtime stream returned retryable status {status_code}"
)),
None,
Some(terminal_unix_secs),
Some(terminal_unix_secs),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Failed,
status_code: Some(status_code),
error_type: Some("retryable_upstream_status".to_string()),
error_message: Some(format!(
"execution runtime stream returned retryable status {status_code}"
)),
latency_ms: None,
started_at_unix_secs: Some(terminal_unix_secs),
finished_at_unix_secs: Some(terminal_unix_secs),
},
)
.await;
warn!(
@@ -276,15 +281,17 @@ async fn execute_stream_from_frame_stream(
state,
&plan,
report_context.as_ref(),
RequestCandidateStatus::Failed,
Some(status_code),
Some("control_fallback".to_string()),
Some(format!(
"stream decision fell back to control after status {status_code}"
)),
None,
Some(terminal_unix_secs),
Some(terminal_unix_secs),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Failed,
status_code: Some(status_code),
error_type: Some("control_fallback".to_string()),
error_message: Some(format!(
"stream decision fell back to control after status {status_code}"
)),
latency_ms: None,
started_at_unix_secs: Some(terminal_unix_secs),
finished_at_unix_secs: Some(terminal_unix_secs),
},
)
.await;
return Ok(None);
@@ -322,15 +329,17 @@ async fn execute_stream_from_frame_stream(
state,
&plan,
report_context.as_ref(),
RequestCandidateStatus::Failed,
Some(status_code),
Some("execution_runtime_stream_error".to_string()),
Some(format!(
"execution runtime stream returned error status {status_code}"
)),
None,
Some(terminal_unix_secs),
Some(terminal_unix_secs),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Failed,
status_code: Some(status_code),
error_type: Some("execution_runtime_stream_error".to_string()),
error_message: Some(format!(
"execution runtime stream returned error status {status_code}"
)),
latency_ms: None,
started_at_unix_secs: Some(terminal_unix_secs),
finished_at_unix_secs: Some(terminal_unix_secs),
},
)
.await;
if let Some(report_kind) = stream_error_finalize_kind {
@@ -632,15 +641,17 @@ async fn execute_stream_from_frame_stream(
state,
&plan,
report_context.as_ref(),
RequestCandidateStatus::Streaming,
Some(status_code),
None,
None,
prefetched_telemetry
.as_ref()
.and_then(|telemetry| telemetry.elapsed_ms),
Some(candidate_started_unix_secs),
None,
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Streaming,
status_code: Some(status_code),
error_type: None,
error_message: None,
latency_ms: prefetched_telemetry
.as_ref()
.and_then(|telemetry| telemetry.elapsed_ms),
started_at_unix_secs: Some(candidate_started_unix_secs),
finished_at_unix_secs: None,
},
)
.await;
@@ -981,13 +992,15 @@ async fn execute_stream_from_frame_stream(
&state_for_report,
&plan_for_report,
report_context_owned.as_ref(),
RequestCandidateStatus::Cancelled,
Some(499),
Some("downstream_disconnect".to_string()),
Some("client disconnected before stream completion".to_string()),
telemetry.as_ref().and_then(|value| value.elapsed_ms),
Some(candidate_started_unix_secs_for_report),
Some(current_request_candidate_unix_secs()),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Cancelled,
status_code: Some(499),
error_type: Some("downstream_disconnect".to_string()),
error_message: Some("client disconnected before stream completion".to_string()),
latency_ms: telemetry.as_ref().and_then(|value| value.elapsed_ms),
started_at_unix_secs: Some(candidate_started_unix_secs_for_report),
finished_at_unix_secs: Some(current_request_candidate_unix_secs()),
},
)
.await;
return;
@@ -1036,13 +1049,15 @@ async fn execute_stream_from_frame_stream(
&state_for_report,
&plan_for_report,
report_context_owned.as_ref(),
RequestCandidateStatus::Success,
Some(status_code),
None,
None,
telemetry.as_ref().and_then(|value| value.elapsed_ms),
Some(candidate_started_unix_secs_for_report),
Some(current_request_candidate_unix_secs()),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Success,
status_code: Some(status_code),
error_type: None,
error_message: None,
latency_ms: telemetry.as_ref().and_then(|value| value.elapsed_ms),
started_at_unix_secs: Some(candidate_started_unix_secs_for_report),
finished_at_unix_secs: Some(current_request_candidate_unix_secs()),
},
)
.await;

23
apps/aether-gateway/src/execution_runtime/stream/execution_failures.rs
View File

@@ -1,5 +1,6 @@
use aether_contracts::{ExecutionError, ExecutionPlan, ExecutionTelemetry};
use aether_data_contracts::repository::candidates::RequestCandidateStatus;
use aether_scheduler_core::SchedulerRequestCandidateStatusUpdate;
use axum::body::Body;
use axum::http::Response;
use base64::Engine as _;
@@ -126,16 +127,18 @@ async fn record_stream_sync_failure(
record_report_request_candidate_status(
state,
report_context,
RequestCandidateStatus::Failed,
Some(failure.status_code),
Some(failure.error_type.clone()),
Some(failure.error_message.clone()),
payload
.telemetry
.as_ref()
.and_then(|telemetry| telemetry.elapsed_ms),
started_at_unix_secs.or(Some(terminal_unix_secs)),
Some(terminal_unix_secs),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Failed,
status_code: Some(failure.status_code),
error_type: Some(failure.error_type.clone()),
error_message: Some(failure.error_message.clone()),
latency_ms: payload
.telemetry
.as_ref()
.and_then(|telemetry| telemetry.elapsed_ms),
started_at_unix_secs: started_at_unix_secs.or(Some(terminal_unix_secs)),
finished_at_unix_secs: Some(terminal_unix_secs),
},
)
.await;
}

78
apps/aether-gateway/src/execution_runtime/sync/execution.rs
View File

@@ -2,7 +2,7 @@ use std::collections::BTreeMap;
use aether_contracts::{ExecutionPlan, ExecutionResult, ExecutionTelemetry};
use aether_data_contracts::repository::candidates::RequestCandidateStatus;
use aether_scheduler_core::execution_error_details;
use aether_scheduler_core::{execution_error_details, SchedulerRequestCandidateStatusUpdate};
use axum::body::Body;
use axum::http::Response;
use base64::Engine as _;
@@ -165,13 +165,15 @@ pub(crate) async fn execute_execution_runtime_sync(
state,
&plan,
report_context.as_ref(),
RequestCandidateStatus::Failed,
Some(result.status_code),
result_error_type.clone(),
result_error_message.clone(),
result_latency_ms,
Some(terminal_unix_secs),
Some(terminal_unix_secs),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Failed,
status_code: Some(result.status_code),
error_type: result_error_type.clone(),
error_message: result_error_message.clone(),
latency_ms: result_latency_ms,
started_at_unix_secs: Some(terminal_unix_secs),
finished_at_unix_secs: Some(terminal_unix_secs),
},
)
.await;
warn!(
@@ -231,13 +233,15 @@ pub(crate) async fn execute_execution_runtime_sync(
state,
&plan,
report_context.as_ref(),
RequestCandidateStatus::Failed,
Some(result.status_code),
result_error_type.clone(),
result_error_message.clone(),
result_latency_ms,
Some(terminal_unix_secs),
Some(terminal_unix_secs),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Failed,
status_code: Some(result.status_code),
error_type: result_error_type.clone(),
error_message: result_error_message.clone(),
latency_ms: result_latency_ms,
started_at_unix_secs: Some(terminal_unix_secs),
finished_at_unix_secs: Some(terminal_unix_secs),
},
)
.await;
return Ok(None);
@@ -252,17 +256,19 @@ pub(crate) async fn execute_execution_runtime_sync(
state,
&plan,
report_context.as_ref(),
if result.status_code >= 400 {
RequestCandidateStatus::Failed
} else {
RequestCandidateStatus::Success
SchedulerRequestCandidateStatusUpdate {
status: if result.status_code >= 400 {
RequestCandidateStatus::Failed
} else {
RequestCandidateStatus::Success
},
status_code: Some(result.status_code),
error_type: result_error_type.clone(),
error_message: result_error_message.clone(),
latency_ms: result_latency_ms,
started_at_unix_secs: Some(terminal_unix_secs),
finished_at_unix_secs: Some(terminal_unix_secs),
},
Some(result.status_code),
result_error_type.clone(),
result_error_message.clone(),
result_latency_ms,
Some(terminal_unix_secs),
Some(terminal_unix_secs),
)
.await;
@@ -603,16 +609,18 @@ async fn execute_sync_via_remote_execution_runtime(
state,
plan,
report_context,
RequestCandidateStatus::Failed,
Some(response.status().as_u16()),
Some("execution_runtime_http_error".to_string()),
Some(format!(
"execution runtime returned HTTP {}",
response.status()
)),
None,
Some(terminal_unix_secs),
Some(terminal_unix_secs),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Failed,
status_code: Some(response.status().as_u16()),
error_type: Some("execution_runtime_http_error".to_string()),
error_message: Some(format!(
"execution runtime returned HTTP {}",
response.status()
)),
latency_ms: None,
started_at_unix_secs: Some(terminal_unix_secs),
finished_at_unix_secs: Some(terminal_unix_secs),
},
)
.await;
return Ok(RemoteSyncFallbackOutcome::ClientResponse(

33
apps/aether-gateway/src/executor/candidate_loop.rs
View File

@@ -2,6 +2,7 @@ use axum::body::Body;
use axum::http::Response;
use aether_data_contracts::repository::candidates::RequestCandidateStatus;
use aether_scheduler_core::SchedulerRequestCandidateStatusUpdate;
use crate::ai_pipeline_api::{LocalStreamPlanAndReport, LocalSyncPlanAndReport};
use crate::control::GatewayControlDecision;
@@ -118,13 +119,15 @@ where
state,
plan_and_report.plan(),
plan_and_report.report_context().as_ref(),
RequestCandidateStatus::Unused,
None,
None,
None,
None,
None,
None,
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Unused,
status_code: None,
error_type: None,
error_message: None,
latency_ms: None,
started_at_unix_secs: None,
finished_at_unix_secs: None,
},
)
.await;
}
@@ -144,13 +147,15 @@ pub(crate) async fn mark_unused_local_candidate_items<T, FPlan, FContext>(
state,
plan(&item),
report_context(&item),
RequestCandidateStatus::Unused,
None,
None,
None,
None,
None,
None,
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Unused,
status_code: None,
error_type: None,
error_message: None,
latency_ms: None,
started_at_unix_secs: None,
finished_at_unix_secs: None,
},
)
.await;
}

6
apps/aether-gateway/src/executor/stream_path.rs
View File

@@ -169,8 +169,10 @@ async fn maybe_execute_local_video_task_content_stream(
&body_json,
)?)),
crate::video_tasks::LocalVideoTaskContentAction::StreamPlan(plan) => {
execute_execution_runtime_stream(state, plan, trace_id, decision, plan_kind, None, None)
.await
execute_execution_runtime_stream(
state, *plan, trace_id, decision, plan_kind, None, None,
)
.await
}
}
}

5
apps/aether-gateway/src/handlers/admin/billing/mod.rs
View File

@@ -285,7 +285,6 @@ pub(crate) async fn maybe_build_local_admin_billing_response(
return Ok(Some(response));
}
match decision.route_kind.as_deref() {
_ => Ok(Some(build_admin_billing_data_unavailable_response())),
}
let _ = decision.route_kind.as_deref();
Ok(Some(build_admin_billing_data_unavailable_response()))
}

8
apps/aether-gateway/src/handlers/admin/observability/monitoring/cache_affinity_reads.rs
View File

@@ -246,9 +246,11 @@ pub(super) async fn build_admin_monitoring_cache_affinity_response(
"缺少 user_identifier",
));
};
let direct_api_key_by_id =
admin_monitoring_list_export_api_key_records_by_ids(state, &[user_identifier.clone()])
.await?;
let direct_api_key_by_id = admin_monitoring_list_export_api_key_records_by_ids(
state,
std::slice::from_ref(&user_identifier),
)
.await?;
let direct_affinity_keys =
std::iter::once(user_identifier.clone()).collect::<std::collections::BTreeSet<_>>();
let direct_affinities =

7
apps/aether-gateway/src/handlers/admin/provider/endpoint_keys/mutations/create.rs
View File

@@ -70,7 +70,12 @@ pub(super) async fn maybe_handle(
.unwrap_or(0);
Ok(Some(
Json(state.build_admin_provider_key_response(&created, now_unix_secs)).into_response(),
Json(state.build_admin_provider_key_response(
&created,
&provider.provider_type,
now_unix_secs,
))
.into_response(),
))
}

7
apps/aether-gateway/src/handlers/admin/provider/endpoint_keys/mutations/update.rs
View File

@@ -88,7 +88,12 @@ pub(super) async fn maybe_handle(
.unwrap_or(0);
Ok(Some(
Json(state.build_admin_provider_key_response(&updated, now_unix_secs)).into_response(),
Json(state.build_admin_provider_key_response(
&updated,
&provider.provider_type,
now_unix_secs,
))
.into_response(),
))
}

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

@@ -1,34 +1,820 @@
use crate::handlers::admin::provider::shared::support::{
AdminProviderPoolConfig, AdminProviderPoolRuntimeState,
};
use aether_admin::provider::pool as admin_provider_pool_pure;
use crate::handlers::admin::shared::{
provider_key_status_snapshot_payload, unix_secs_to_rfc3339,
};
use crate::handlers::admin::request::AdminAppState;
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use serde_json::json;
pub(super) fn admin_pool_api_formats(key: &StoredProviderCatalogKey) -> Vec<String> {
key.api_formats
.as_ref()
.and_then(serde_json::Value::as_array)
.map(|values| {
values
.iter()
.filter_map(serde_json::Value::as_str)
.map(ToOwned::to_owned)
.collect::<Vec<_>>()
})
.unwrap_or_default()
}
fn admin_pool_string_list(value: Option<&serde_json::Value>) -> Option<Vec<String>> {
let values = value
.and_then(serde_json::Value::as_array)
.map(|items| {
items
.iter()
.filter_map(serde_json::Value::as_str)
.map(ToOwned::to_owned)
.collect::<Vec<_>>()
})
.unwrap_or_default();
if values.is_empty() {
None
} else {
Some(values)
}
}
fn admin_pool_json_object(
value: Option<&serde_json::Value>,
) -> Option<serde_json::Map<String, serde_json::Value>> {
value
.and_then(serde_json::Value::as_object)
.cloned()
.filter(|value| !value.is_empty())
}
fn admin_pool_json_to_f64(value: Option<&serde_json::Value>) -> Option<f64> {
let parsed = match value {
Some(serde_json::Value::Number(number)) => number.as_f64(),
Some(serde_json::Value::String(text)) => text.trim().parse::<f64>().ok(),
_ => None,
}?;
if parsed.is_finite() {
Some(parsed)
} else {
None
}
}
fn admin_pool_json_to_u64(value: Option<&serde_json::Value>) -> Option<u64> {
let mut parsed = match value {
Some(serde_json::Value::Number(number)) => number.as_f64(),
Some(serde_json::Value::String(text)) => text.trim().parse::<f64>().ok(),
_ => None,
}?;
if !parsed.is_finite() || parsed <= 0.0 {
return None;
}
if parsed > 1_000_000_000_000.0 {
parsed /= 1000.0;
}
Some(parsed.floor() as u64)
}
fn admin_pool_trimmed_string(value: Option<&serde_json::Value>) -> Option<String> {
value
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
fn admin_pool_trimmed_string_from_map(
value: Option<&serde_json::Map<String, serde_json::Value>>,
field: &str,
) -> Option<String> {
admin_pool_trimmed_string(value.and_then(|object| object.get(field)))
}
fn admin_pool_oauth_organizations(
auth_config: Option<&serde_json::Map<String, serde_json::Value>>,
) -> Vec<serde_json::Value> {
auth_config
.and_then(|config| config.get("organizations"))
.and_then(serde_json::Value::as_array)
.cloned()
.unwrap_or_default()
}
fn admin_pool_normalize_oauth_plan_type(value: &str, provider_type: &str) -> Option<String> {
let mut normalized = value.trim().to_string();
if normalized.is_empty() {
return None;
}
let provider_type = provider_type.trim().to_ascii_lowercase();
if !provider_type.is_empty() && normalized.to_ascii_lowercase().starts_with(&provider_type) {
normalized = normalized[provider_type.len()..]
.trim_matches(|ch: char| [' ', ':', '-', '_'].contains(&ch))
.to_string();
}
let normalized = normalized.trim().to_ascii_lowercase();
if normalized.is_empty() {
None
} else {
Some(normalized)
}
}
fn admin_pool_derive_oauth_expires_at(
key: &StoredProviderCatalogKey,
auth_config: Option<&serde_json::Map<String, serde_json::Value>>,
) -> Option<u64> {
if !key.auth_type.trim().eq_ignore_ascii_case("oauth") {
return None;
}
for field in ["expires_at", "expiresAt", "expiry", "exp"] {
let expires_at = admin_pool_json_to_u64(auth_config.and_then(|config| config.get(field)));
if expires_at.is_some() {
return expires_at;
}
}
key.expires_at_unix_secs
}
fn admin_pool_derive_oauth_plan_type(
key: &StoredProviderCatalogKey,
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") {
return None;
}
if let Some(upstream_metadata) = key
.upstream_metadata
.as_ref()
.and_then(serde_json::Value::as_object)
{
let provider_bucket = upstream_metadata
.get(&provider_type.trim().to_ascii_lowercase())
.and_then(serde_json::Value::as_object);
for source in provider_bucket
.into_iter()
.chain(std::iter::once(upstream_metadata))
{
for field in [
"plan_type",
"tier",
"subscription_title",
"subscription_plan",
] {
if let Some(value) = source.get(field).and_then(serde_json::Value::as_str) {
let normalized = admin_pool_normalize_oauth_plan_type(value, provider_type);
if normalized.is_some() {
return normalized;
}
}
}
}
}
if let Some(config) = auth_config {
for field in ["plan_type", "tier", "plan", "subscription_plan"] {
if let Some(value) = config.get(field).and_then(serde_json::Value::as_str) {
let normalized = admin_pool_normalize_oauth_plan_type(value, provider_type);
if normalized.is_some() {
return normalized;
}
}
}
}
None
}
fn admin_pool_format_percent(value: f64) -> String {
format!("{:.1}%", value.clamp(0.0, 100.0))
}
fn admin_pool_format_quota_value(value: f64) -> String {
let rounded = value.round();
if (value - rounded).abs() < 1e-6 {
rounded.to_string()
} else {
format!("{value:.1}")
}
}
fn admin_pool_has_quota_consumption(used_percent: Option<f64>) -> bool {
used_percent
.map(|value| value.clamp(0.0, 100.0) > 1e-6)
.unwrap_or(false)
}
fn admin_pool_format_reset_after(seconds: f64) -> Option<String> {
if !seconds.is_finite() {
return None;
}
let total_seconds = seconds.floor() as i64;
if total_seconds <= 0 {
return Some("已重置".to_string());
}
let days = total_seconds / 86_400;
let hours = (total_seconds % 86_400) / 3_600;
let minutes = (total_seconds % 3_600) / 60;
if days > 0 {
return Some(format!("{days}天{hours}小时后重置"));
}
if hours > 0 {
return Some(format!("{hours}小时{minutes}分钟后重置"));
}
if minutes > 0 {
return Some(format!("{minutes}分钟后重置"));
}
Some("即将重置".to_string())
}
fn admin_pool_build_codex_account_quota(
data: &serde_json::Map<String, serde_json::Value>,
) -> Option<String> {
fn codex_reset_seconds(
data: &serde_json::Map<String, serde_json::Value>,
reset_seconds_key: &str,
reset_after_seconds_key: &str,
reset_at_key: &str,
) -> Option<f64> {
admin_pool_json_to_f64(data.get(reset_seconds_key))
.or_else(|| admin_pool_json_to_f64(data.get(reset_after_seconds_key)))
.or_else(|| {
let reset_at = admin_pool_json_to_u64(data.get(reset_at_key))?;
let now_unix_secs = chrono::Utc::now().timestamp().max(0) as u64;
Some(reset_at.saturating_sub(now_unix_secs) as f64)
})
}
let mut parts = Vec::new();
let primary_used = admin_pool_json_to_f64(data.get("primary_used_percent"));
if let Some(primary_used) = primary_used {
let mut part = format!("周剩余 {}", admin_pool_format_percent(100.0 - primary_used));
if admin_pool_has_quota_consumption(Some(primary_used)) {
if let Some(reset_text) = codex_reset_seconds(
data,
"primary_reset_seconds",
"primary_reset_after_seconds",
"primary_reset_at",
)
.and_then(admin_pool_format_reset_after)
{
part.push_str(&format!(" ({reset_text})"));
}
}
parts.push(part);
}
let secondary_used = admin_pool_json_to_f64(data.get("secondary_used_percent"));
if let Some(secondary_used) = secondary_used {
let mut part = format!(
"5H剩余 {}",
admin_pool_format_percent(100.0 - secondary_used)
);
if admin_pool_has_quota_consumption(Some(secondary_used)) {
if let Some(reset_text) = codex_reset_seconds(
data,
"secondary_reset_seconds",
"secondary_reset_after_seconds",
"secondary_reset_at",
)
.and_then(admin_pool_format_reset_after)
{
part.push_str(&format!(" ({reset_text})"));
}
}
parts.push(part);
}
if !parts.is_empty() {
return Some(parts.join(" | "));
}
let has_credits = data
.get("has_credits")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false);
let credits_balance = admin_pool_json_to_f64(data.get("credits_balance"));
if has_credits && credits_balance.is_some() {
return credits_balance.map(|value| format!("积分 {value:.2}"));
}
if has_credits {
return Some("有积分".to_string());
}
None
}
fn admin_pool_build_kiro_account_quota(
data: &serde_json::Map<String, serde_json::Value>,
) -> Option<String> {
if data
.get("is_banned")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false)
{
return Some("账号已封禁".to_string());
}
let usage_percentage = admin_pool_json_to_f64(data.get("usage_percentage"));
if let Some(usage_percentage) = usage_percentage {
let remaining = 100.0 - usage_percentage;
let current_usage = admin_pool_json_to_f64(data.get("current_usage"));
let usage_limit = admin_pool_json_to_f64(data.get("usage_limit"));
if let (Some(current_usage), Some(usage_limit)) = (current_usage, usage_limit) {
if usage_limit > 0.0 {
return Some(format!(
"剩余 {} ({}/{})",
admin_pool_format_percent(remaining),
admin_pool_format_quota_value(current_usage),
admin_pool_format_quota_value(usage_limit),
));
}
}
return Some(format!("剩余 {}", admin_pool_format_percent(remaining)));
}
let remaining = admin_pool_json_to_f64(data.get("remaining"));
let usage_limit = admin_pool_json_to_f64(data.get("usage_limit"));
match (remaining, usage_limit) {
(Some(remaining), Some(usage_limit)) if usage_limit > 0.0 => Some(format!(
"剩余 {}/{}",
admin_pool_format_quota_value(remaining),
admin_pool_format_quota_value(usage_limit),
)),
_ => None,
}
}
fn admin_pool_quota_by_model(
data: &serde_json::Map<String, serde_json::Value>,
) -> Option<&serde_json::Map<String, serde_json::Value>> {
data.get("quota_by_model")?.as_object()
}
fn admin_pool_build_antigravity_account_quota(
data: &serde_json::Map<String, serde_json::Value>,
) -> Option<String> {
if data
.get("is_forbidden")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false)
{
return Some("访问受限".to_string());
}
let remaining_list = admin_pool_quota_by_model(data)?
.values()
.filter_map(serde_json::Value::as_object)
.filter_map(|item| {
let used_percent = admin_pool_json_to_f64(item.get("used_percent")).or_else(|| {
admin_pool_json_to_f64(item.get("remaining_fraction"))
.map(|value| (1.0 - value) * 100.0)
})?;
Some((100.0 - used_percent).clamp(0.0, 100.0))
})
.collect::<Vec<_>>();
if remaining_list.is_empty() {
return None;
}
let min_remaining = remaining_list.iter().copied().fold(100.0_f64, f64::min);
if remaining_list.len() == 1 {
return Some(format!("剩余 {}", admin_pool_format_percent(min_remaining)));
}
Some(format!(
"最低剩余 {} ({} 模型)",
admin_pool_format_percent(min_remaining),
remaining_list.len()
))
}
fn admin_pool_gemini_reset_at(item: &serde_json::Map<String, serde_json::Value>) -> Option<i64> {
let reset_at = admin_pool_json_to_u64(item.get("reset_at"))?;
Some(reset_at as i64)
}
fn admin_pool_gemini_model_exhausted(item: &serde_json::Map<String, serde_json::Value>) -> bool {
if item
.get("is_exhausted")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false)
{
return true;
}
if admin_pool_json_to_f64(item.get("remaining_fraction")).is_some_and(|value| value <= 0.0) {
return true;
}
admin_pool_json_to_f64(item.get("used_percent")).is_some_and(|value| value >= 100.0 - 1e-6)
}
fn admin_pool_build_gemini_cli_account_quota(
data: &serde_json::Map<String, serde_json::Value>,
) -> Option<String> {
let now = chrono::Utc::now().timestamp();
let mut active = admin_pool_quota_by_model(data)?
.iter()
.filter_map(|(model_name, item)| {
let item = item.as_object()?;
if !admin_pool_gemini_model_exhausted(item) {
return None;
}
let reset_at = admin_pool_gemini_reset_at(item);
if reset_at.is_some_and(|value| value <= now) {
return None;
}
Some((model_name.as_str(), reset_at))
})
.collect::<Vec<_>>();
if active.is_empty() {
return None;
}
active.sort_by_key(|(_, reset_at)| reset_at.unwrap_or(i64::MAX));
let (first_model, first_reset_at) = active[0];
if active.len() == 1 {
if let Some(reset_at) = first_reset_at {
if let Some(reset_text) = admin_pool_format_reset_after((reset_at - now) as f64) {
return Some(format!("{first_model} 冷却中 ({reset_text})"));
}
}
return Some(format!("{first_model} 冷却中"));
}
if let Some(reset_at) = first_reset_at {
if let Some(reset_text) = admin_pool_format_reset_after((reset_at - now) as f64) {
return Some(format!(
"{} 个模型冷却中（最早 {reset_text}）",
active.len()
));
}
}
Some(format!("{} 个模型冷却中", active.len()))
}
fn admin_pool_build_account_quota(
provider_type: &str,
upstream_metadata: Option<&serde_json::Value>,
) -> Option<String> {
let normalized_provider_type = provider_type.trim().to_ascii_lowercase();
let upstream_metadata = upstream_metadata?.as_object()?;
let data = upstream_metadata
.get(&normalized_provider_type)?
.as_object()?;
match normalized_provider_type.as_str() {
"codex" => admin_pool_build_codex_account_quota(data),
"kiro" => admin_pool_build_kiro_account_quota(data),
"antigravity" => admin_pool_build_antigravity_account_quota(data),
"gemini_cli" => admin_pool_build_gemini_cli_account_quota(data),
_ => None,
}
}
fn admin_pool_health_score(key: &StoredProviderCatalogKey) -> f64 {
let scores = key
.health_by_format
.as_ref()
.and_then(serde_json::Value::as_object)
.map(|formats| {
formats
.values()
.filter_map(serde_json::Value::as_object)
.filter_map(|item| item.get("health_score"))
.filter_map(serde_json::Value::as_f64)
.collect::<Vec<_>>()
})
.unwrap_or_default();
if scores.is_empty() {
1.0
} else {
scores.into_iter().fold(1.0, f64::min)
}
}
fn admin_pool_circuit_breaker_open(key: &StoredProviderCatalogKey) -> bool {
key.circuit_breaker_by_format
.as_ref()
.and_then(serde_json::Value::as_object)
.map(|formats| {
formats
.values()
.filter_map(serde_json::Value::as_object)
.any(|item| {
item.get("open")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false)
})
})
.unwrap_or(false)
}
fn admin_pool_scheduling_payload(
key: &StoredProviderCatalogKey,
cooldown_reason: Option<&str>,
cooldown_ttl_seconds: Option<u64>,
health_score: f64,
circuit_breaker_open: bool,
) -> (String, String, String, Vec<serde_json::Value>) {
if !key.is_active {
return (
"blocked".to_string(),
"inactive".to_string(),
"已禁用".to_string(),
vec![json!({
"code": "inactive",
"label": "已禁用",
"blocking": true,
"source": "manual",
"ttl_seconds": serde_json::Value::Null,
"detail": serde_json::Value::Null,
})],
);
}
if let Some(reason) = cooldown_reason {
return (
"degraded".to_string(),
"cooldown".to_string(),
"冷却中".to_string(),
vec![json!({
"code": "cooldown",
"label": "冷却中",
"blocking": true,
"source": "pool",
"ttl_seconds": cooldown_ttl_seconds,
"detail": reason,
})],
);
}
if circuit_breaker_open {
return (
"degraded".to_string(),
"circuit_breaker".to_string(),
"熔断中".to_string(),
vec![json!({
"code": "circuit_breaker",
"label": "熔断中",
"blocking": true,
"source": "health",
"ttl_seconds": serde_json::Value::Null,
"detail": serde_json::Value::Null,
})],
);
}
if health_score < 0.5 {
return (
"degraded".to_string(),
"health_low".to_string(),
"健康度较低".to_string(),
vec![json!({
"code": "health_low",
"label": "健康度较低",
"blocking": false,
"source": "health",
"ttl_seconds": serde_json::Value::Null,
"detail": serde_json::Value::Null,
})],
);
}
(
"available".to_string(),
"available".to_string(),
"可用".to_string(),
Vec::new(),
)
}
pub(super) fn build_admin_pool_key_payload(
state: &AdminAppState<'_>,
provider_type: &str,
key: &StoredProviderCatalogKey,
runtime: &AdminProviderPoolRuntimeState,
pool_config: Option<AdminProviderPoolConfig>,
) -> serde_json::Value {
admin_provider_pool_pure::build_admin_pool_key_payload(
key,
&admin_provider_pool_pure::AdminPoolKeyPayloadContext {
cooldown_reason: runtime.cooldown_reason_by_key.get(&key.id).cloned(),
cooldown_ttl_seconds: runtime
.cooldown_reason_by_key
.get(&key.id)
.and_then(|_| runtime.cooldown_ttl_by_key.get(&key.id).copied()),
cost_window_usage: runtime
.cost_window_usage_by_key
.get(&key.id)
.copied()
.unwrap_or(0),
sticky_sessions: runtime
.sticky_sessions_by_key
.get(&key.id)
.copied()
.unwrap_or(0),
lru_score: runtime.lru_score_by_key.get(&key.id).copied(),
cost_limit: pool_config.and_then(|config| config.cost_limit_per_key_tokens),
},
)
let cooldown_reason = runtime.cooldown_reason_by_key.get(&key.id).cloned();
let cooldown_ttl_seconds = cooldown_reason
.as_ref()
.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 (scheduling_status, scheduling_reason, scheduling_label, scheduling_reasons) =
admin_pool_scheduling_payload(
key,
cooldown_reason.as_deref(),
cooldown_ttl_seconds,
health_score,
circuit_breaker_open,
);
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_plan_type =
admin_pool_derive_oauth_plan_type(key, provider_type, auth_config.as_ref());
let status_snapshot = provider_key_status_snapshot_payload(key);
let account_snapshot = status_snapshot
.get("account")
.and_then(serde_json::Value::as_object);
let quota_snapshot = status_snapshot
.get("quota")
.and_then(serde_json::Value::as_object);
let oauth_snapshot = status_snapshot
.get("oauth")
.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 =
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());
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")));
let account_status_reason =
admin_pool_trimmed_string(account_snapshot.and_then(|item| item.get("reason")));
let account_status_blocked = account_snapshot
.and_then(|item| item.get("blocked"))
.and_then(serde_json::Value::as_bool)
.unwrap_or(false);
let account_status_recoverable = account_snapshot
.and_then(|item| item.get("recoverable"))
.and_then(serde_json::Value::as_bool)
.unwrap_or(false);
let account_status_source =
admin_pool_trimmed_string(account_snapshot.and_then(|item| item.get("source")));
let mut payload = serde_json::Map::new();
payload.insert("key_id".to_string(), json!(key.id));
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("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),
);
payload.insert("oauth_plan_type".to_string(), json!(oauth_plan_type));
payload.insert("oauth_account_id".to_string(), json!(oauth_account_id));
payload.insert("oauth_account_name".to_string(), json!(oauth_account_name));
payload.insert(
"oauth_account_user_id".to_string(),
json!(oauth_account_user_id),
);
payload.insert(
"oauth_organizations".to_string(),
serde_json::Value::Array(oauth_organizations),
);
payload.insert(
"account_status_code".to_string(),
json!(account_status_code),
);
payload.insert(
"account_status_label".to_string(),
json!(account_status_label),
);
payload.insert(
"account_status_reason".to_string(),
json!(account_status_reason),
);
payload.insert(
"account_status_blocked".to_string(),
json!(account_status_blocked),
);
payload.insert(
"account_status_recoverable".to_string(),
json!(account_status_recoverable),
);
payload.insert(
"account_status_source".to_string(),
json!(account_status_source),
);
payload.insert("status_snapshot".to_string(), status_snapshot);
payload.insert("quota_updated_at".to_string(), json!(quota_updated_at));
payload.insert("health_score".to_string(), json!(health_score));
payload.insert(
"circuit_breaker_open".to_string(),
json!(circuit_breaker_open),
);
payload.insert(
"api_formats".to_string(),
json!(admin_pool_api_formats(key)),
);
payload.insert(
"rate_multipliers".to_string(),
json!(admin_pool_json_object(key.rate_multipliers.as_ref())),
);
payload.insert(
"internal_priority".to_string(),
json!(key.internal_priority),
);
payload.insert("rpm_limit".to_string(), json!(key.rpm_limit));
payload.insert(
"cache_ttl_minutes".to_string(),
json!(key.cache_ttl_minutes),
);
payload.insert(
"max_probe_interval_minutes".to_string(),
json!(key.max_probe_interval_minutes),
);
payload.insert("note".to_string(), json!(key.note));
payload.insert(
"allowed_models".to_string(),
json!(admin_pool_string_list(key.allowed_models.as_ref())),
);
payload.insert(
"capabilities".to_string(),
json!(admin_pool_json_object(key.capabilities.as_ref())),
);
payload.insert(
"auto_fetch_models".to_string(),
json!(key.auto_fetch_models),
);
payload.insert(
"locked_models".to_string(),
json!(admin_pool_string_list(key.locked_models.as_ref())),
);
payload.insert(
"model_include_patterns".to_string(),
json!(admin_pool_string_list(key.model_include_patterns.as_ref())),
);
payload.insert(
"model_exclude_patterns".to_string(),
json!(admin_pool_string_list(key.model_exclude_patterns.as_ref())),
);
payload.insert("proxy".to_string(), json!(key.proxy.clone()));
payload.insert("fingerprint".to_string(), json!(key.fingerprint.clone()));
payload.insert(
"account_quota".to_string(),
json!(admin_pool_build_account_quota(
provider_type,
key.upstream_metadata.as_ref(),
)),
);
payload.insert("cooldown_reason".to_string(), json!(cooldown_reason));
payload.insert(
"cooldown_ttl_seconds".to_string(),
json!(cooldown_ttl_seconds),
);
payload.insert(
"cost_window_usage".to_string(),
json!(runtime
.cost_window_usage_by_key
.get(&key.id)
.copied()
.unwrap_or(0)),
);
payload.insert(
"cost_limit".to_string(),
json!(pool_config.map(|config| config.cost_limit_per_key_tokens)),
);
payload.insert(
"request_count".to_string(),
json!(key.request_count.unwrap_or(0)),
);
payload.insert("total_tokens".to_string(), json!(key.total_tokens));
payload.insert(
"total_cost_usd".to_string(),
json!(format!("{:.8}", key.total_cost_usd)),
);
payload.insert(
"sticky_sessions".to_string(),
json!(runtime
.sticky_sessions_by_key
.get(&key.id)
.copied()
.unwrap_or(0)),
);
payload.insert(
"lru_score".to_string(),
json!(runtime.lru_score_by_key.get(&key.id).copied()),
);
payload.insert(
"created_at".to_string(),
json!(key.created_at_unix_secs.and_then(unix_secs_to_rfc3339)),
);
payload.insert(
"last_used_at".to_string(),
json!(key.last_used_at_unix_secs.and_then(unix_secs_to_rfc3339)),
);
payload.insert("scheduling_status".to_string(), json!(scheduling_status));
payload.insert("scheduling_reason".to_string(), json!(scheduling_reason));
payload.insert("scheduling_label".to_string(), json!(scheduling_label));
payload.insert("scheduling_reasons".to_string(), json!(scheduling_reasons));
serde_json::Value::Object(payload)
}

10
apps/aether-gateway/src/handlers/admin/provider/pool_admin/read_routes/keys.rs
View File

@@ -133,7 +133,15 @@ pub(super) async fn build_admin_pool_list_keys_response(
let items = keys
.into_iter()
.map(|key| pool_payloads::build_admin_pool_key_payload(&key, &runtime, pool_config))
.map(|key| {
pool_payloads::build_admin_pool_key_payload(
state,
&provider.provider_type,
&key,
&runtime,
pool_config,
)
})
.collect::<Vec<_>>();
Ok(Json(json!({

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

@@ -48,6 +48,33 @@ fn admin_pool_derive_oauth_plan_type(
return None;
}
if let Some(upstream_metadata) = key
.upstream_metadata
.as_ref()
.and_then(serde_json::Value::as_object)
{
let provider_bucket = upstream_metadata
.get(&provider_type.trim().to_ascii_lowercase())
.and_then(serde_json::Value::as_object);
for source in provider_bucket
.into_iter()
.chain(std::iter::once(upstream_metadata))
{
for plan_key in [
"plan_type",
"tier",
"subscription_title",
"subscription_plan",
] {
if let Some(value) = source.get(plan_key).and_then(serde_json::Value::as_str) {
if let Some(normalized) = normalize(value) {
return Some(normalized);
}
}
}
}
}
if let Some(auth_config) = admin_pool_parse_auth_config_json(state, key) {
for plan_key in ["plan_type", "tier", "plan", "subscription_plan"] {
if let Some(value) = auth_config
@@ -61,28 +88,6 @@ fn admin_pool_derive_oauth_plan_type(
}
}
let upstream_metadata = key.upstream_metadata.as_ref()?.as_object()?;
let provider_bucket = upstream_metadata
.get(&provider_type.trim().to_ascii_lowercase())
.and_then(serde_json::Value::as_object);
for source in provider_bucket
.into_iter()
.chain(std::iter::once(upstream_metadata))
{
for plan_key in [
"plan_type",
"tier",
"subscription_title",
"subscription_plan",
] {
if let Some(value) = source.get(plan_key).and_then(serde_json::Value::as_str) {
if let Some(normalized) = normalize(value) {
return Some(normalized);
}
}
}
}
None
}

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

@@ -1,71 +1,223 @@
use super::payload::{provider_query_extract_api_key_id, provider_query_extract_provider_id};
use super::payload::{
provider_query_extract_api_key_id, provider_query_extract_force_refresh,
provider_query_extract_provider_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_NO_LOCAL_MODELS_DETAIL, ADMIN_PROVIDER_QUERY_PROVIDER_ID_REQUIRED_DETAIL,
ADMIN_PROVIDER_QUERY_PROVIDER_ID_REQUIRED_DETAIL,
ADMIN_PROVIDER_QUERY_PROVIDER_NOT_FOUND_DETAIL,
};
use crate::execution_runtime;
use crate::model_fetch::ModelFetchRuntimeState;
use crate::handlers::admin::request::AdminAppState;
use crate::GatewayError;
use crate::{AppState, GatewayError};
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use aether_model_fetch::{
aggregate_models_for_cache, build_models_fetch_execution_plan, extract_error_message,
parse_models_response,
};
use axum::{body::Body, http::Response, response::IntoResponse, Json};
use serde_json::json;
use std::collections::{BTreeMap, BTreeSet};
use serde_json::{json, Value};
use std::collections::BTreeMap;
pub(crate) const ADMIN_PROVIDER_QUERY_LOCAL_TEST_MODEL_MESSAGE: &str =
"Rust local provider-query model test is not configured";
pub(crate) const ADMIN_PROVIDER_QUERY_LOCAL_TEST_MODEL_FAILOVER_MESSAGE: &str =
"Rust local provider-query failover simulation is not configured";
const ADMIN_PROVIDER_QUERY_NO_ACTIVE_ENDPOINT_DETAIL: &str =
"No active endpoints found for this provider";
const ADMIN_PROVIDER_QUERY_NO_MODELS_FROM_ENDPOINT_DETAIL: &str =
"No models returned from any endpoint";
const PROVIDER_QUERY_FETCH_FORMAT_PRIORITY: &[&[&str]] = &[
&["openai:chat", "openai:cli", "openai:compact"],
&["claude:chat", "claude:cli"],
&["gemini:chat", "gemini:cli"],
];
fn provider_query_string_list(value: Option<&serde_json::Value>) -> Vec<String> {
value
.and_then(serde_json::Value::as_array)
.map(|items| {
items
.iter()
.filter_map(serde_json::Value::as_str)
.map(str::trim)
.filter(|item| !item.is_empty())
.map(ToOwned::to_owned)
.collect::<Vec<_>>()
})
.unwrap_or_default()
#[derive(Debug)]
struct ProviderQueryKeyFetchResult {
models: Vec<Value>,
error: Option<String>,
from_cache: bool,
}
fn provider_query_resolved_api_formats(
endpoints: &[StoredProviderCatalogEndpoint],
selected_key: Option<&StoredProviderCatalogKey>,
) -> Vec<String> {
let mut seen = BTreeSet::new();
let key_formats = selected_key
.map(|key| provider_query_string_list(key.api_formats.as_ref()))
.unwrap_or_default();
let mut formats = Vec::new();
fn provider_query_provider_payload(provider: &StoredProviderCatalogProvider) -> Value {
json!({
"id": provider.id.clone(),
"name": provider.name.clone(),
"display_name": provider.name.clone(),
})
}
fn provider_query_key_display_name(key: &StoredProviderCatalogKey) -> String {
let trimmed = key.name.trim();
if trimmed.is_empty() {
key.id.clone()
} else {
trimmed.to_string()
}
}
fn provider_query_normalize_api_format(value: &str) -> String {
value.trim().to_ascii_lowercase()
}
fn provider_query_selected_fetch_endpoints(
endpoints: &[StoredProviderCatalogEndpoint],
) -> Vec<StoredProviderCatalogEndpoint> {
let mut by_format = BTreeMap::<String, StoredProviderCatalogEndpoint>::new();
for endpoint in endpoints.iter().filter(|endpoint| endpoint.is_active) {
let api_format = endpoint.api_format.trim();
let api_format = provider_query_normalize_api_format(&endpoint.api_format);
if api_format.is_empty() {
continue;
}
if !key_formats.is_empty() && !key_formats.iter().any(|value| value == api_format) {
by_format.insert(api_format, endpoint.clone());
}
// 与 Python 版本保持一致：同族优先使用 chat 端点，其次才回退到 cli/compact。
PROVIDER_QUERY_FETCH_FORMAT_PRIORITY
.iter()
.filter_map(|candidates| {
candidates
.iter()
.find_map(|api_format| by_format.get(*api_format))
.cloned()
})
.collect()
}
async fn provider_query_read_cached_models(
state: &AppState,
provider_id: &str,
key_id: &str,
) -> Option<Vec<Value>> {
let runner = state.redis_kv_runner()?;
let cache_key = runner
.keyspace()
.key(&format!("upstream_models:{provider_id}:{key_id}"));
let mut connection = runner
.client()
.get_multiplexed_async_connection()
.await
.ok()?;
let raw = redis::cmd("GET")
.arg(&cache_key)
.query_async::<Option<String>>(&mut connection)
.await
.ok()??;
let parsed = serde_json::from_str::<Vec<Value>>(&raw).ok()?;
Some(aggregate_models_for_cache(&parsed))
}
async fn provider_query_fetch_models_from_transport(
state: &AppState,
transport: &crate::provider_transport::GatewayProviderTransportSnapshot,
) -> Result<Vec<Value>, String> {
let plan = build_models_fetch_execution_plan(state, transport).await?;
let result = execution_runtime::execute_execution_runtime_sync_plan(state, None, &plan)
.await
.map_err(|err| format!("{err:?}"))?;
if result.status_code != 200 {
let message = result
.body
.as_ref()
.and_then(|body| body.json_body.as_ref())
.and_then(extract_error_message)
.or_else(|| {
result.error.as_ref().and_then(|error| {
let message = error.message.trim();
(!message.is_empty()).then_some(message.to_string())
})
})
.unwrap_or_else(|| format!("upstream returned status {}", result.status_code));
return Err(message);
}
let body_json = result
.body
.as_ref()
.and_then(|body| body.json_body.as_ref())
.ok_or_else(|| "models fetch response body is missing JSON payload".to_string())?;
let parsed = parse_models_response(&transport.endpoint.api_format, body_json)?;
Ok(parsed.cached_models)
}
async fn provider_query_fetch_models_for_key(
state: &AppState,
provider: &StoredProviderCatalogProvider,
endpoints: &[StoredProviderCatalogEndpoint],
key: &StoredProviderCatalogKey,
force_refresh: bool,
) -> Result<ProviderQueryKeyFetchResult, GatewayError> {
if !force_refresh {
if let Some(cached_models) =
provider_query_read_cached_models(state, &provider.id, &key.id).await
{
return Ok(ProviderQueryKeyFetchResult {
models: cached_models,
error: None,
from_cache: true,
});
}
}
let selected_endpoints = provider_query_selected_fetch_endpoints(endpoints);
if selected_endpoints.is_empty() {
return Ok(ProviderQueryKeyFetchResult {
models: Vec::new(),
error: Some(ADMIN_PROVIDER_QUERY_NO_ACTIVE_ENDPOINT_DETAIL.to_string()),
from_cache: false,
});
}
let mut all_models = Vec::new();
let mut all_errors = Vec::new();
for endpoint in selected_endpoints {
let Some(transport) = state
.read_provider_transport_snapshot(&provider.id, &endpoint.id, &key.id)
.await?
else {
all_errors.push(format!(
"{} transport snapshot unavailable",
endpoint.api_format.trim()
));
continue;
}
if seen.insert(api_format.to_string()) {
formats.push(api_format.to_string());
};
match provider_query_fetch_models_from_transport(state, &transport).await {
Ok(models) => all_models.extend(models),
Err(err) => all_errors.push(err),
}
}
if formats.is_empty() {
for api_format in key_formats {
if seen.insert(api_format.clone()) {
formats.push(api_format);
}
}
let unique_models = aggregate_models_for_cache(&all_models);
if !unique_models.is_empty() {
<AppState as ModelFetchRuntimeState>::write_upstream_models_cache(
state,
&provider.id,
&key.id,
&unique_models,
)
.await;
}
formats
let mut error = if all_errors.is_empty() {
None
} else {
Some(all_errors.join("; "))
};
if unique_models.is_empty() && error.is_none() {
error = Some(ADMIN_PROVIDER_QUERY_NO_MODELS_FROM_ENDPOINT_DETAIL.to_string());
}
Ok(ProviderQueryKeyFetchResult {
models: unique_models,
error,
from_cache: false,
})
}
pub(crate) async fn build_admin_provider_query_models_response(
@@ -99,96 +251,89 @@ pub(crate) async fn build_admin_provider_query_models_response(
.app()
.list_provider_catalog_keys_by_provider_ids(&provider_ids)
.await?;
let selected_key = if let Some(api_key_id) = provider_query_extract_api_key_id(payload) {
let Some(key) = keys.iter().find(|key| key.id == api_key_id) else {
let force_refresh = provider_query_extract_force_refresh(payload);
if let Some(api_key_id) = provider_query_extract_api_key_id(payload) {
let Some(selected_key) = keys.iter().find(|key| key.id == api_key_id) else {
return Ok(build_admin_provider_query_not_found_response(
ADMIN_PROVIDER_QUERY_API_KEY_NOT_FOUND_DETAIL,
));
};
Some(key)
} else {
None
};
let active_keys = keys.iter().filter(|key| key.is_active).count();
if selected_key.is_none() && active_keys == 0 {
let result = provider_query_fetch_models_for_key(
state.app(),
&provider,
&endpoints,
selected_key,
force_refresh,
)
.await?;
let success = !result.models.is_empty();
return Ok(Json(json!({
"success": success,
"data": {
"models": result.models,
"error": result.error,
"from_cache": result.from_cache,
},
"provider": provider_query_provider_payload(&provider),
}))
.into_response());
}
let active_keys = keys.iter().filter(|key| key.is_active).collect::<Vec<_>>();
if active_keys.is_empty() {
return Ok(build_admin_provider_query_bad_request_response(
ADMIN_PROVIDER_QUERY_NO_ACTIVE_API_KEY_DETAIL,
));
}
let active_key_count = active_keys.len();
let resolved_api_formats = provider_query_resolved_api_formats(&endpoints, selected_key);
let provider_models = state
.app()
.list_admin_provider_available_source_models(&provider.id)
.await?;
let mut grouped: BTreeMap<
String,
(
aether_data_contracts::repository::global_models::StoredAdminProviderModel,
BTreeSet<String>,
),
> = BTreeMap::new();
for model in provider_models {
let entry = grouped
.entry(model.provider_model_name.clone())
.or_insert_with(|| (model.clone(), BTreeSet::new()));
for api_format in &resolved_api_formats {
entry.1.insert(api_format.clone());
let mut all_models = Vec::new();
let mut all_errors = Vec::new();
let mut cache_hit_count = 0usize;
let mut fetch_count = 0usize;
for key in active_keys {
let result =
provider_query_fetch_models_for_key(state.app(), &provider, &endpoints, key, force_refresh)
.await?;
all_models.extend(result.models);
if let Some(error) = result.error {
all_errors.push(format!(
"Key {}: {}",
provider_query_key_display_name(key),
error
));
}
if result.from_cache {
cache_hit_count += 1;
} else {
fetch_count += 1;
}
}
let models: Vec<_> = grouped
.into_iter()
.map(|(model_id, (model, api_formats))| {
let display_name = model
.global_model_display_name
.clone()
.or(model.global_model_name.clone())
.unwrap_or_else(|| model_id.clone());
let api_formats: Vec<_> = api_formats.into_iter().collect();
json!({
"id": model_id,
"object": "model",
"created": model.created_at_unix_secs,
"owned_by": provider.name,
"display_name": display_name,
"api_format": api_formats.first().cloned(),
"api_formats": api_formats,
"provider_model_name": model.provider_model_name,
"global_model_id": model.global_model_id,
"global_model_name": model.global_model_name,
"supports_streaming": model.supports_streaming,
"supports_function_calling": model.supports_function_calling,
"supports_vision": model.supports_vision,
"supports_extended_thinking": model.supports_extended_thinking,
"supports_image_generation": model.supports_image_generation,
"is_available": model.is_available,
})
})
.collect();
let models = aggregate_models_for_cache(&all_models);
let success = !models.is_empty();
let error = if success {
let mut error = if all_errors.is_empty() {
None
} else {
Some(ADMIN_PROVIDER_QUERY_NO_LOCAL_MODELS_DETAIL)
Some(all_errors.join("; "))
};
if !success && error.is_none() {
error = Some("No models returned from any key".to_string());
}
Ok(Json(json!({
"success": success,
"data": {
"models": models,
"error": error,
"from_cache": true,
"keys_total": active_keys,
"keys_cached": 0,
"keys_fetched": 0,
},
"provider": {
"id": provider.id,
"name": provider.name,
"display_name": provider.name,
"from_cache": fetch_count == 0 && cache_hit_count > 0,
"keys_total": active_key_count,
"keys_cached": cache_hit_count,
"keys_fetched": fetch_count,
},
"provider": provider_query_provider_payload(&provider),
}))
.into_response())
}

7
apps/aether-gateway/src/handlers/admin/provider/query/payload.rs
View File

@@ -36,6 +36,13 @@ pub(crate) fn provider_query_extract_api_key_id(payload: &serde_json::Value) ->
.map(ToOwned::to_owned)
}
pub(crate) fn provider_query_extract_force_refresh(payload: &serde_json::Value) -> bool {
payload
.get("force_refresh")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false)
}
pub(crate) fn provider_query_extract_model(payload: &serde_json::Value) -> Option<String> {
payload
.get("model")

4
apps/aether-gateway/src/handlers/admin/provider/write/keys/payload.rs
View File

@@ -39,7 +39,9 @@ pub(crate) async fn build_admin_provider_keys_payload(
keys.into_iter()
.skip(skip)
.take(limit)
.map(|key| state.build_admin_provider_key_response(&key, now_unix_secs))
.map(|key| {
state.build_admin_provider_key_response(&key, &provider.provider_type, now_unix_secs)
})
.collect(),
))
}

2
apps/aether-gateway/src/handlers/admin/request/provider/builders.rs
View File

@@ -40,11 +40,13 @@ impl<'a> AdminAppState<'a> {
pub(crate) fn build_admin_provider_key_response(
&self,
key: &aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey,
provider_type: &str,
now_unix_secs: u64,
) -> serde_json::Value {
crate::handlers::admin::shared::build_admin_provider_key_response(
self.app,
key,
provider_type,
now_unix_secs,
)
}

3
apps/aether-gateway/src/handlers/public/support.rs
View File

@@ -672,8 +672,7 @@ pub(crate) async fn maybe_build_local_public_support_response(
providers.into_iter().next()
} else {
None
}
.or_else(|| None);
};
let provider = match provider {
Some(provider) => provider,
None => state

88
apps/aether-gateway/src/handlers/public/support/models/route.rs
View File

@@ -1,5 +1,6 @@
use std::time::{SystemTime, UNIX_EPOCH};
use aether_data_contracts::repository::candidate_selection::StoredMinimalCandidateSelectionRow;
use axum::{body::Body, response::Response};
use super::models_responses::{
@@ -9,9 +10,74 @@ use super::models_responses::{
build_models_not_found_response, build_openai_model_detail_response,
build_openai_models_list_response,
};
use super::models_shared::{filter_rows_for_models, models_api_format, models_detail_id};
use super::models_shared::{
filter_rows_for_models, models_api_format, models_detail_id, models_query_api_formats,
};
use super::{query_param_value, AppState, GatewayPublicRequestContext};
fn sort_and_dedup_model_rows(
mut rows: Vec<StoredMinimalCandidateSelectionRow>,
) -> Vec<StoredMinimalCandidateSelectionRow> {
rows.sort_by(|left, right| {
left.global_model_name
.cmp(&right.global_model_name)
.then(left.provider_priority.cmp(&right.provider_priority))
.then(left.key_internal_priority.cmp(&right.key_internal_priority))
.then(left.provider_id.cmp(&right.provider_id))
.then(left.endpoint_id.cmp(&right.endpoint_id))
.then(left.key_id.cmp(&right.key_id))
.then(left.model_id.cmp(&right.model_id))
});
let mut deduped = Vec::with_capacity(rows.len());
let mut last_model_name: Option<String> = None;
for row in rows {
if last_model_name.as_deref() == Some(row.global_model_name.as_str()) {
continue;
}
last_model_name = Some(row.global_model_name.clone());
deduped.push(row);
}
deduped
}
async fn list_model_rows_for_client_format(
state: &AppState,
api_format: &str,
auth_snapshot: Option<&crate::data::auth::GatewayAuthApiKeySnapshot>,
) -> Option<Vec<StoredMinimalCandidateSelectionRow>> {
let mut collected = Vec::new();
for query_format in models_query_api_formats(api_format) {
let rows = state
.list_minimal_candidate_selection_rows_for_api_format(query_format)
.await
.ok()?;
let mut filtered = filter_rows_for_models(rows, auth_snapshot, query_format);
collected.append(&mut filtered);
}
Some(sort_and_dedup_model_rows(collected))
}
async fn list_model_rows_for_client_format_and_global_model(
state: &AppState,
api_format: &str,
global_model_name: &str,
auth_snapshot: Option<&crate::data::auth::GatewayAuthApiKeySnapshot>,
) -> Option<Vec<StoredMinimalCandidateSelectionRow>> {
let mut collected = Vec::new();
for query_format in models_query_api_formats(api_format) {
let rows = state
.list_minimal_candidate_selection_rows_for_api_format_and_global_model(
query_format,
global_model_name,
)
.await
.ok()?;
let mut filtered = filter_rows_for_models(rows, auth_snapshot, query_format);
collected.append(&mut filtered);
}
Some(sort_and_dedup_model_rows(collected))
}
pub(super) async fn maybe_build_local_models_route_response(
state: &AppState,
request_context: &GatewayPublicRequestContext,
@@ -44,11 +110,7 @@ pub(super) async fn maybe_build_local_models_route_response(
match decision.route_kind.as_deref() {
Some("list") => {
let rows = state
.list_minimal_candidate_selection_rows_for_api_format(api_format)
.await
.ok()?;
let rows = filter_rows_for_models(rows, auth_snapshot, api_format);
let rows = list_model_rows_for_client_format(state, api_format, auth_snapshot).await?;
if rows.is_empty() {
return Some(build_empty_models_list_response(api_format));
}
@@ -94,13 +156,13 @@ pub(super) async fn maybe_build_local_models_route_response(
}
Some("detail") => {
let model_id = models_detail_id(&request_context.request_path)?;
let rows = state
.list_minimal_candidate_selection_rows_for_api_format_and_global_model(
api_format, &model_id,
)
.await
.ok()?;
let rows = filter_rows_for_models(rows, auth_snapshot, api_format);
let rows = list_model_rows_for_client_format_and_global_model(
state,
api_format,
&model_id,
auth_snapshot,
)
.await?;
let Some(row) = rows.first() else {
return Some(build_models_not_found_response(&model_id, api_format));
};

17
apps/aether-gateway/src/handlers/public/support/models/shared.rs
View File

@@ -13,6 +13,23 @@ pub(crate) fn models_api_format(request_context: &GatewayPublicRequestContext) -
.filter(|signature| matches!(*signature, "openai:chat" | "claude:chat" | "gemini:chat"))
}
const MODELS_CROSS_FORMAT_QUERY_API_FORMATS: &[&str] = &[
"openai:chat",
"openai:cli",
"openai:compact",
"claude:chat",
"claude:cli",
"gemini:chat",
"gemini:cli",
];
pub(super) fn models_query_api_formats(api_format: &str) -> &'static [&'static str] {
match api_format.trim().to_ascii_lowercase().as_str() {
"openai:chat" | "claude:chat" | "gemini:chat" => MODELS_CROSS_FORMAT_QUERY_API_FORMATS,
_ => &[],
}
}
pub(super) fn models_detail_id(request_path: &str) -> Option<String> {
let raw = if let Some(value) = request_path.strip_prefix("/v1/models/") {
value

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

@@ -231,9 +231,98 @@ pub(crate) fn provider_key_health_summary(
)
}
fn normalize_catalog_oauth_plan_type(value: &str, provider_type: &str) -> Option<String> {
let mut normalized = value.trim().to_string();
if normalized.is_empty() {
return None;
}
let provider_type = provider_type.trim().to_ascii_lowercase();
if !provider_type.is_empty() && normalized.to_ascii_lowercase().starts_with(&provider_type) {
normalized = normalized[provider_type.len()..]
.trim_matches(|ch: char| [' ', ':', '-', '_'].contains(&ch))
.to_string();
}
let normalized = normalized.trim().to_ascii_lowercase();
if normalized.is_empty() {
None
} else {
Some(normalized)
}
}
fn catalog_oauth_plan_type_from_source(
source: &serde_json::Map<String, serde_json::Value>,
provider_type: &str,
fields: &[&str],
) -> Option<String> {
for field in fields {
let Some(value) = source.get(*field).and_then(serde_json::Value::as_str) else {
continue;
};
if let Some(normalized) = normalize_catalog_oauth_plan_type(value, provider_type) {
return Some(normalized);
}
}
None
}
fn derive_catalog_oauth_plan_type(
key: &StoredProviderCatalogKey,
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") {
return None;
}
let provider_type_key = provider_type.trim().to_ascii_lowercase();
if let Some(upstream_metadata) = key
.upstream_metadata
.as_ref()
.and_then(serde_json::Value::as_object)
{
let provider_bucket = if provider_type_key.is_empty() {
None
} else {
upstream_metadata
.get(&provider_type_key)
.and_then(serde_json::Value::as_object)
};
for source in provider_bucket
.into_iter()
.chain(std::iter::once(upstream_metadata))
{
if let Some(plan_type) = catalog_oauth_plan_type_from_source(
source,
provider_type,
&[
"plan_type",
"tier",
"subscription_title",
"subscription_plan",
"plan",
],
) {
return Some(plan_type);
}
}
}
auth_config.and_then(|source| {
catalog_oauth_plan_type_from_source(
source,
provider_type,
&["plan_type", "tier", "plan", "subscription_plan"],
)
})
}
pub(crate) fn build_admin_provider_key_response(
state: &AppState,
key: &StoredProviderCatalogKey,
provider_type: &str,
now_unix_secs: u64,
) -> serde_json::Value {
let request_count = u64::from(key.request_count.unwrap_or(0));
@@ -257,16 +346,7 @@ pub(crate) fn build_admin_provider_key_response(
.and_then(serde_json::Value::as_array)
.cloned()
.unwrap_or_default();
let oauth_plan_type = auth_config
.as_ref()
.and_then(|config| config.get("plan_type").and_then(serde_json::Value::as_str))
.map(ToOwned::to_owned)
.or_else(|| {
auth_config
.as_ref()
.and_then(|config| config.get("tier").and_then(serde_json::Value::as_str))
.map(|value| value.to_ascii_lowercase())
});
let oauth_plan_type = derive_catalog_oauth_plan_type(key, provider_type, auth_config.as_ref());
let (
health_score,
consecutive_failures,

2
apps/aether-gateway/src/lib.rs
View File

@@ -79,7 +79,7 @@ pub(crate) use self::fallback_metrics::{GatewayFallbackMetricKind, GatewayFallba
pub use self::middleware::strip_cf_headers_middleware;
pub use self::rate_limit::FrontdoorUserRpmConfig;
pub(crate) use self::rate_limit::FrontdoorUserRpmOutcome;
pub use self::router::{build_router, build_router_with_state, serve_tcp};
pub use self::router::{attach_static_frontend, build_router, build_router_with_state, serve_tcp};
pub(crate) use self::state::{
AdminBillingCollectorRecord, AdminBillingCollectorWriteInput, AdminBillingRuleRecord,
AdminBillingRuleWriteInput, AdminWalletMutationOutcome, AdminWalletPaymentOrderRecord,

146
apps/aether-gateway/src/main.rs
View File

@@ -5,8 +5,8 @@ use aether_crypto::warm_python_fernet_secret;
use aether_data::postgres::PostgresPoolConfig;
use aether_data::redis::RedisClientConfig;
use aether_gateway::{
build_router_with_state, AppState, FrontdoorCorsConfig, FrontdoorUserRpmConfig,
GatewayDataConfig, UsageRuntimeConfig, VideoTaskTruthSourceMode,
attach_static_frontend, build_router_with_state, AppState, FrontdoorCorsConfig,
FrontdoorUserRpmConfig, GatewayDataConfig, UsageRuntimeConfig, VideoTaskTruthSourceMode,
};
use aether_runtime::{
init_service_runtime, DistributedConcurrencyGate, FileLoggingConfig, LogDestination, LogFormat,
@@ -494,6 +494,18 @@ struct Args {
#[arg(long, env = "AETHER_GATEWAY_BIND", default_value = "0.0.0.0:80")]
bind: String,
/// 容器内健康检查入口：根据当前 bind 端口探测本地 /health。
#[arg(long, hide = true, default_value_t = false)]
healthcheck: bool,
#[arg(
long,
hide = true,
env = "AETHER_GATEWAY_HEALTHCHECK_TIMEOUT_MS",
default_value_t = 3_000
)]
healthcheck_timeout_ms: u64,
#[arg(
long,
env = "AETHER_GATEWAY_DEPLOYMENT_TOPOLOGY",
@@ -609,6 +621,70 @@ fn resolve_gateway_log_instance_id() -> String {
.unwrap_or_else(|| "local".to_string())
}
fn resolve_healthcheck_url(bind: &str) -> Result<String, std::io::Error> {
let trimmed = bind.trim();
if trimmed.is_empty() {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"AETHER_GATEWAY_BIND cannot be empty when --healthcheck is enabled",
));
}
if let Ok(socket_addr) = trimmed.parse::<std::net::SocketAddr>() {
let host = match socket_addr.ip() {
std::net::IpAddr::V4(ip) if ip.is_unspecified() => "127.0.0.1".to_string(),
std::net::IpAddr::V4(ip) => ip.to_string(),
std::net::IpAddr::V6(ip) if ip.is_unspecified() => "[::1]".to_string(),
std::net::IpAddr::V6(ip) => format!("[{ip}]"),
};
return Ok(format!("http://{host}:{}/health", socket_addr.port()));
}
let (host, port) = trimmed.rsplit_once(':').ok_or_else(|| {
std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!(
"AETHER_GATEWAY_BIND must include a port when --healthcheck is enabled: {trimmed}"
),
)
})?;
let port = port.parse::<u16>().map_err(|error| {
std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!("invalid healthcheck port in AETHER_GATEWAY_BIND={trimmed}: {error}"),
)
})?;
let host = host.trim();
if host.is_empty() {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!("invalid host in AETHER_GATEWAY_BIND={trimmed}"),
));
}
let host = if host.contains(':') && !host.starts_with('[') {
format!("[{host}]")
} else {
host.to_string()
};
Ok(format!("http://{host}:{port}/health"))
}
async fn run_healthcheck(args: &Args) -> Result<(), Box<dyn std::error::Error>> {
let url = resolve_healthcheck_url(&args.bind)?;
reqwest::Client::builder()
.timeout(std::time::Duration::from_millis(
args.healthcheck_timeout_ms.max(1),
))
.build()?
.get(url)
.send()
.await?
.error_for_status()?;
Ok(())
}
fn validate_deployment_topology(
args: &Args,
data_postgres_url: Option<&str>,
@@ -685,6 +761,9 @@ fn validate_deployment_topology(
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
let args = Args::parse();
if args.healthcheck {
return run_healthcheck(&args).await;
}
init_service_runtime(args.runtime_config()?)?;
let data_postgres_url = args.data.effective_postgres_url();
let data_redis_url = args.data.effective_redis_url();
@@ -776,7 +855,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.or_else(|| data_redis_url.as_deref())
.or(data_redis_url.as_deref())
.ok_or_else(|| {
std::io::Error::new(
std::io::ErrorKind::InvalidInput,
@@ -847,17 +926,9 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
// Compose the final router: API routes + optional static file serving + CF header stripping
let router = if let Some(ref static_dir) = args.static_dir {
use tower_http::compression::CompressionLayer;
use tower_http::services::{ServeDir, ServeFile};
let static_path = std::path::PathBuf::from(static_dir);
let index_html = static_path.join("index.html");
info!(static_dir = %static_dir, "serving frontend static files");
// ServeDir with SPA fallback: if no static file matches, serve index.html
let serve_dir = ServeDir::new(&static_path).not_found_service(ServeFile::new(&index_html));
api_router
.fallback_service(serve_dir)
attach_static_frontend(api_router, static_dir)
.layer(CompressionLayer::new())
.layer(axum::middleware::from_fn(
aether_gateway::strip_cf_headers_middleware,
@@ -878,3 +949,54 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::resolve_healthcheck_url;
#[test]
fn resolves_ipv4_healthcheck_url() {
assert_eq!(
resolve_healthcheck_url("0.0.0.0:80").unwrap(),
"http://127.0.0.1:80/health"
);
}
#[test]
fn resolves_ipv6_healthcheck_url() {
assert_eq!(
resolve_healthcheck_url("[::]:8080").unwrap(),
"http://[::1]:8080/health"
);
}
#[test]
fn preserves_explicit_ipv4_bind_for_healthcheck_url() {
assert_eq!(
resolve_healthcheck_url("172.18.0.2:9000").unwrap(),
"http://172.18.0.2:9000/health"
);
}
#[test]
fn preserves_explicit_ipv6_bind_for_healthcheck_url() {
assert_eq!(
resolve_healthcheck_url("[2001:db8::2]:9000").unwrap(),
"http://[2001:db8::2]:9000/health"
);
}
#[test]
fn preserves_hostname_healthcheck_url() {
assert_eq!(
resolve_healthcheck_url("gateway.internal:9000").unwrap(),
"http://gateway.internal:9000/health"
);
}
#[test]
fn rejects_bind_without_port() {
let error = resolve_healthcheck_url("not-a-socket").unwrap_err();
assert_eq!(error.kind(), std::io::ErrorKind::InvalidInput);
}
}

3
apps/aether-gateway/src/maintenance/runtime/config.rs
View File

@@ -88,8 +88,7 @@ pub(super) async fn pending_cleanup_batch_size(
) -> Result<usize, DataLayerError> {
Ok(system_config_usize(data, "cleanup_batch_size", 1_000)
.await?
.max(1)
.min(200))
.clamp(1, 200))
}
pub(super) async fn usage_cleanup_settings(

2
apps/aether-gateway/src/maintenance/runtime/usage_cleanup.rs
View File

@@ -205,7 +205,7 @@ async fn compress_usage_body_fields(
let mut total_compressed = 0usize;
let mut no_progress_count = 0usize;
let batch_size = batch_size.max(1).min(25);
let batch_size = batch_size.clamp(1, 25);
loop {
let rows = sqlx::query(SELECT_USAGE_BODY_COMPRESSION_BATCH_SQL)
.bind(cutoff_time)

3
apps/aether-gateway/src/model_fetch/tests.rs
View File

@@ -122,6 +122,9 @@ fn sample_global_model(id: &str, name: &str, mappings: &[&str]) -> StoredAdminGl
None,
None,
Some(json!({ "model_mappings": mappings })),
0,
0,
0,
Some(1_711_000_000),
Some(1_711_000_000),
)

76
apps/aether-gateway/src/request_candidate_runtime.rs
View File

@@ -7,7 +7,9 @@ use aether_scheduler_core::{
build_report_request_candidate_status_record,
finalize_execution_request_candidate_report_context, parse_request_candidate_report_context,
resolve_report_request_candidate_slot as resolve_report_request_candidate_slot_from_candidates,
SchedulerMinimalCandidateSelectionCandidate, SchedulerResolvedReportRequestCandidateSlot,
LocalRequestCandidateStatusRecordInput, ReportRequestCandidateStatusRecordInput,
SchedulerMinimalCandidateSelectionCandidate, SchedulerRequestCandidateStatusUpdate,
SchedulerResolvedReportRequestCandidateSlot,
};
use aether_usage_runtime::build_locally_actionable_report_context_from_request_candidate;
use async_trait::async_trait;
@@ -41,25 +43,15 @@ pub(crate) async fn record_local_request_candidate_status(
state: &(impl RequestCandidateRuntimeWriter + ?Sized),
plan: &ExecutionPlan,
report_context: Option<&Value>,
status: RequestCandidateStatus,
status_code: Option<u16>,
error_type: Option<String>,
error_message: Option<String>,
latency_ms: Option<u64>,
started_at_unix_secs: Option<u64>,
finished_at_unix_secs: Option<u64>,
status_update: SchedulerRequestCandidateStatusUpdate,
) {
let Some(record) = build_local_request_candidate_status_record(
plan,
report_context,
status,
status_code,
error_type,
error_message,
latency_ms,
started_at_unix_secs,
finished_at_unix_secs,
) else {
let Some(record) =
build_local_request_candidate_status_record(LocalRequestCandidateStatusRecordInput {
plan,
report_context,
status_update,
})
else {
return;
};
let candidate_id = record.id.clone();
@@ -80,13 +72,7 @@ pub(crate) async fn record_local_request_candidate_status(
pub(crate) async fn record_report_request_candidate_status(
state: &(impl RequestCandidateRuntimeReader + RequestCandidateRuntimeWriter + ?Sized),
report_context: Option<&Value>,
status: RequestCandidateStatus,
status_code: Option<u16>,
error_type: Option<String>,
error_message: Option<String>,
latency_ms: Option<u64>,
started_at_unix_secs: Option<u64>,
finished_at_unix_secs: Option<u64>,
status_update: SchedulerRequestCandidateStatusUpdate,
) {
let Some(slot) = resolve_report_request_candidate_slot(state, report_context).await else {
return;
@@ -95,17 +81,12 @@ pub(crate) async fn record_report_request_candidate_status(
let request_id_for_log = short_request_id(request_id.as_str());
let candidate_index = slot.candidate_index;
let retry_index = slot.retry_index;
let record = build_report_request_candidate_status_record(
slot,
status,
status_code,
error_type,
error_message,
latency_ms,
started_at_unix_secs,
finished_at_unix_secs,
current_unix_secs(),
);
let record =
build_report_request_candidate_status_record(ReportRequestCandidateStatusRecordInput {
slot,
status_update,
now_unix_secs: current_unix_secs(),
});
if let Err(err) = state.upsert_request_candidate(record).await {
warn!(
@@ -326,7 +307,10 @@ mod tests {
};
use serde_json::json;
use super::{ensure_execution_request_candidate_slot, record_report_request_candidate_status};
use super::{
ensure_execution_request_candidate_slot, record_report_request_candidate_status,
SchedulerRequestCandidateStatusUpdate,
};
use crate::data::GatewayDataState;
use crate::AppState;
@@ -494,13 +478,15 @@ mod tests {
record_report_request_candidate_status(
&state,
Some(&report_context),
RequestCandidateStatus::Success,
Some(200),
None,
None,
Some(25),
Some(101),
Some(102),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Success,
status_code: Some(200),
error_type: None,
error_message: None,
latency_ms: Some(25),
started_at_unix_secs: Some(101),
finished_at_unix_secs: Some(102),
},
)
.await;

83
apps/aether-gateway/src/router.rs
View File

@@ -1,5 +1,13 @@
use std::path::PathBuf;
use axum::extract::Request;
use axum::http::Method;
use axum::response::{IntoResponse, Response};
use axum::routing::any;
use axum::Router;
use tower::ServiceExt;
use tower_http::services::{ServeDir, ServeFile};
use tracing::warn;
use aether_runtime::{prometheus_response, ConcurrencyError, DistributedConcurrencyError};
@@ -9,6 +17,12 @@ pub fn build_router() -> Result<Router, reqwest::Error> {
Ok(build_router_with_state(AppState::new()?))
}
#[derive(Clone, Debug)]
struct FrontendStaticState {
static_dir: PathBuf,
index_html: PathBuf,
}
pub fn build_router_with_state(state: AppState) -> Router {
let cors_state = state.clone();
let mut router = Router::<AppState>::new();
@@ -32,6 +46,75 @@ pub fn build_router_with_state(state: AppState) -> Router {
router
}
pub fn attach_static_frontend(router: Router, static_dir: impl Into<PathBuf>) -> Router {
let static_dir = static_dir.into();
let index_html = static_dir.join("index.html");
router.layer(axum::middleware::from_fn_with_state(
FrontendStaticState {
static_dir,
index_html,
},
frontend_static_middleware,
))
}
async fn frontend_static_middleware(
axum::extract::State(frontend): axum::extract::State<FrontendStaticState>,
request: Request,
next: axum::middleware::Next,
) -> Response {
let path = request.uri().path().to_string();
if !matches!(request.method(), &Method::GET | &Method::HEAD)
|| frontend_path_bypasses_static(&path)
{
return next.run(request).await;
}
if frontend_path_targets_static_asset(&path) {
return serve_static_asset(&frontend.static_dir, request).await;
}
serve_frontend_index(&frontend.index_html, request).await
}
fn frontend_path_bypasses_static(path: &str) -> bool {
matches!(
path,
"/health" | "/test-connection" | crate::constants::READYZ_PATH
) || path.starts_with("/api/")
|| path.starts_with("/v1/")
|| path.starts_with("/v1beta/")
|| path.starts_with("/upload/")
|| path.starts_with("/_gateway/")
|| path.starts_with("/.well-known/")
}
fn frontend_path_targets_static_asset(path: &str) -> bool {
path.rsplit('/')
.next()
.is_some_and(|segment| !segment.is_empty() && segment.contains('.'))
}
async fn serve_static_asset(static_dir: &PathBuf, request: Request) -> Response {
match ServeDir::new(static_dir).oneshot(request).await {
Ok(response) => response.into_response(),
Err(err) => {
warn!(error = %err, "failed to serve frontend static asset");
axum::http::StatusCode::INTERNAL_SERVER_ERROR.into_response()
}
}
}
async fn serve_frontend_index(index_html: &PathBuf, request: Request) -> Response {
match ServeFile::new(index_html).oneshot(request).await {
Ok(response) => response.into_response(),
Err(err) => {
warn!(error = %err, "failed to serve frontend index");
axum::http::StatusCode::INTERNAL_SERVER_ERROR.into_response()
}
}
}
pub(crate) async fn metrics(
axum::extract::State(state): axum::extract::State<AppState>,
) -> impl axum::response::IntoResponse {

13
apps/aether-gateway/src/scheduler/candidate/runtime.rs
View File

@@ -4,7 +4,8 @@ use aether_data_contracts::repository::candidates::StoredRequestCandidate;
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use aether_scheduler_core::{
auth_api_key_concurrency_limit_reached, build_provider_concurrent_limit_map,
candidate_is_selectable_with_runtime_state, SchedulerAffinityTarget,
candidate_is_selectable_with_runtime_state, CandidateRuntimeSelectabilityInput,
SchedulerAffinityTarget,
};
use crate::data::auth::GatewayAuthApiKeySnapshot;
@@ -87,16 +88,16 @@ pub(super) fn is_candidate_selectable(
.copied()
.flatten();
candidate_is_selectable_with_runtime_state(
candidate_is_selectable_with_runtime_state(CandidateRuntimeSelectabilityInput {
candidate,
&snapshot.recent_candidates,
&snapshot.provider_concurrent_limits,
&snapshot.provider_key_rpm_states,
recent_candidates: &snapshot.recent_candidates,
provider_concurrent_limits: &snapshot.provider_concurrent_limits,
provider_key_rpm_states: &snapshot.provider_key_rpm_states,
now_unix_secs,
cached_affinity_target,
provider_quota_blocks_requests,
rpm_reset_at,
)
})
}
pub(super) async fn read_provider_concurrent_limits(

1
apps/aether-gateway/src/state/runtime/billing/finance_queries.rs
View File

@@ -160,7 +160,6 @@ impl AppState {
page.items
.into_iter()
.map(stored_admin_payment_callback_to_gateway)
.map(Into::into)
.collect(),
page.total,
)))

2
apps/aether-gateway/src/tests/architecture/admin_provider.rs
View File

@@ -1091,7 +1091,7 @@ fn admin_provider_write_uses_specific_local_owners() {
] {
assert!(
endpoint_keys_mutations.contains(pattern),
"handlers/admin/provider/endpoint_keys/mutations/mod.rs should expose explicit mutation owner {pattern}"
"handlers/admin/provider/endpoint_keys/mutations/mod.rs should expose explicit mutation owner {pattern}"
);
}

3
apps/aether-gateway/src/tests/architecture/usage.rs
View File

@@ -33,7 +33,8 @@ fn usage_runtime_paths_depend_on_shared_crates_not_app_runtime_shims() {
);
}
for path in ["apps/aether-gateway/src/async_task/runtime.rs"] {
{
let path = "apps/aether-gateway/src/async_task/runtime.rs";
let source = read_workspace_file(path);
assert!(
source.contains("aether_billing"),

84
apps/aether-gateway/src/tests/control/admin/endpoints/keys.rs
View File

@@ -116,6 +116,90 @@ async fn gateway_handles_admin_provider_keys_locally_with_trusted_admin_principa
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_admin_provider_keys_prefers_upstream_plan_type_over_auth_config() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/endpoints/providers/provider-codex/keys",
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 mut provider = sample_provider("provider-codex", "codex", 10);
provider.provider_type = "codex".to_string();
let mut key = sample_key(
"key-codex-oauth",
"provider-codex",
"openai:cli",
"oauth-placeholder",
);
key.auth_type = "oauth".to_string();
key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
&json!({
"plan_type": "free",
"account_id": "acct-codex-legacy"
})
.to_string(),
)
.expect("auth config should encrypt"),
);
key.upstream_metadata = Some(json!({
"codex": {
"plan_type": "plus",
"updated_at": 1_775_553_285u64
}
}));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![],
vec![key],
));
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::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/endpoints/providers/provider-codex/keys?skip=0&limit=50"
))
.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::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
let items = payload.as_array().expect("payload should be an array");
assert_eq!(items.len(), 1);
assert_eq!(items[0]["oauth_plan_type"], "plus");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_creates_admin_provider_key_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));

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

@@ -92,6 +92,9 @@ async fn gateway_handles_admin_global_models_locally_with_trusted_admin_principa
payload["models"].as_array().expect("models array")[0]["name"],
"gpt-4.1"
);
assert_eq!(payload["models"][0]["provider_count"], 1);
assert_eq!(payload["models"][0]["active_provider_count"], 1);
assert_eq!(payload["models"][0]["usage_count"], 0);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
@@ -311,6 +314,9 @@ async fn gateway_handles_admin_global_model_detail_locally_with_trusted_admin_pr
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["id"], "global-gpt-5");
assert_eq!(payload["provider_count"], 1);
assert_eq!(payload["active_provider_count"], 1);
assert_eq!(payload["usage_count"], 0);
assert_eq!(payload["total_models"], 1);
assert_eq!(payload["total_providers"], 1);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);

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

@@ -1,9 +1,9 @@
use std::sync::{Arc, Mutex};
use aether_crypto::DEVELOPMENT_ENCRYPTION_KEY;
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data_contracts::repository::provider_catalog::ProviderCatalogReadRepository;
use axum::body::{Body, Bytes};
use axum::body::{to_bytes, Body, Bytes};
use axum::routing::{any, get, post};
use axum::{extract::Request, Router};
use http::{HeaderMap, HeaderValue, StatusCode};
@@ -405,6 +405,68 @@ async fn gateway_handles_admin_pool_trailing_slash_routes_locally_with_trusted_a
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_pool_list_includes_usage_totals_and_nullable_lru_score() {
let provider = sample_provider("provider-openai", "openai", 10).with_transport_fields(
true,
false,
true,
None,
None,
None,
None,
None,
Some(json!({
"pool_advanced": {
"enabled": true
}
})),
);
let mut key = sample_key(
"key-openai-usage",
"provider-openai",
"openai:chat",
"sk-usage",
);
key.name = "usage key".to_string();
key.request_count = Some(1566);
key.total_tokens = 187_327_321;
key.total_cost_usd = 93.1319297;
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-openai/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.len(), 1);
assert_eq!(keys[0]["request_count"], json!(1566));
assert_eq!(keys[0]["total_tokens"], json!(187_327_321u64));
assert_eq!(keys[0]["total_cost_usd"], json!("93.13192970"));
assert!(keys[0]["lru_score"].is_null());
}
#[tokio::test]
async fn gateway_handles_admin_pool_list_keys_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
@@ -563,6 +625,495 @@ async fn gateway_handles_admin_pool_list_keys_locally_with_trusted_admin_princip
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_pool_list_keys_with_quota_compatibility_fields() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/pool/provider-antigravity/keys",
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 mut provider = sample_provider("provider-antigravity", "antigravity", 10)
.with_transport_fields(
true,
false,
true,
None,
None,
None,
None,
None,
Some(json!({
"pool_advanced": {
"enabled": true
}
})),
);
provider.provider_type = "antigravity".to_string();
let mut key = sample_key(
"key-antigravity-a",
"provider-antigravity",
"gemini:chat",
"sk-antigravity",
);
key.name = "quota-key".to_string();
key.auth_type = "oauth".to_string();
key.expires_at_unix_secs = Some(1_775_556_730);
key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"plan_type":"pro","account_id":"acct-antigravity-1","account_name":"quota-user","account_user_id":"quota-user-1","organizations":[{"id":"org-1","name":"Org One"}]}"#,
)
.expect("auth config ciphertext should build"),
);
key.status_snapshot = Some(json!({
"oauth": {
"code": "expired",
"label": "已过期",
"reason": "Token 已过期，请重新授权",
"expires_at": 1775556730u64,
"invalid_at": null,
"source": "expires_at",
"requires_reauth": true,
"expiring_soon": false
},
"account": {
"code": "ok",
"label": null,
"reason": null,
"blocked": false,
"source": null,
"recoverable": false
},
"quota": {
"code": "ok",
"label": null,
"reason": null,
"exhausted": false,
"usage_ratio": 0.0,
"updated_at": 1775553285u64,
"reset_seconds": null,
"plan_type": null
}
}));
key.upstream_metadata = Some(json!({
"antigravity": {
"updated_at": 1775553285u64,
"quota_by_model": {
"gemini-2.5-flash": { "used_percent": 0.0 },
"gemini-2.5-pro": { "used_percent": 0.0 }
}
}
}));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
Vec::new(),
vec![key],
));
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::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/pool/provider-antigravity/keys?page=1&page_size=10&status=all"
))
.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::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
let keys = payload["keys"].as_array().expect("keys should be array");
assert_eq!(keys.len(), 1);
assert_eq!(keys[0]["account_quota"], json!("最低剩余 100.0% (2 模型)"));
assert_eq!(keys[0]["quota_updated_at"], json!(1775553285u64));
assert_eq!(keys[0]["oauth_expires_at"], json!(1775556730u64));
assert_eq!(keys[0]["oauth_plan_type"], json!("pro"));
assert_eq!(keys[0]["oauth_account_id"], json!("acct-antigravity-1"));
assert_eq!(keys[0]["oauth_account_name"], json!("quota-user"));
assert_eq!(keys[0]["oauth_account_user_id"], json!("quota-user-1"));
assert_eq!(keys[0]["oauth_organizations"][0]["id"], json!("org-1"));
assert_eq!(keys[0]["account_status_code"], json!("ok"));
assert_eq!(keys[0]["account_status_blocked"], json!(false));
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_includes_pool_quota_and_compat_fields_in_list_keys_response() {
let mut provider = sample_provider("provider-antigravity", "antigravity", 10)
.with_transport_fields(
true,
false,
true,
None,
None,
None,
None,
None,
Some(json!({
"pool_advanced": {
"enabled": true
}
})),
);
provider.provider_type = "antigravity".to_string();
let mut key = sample_key(
"key-antigravity-oauth",
"provider-antigravity",
"gemini:chat",
"oauth-placeholder",
);
key.name = "quota key".to_string();
key.auth_type = "oauth".to_string();
key.expires_at_unix_secs = Some(1_775_556_730);
key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
&json!({
"plan_type": "pro",
"account_id": "acct-demo-001",
"account_name": "Demo Account",
"account_user_id": "user-demo-001",
"organizations": [],
})
.to_string(),
)
.expect("auth config should encrypt"),
);
key.upstream_metadata = Some(json!({
"antigravity": {
"updated_at": 1_775_553_285u64,
"quota_by_model": {
"gemini-2.5-pro": { "used_percent": 0 },
"gemini-2.5-flash": { "used_percent": 0 }
}
}
}));
key.status_snapshot = Some(json!({
"oauth": {
"code": "expired",
"label": "已过期",
"reason": "Token 已过期，请重新授权",
"expires_at": 1_775_556_730u64,
"invalid_at": serde_json::Value::Null,
"source": "expires_at",
"requires_reauth": true,
"expiring_soon": false
},
"account": {
"code": "ok",
"label": serde_json::Value::Null,
"reason": serde_json::Value::Null,
"blocked": false,
"source": serde_json::Value::Null,
"recoverable": false
},
"quota": {
"code": "ok",
"label": serde_json::Value::Null,
"reason": serde_json::Value::Null,
"exhausted": false,
"usage_ratio": 0.0,
"updated_at": 1_775_553_285u64,
"reset_seconds": serde_json::Value::Null,
"plan_type": serde_json::Value::Null
}
}));
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::GET,
"/api/admin/pool/provider-antigravity/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.len(), 1);
assert_eq!(keys[0]["account_quota"], "最低剩余 100.0% (2 模型)");
assert_eq!(keys[0]["quota_updated_at"], json!(1_775_553_285u64));
assert_eq!(keys[0]["oauth_expires_at"], json!(1_775_556_730u64));
assert_eq!(keys[0]["oauth_plan_type"], "pro");
assert_eq!(keys[0]["oauth_account_id"], "acct-demo-001");
assert_eq!(keys[0]["oauth_account_name"], "Demo Account");
assert_eq!(keys[0]["oauth_account_user_id"], "user-demo-001");
assert_eq!(keys[0]["oauth_organizations"], json!([]));
assert_eq!(keys[0]["account_status_code"], "ok");
assert_eq!(keys[0]["account_status_blocked"], json!(false));
}
#[tokio::test]
async fn gateway_formats_codex_quota_countdown_from_reset_after_seconds() {
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
}
})),
);
provider.provider_type = "codex".to_string();
let mut key = sample_key(
"key-codex-oauth",
"provider-codex",
"openai:cli",
"oauth-placeholder",
);
key.name = "codex quota key".to_string();
key.auth_type = "oauth".to_string();
key.upstream_metadata = Some(json!({
"codex": {
"plan_type": "plus",
"updated_at": 1_775_553_285u64,
"primary_used_percent": 10.0,
"primary_reset_after_seconds": 266_400,
"secondary_used_percent": 33.0,
"secondary_reset_after_seconds": 13_800
}
}));
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.len(), 1);
assert_eq!(
keys[0]["account_quota"],
"周剩余 90.0% (3天2小时后重置) | 5H剩余 67.0% (3小时50分钟后重置)"
);
}
#[tokio::test]
async fn gateway_pool_prefers_upstream_plan_type_over_auth_config() {
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
}
})),
);
provider.provider_type = "codex".to_string();
let mut key = sample_key(
"key-codex-precedence",
"provider-codex",
"openai:cli",
"oauth-placeholder",
);
key.auth_type = "oauth".to_string();
key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
&json!({
"plan_type": "free",
"account_id": "acct-codex-legacy"
})
.to_string(),
)
.expect("auth config should encrypt"),
);
key.upstream_metadata = Some(json!({
"codex": {
"plan_type": "plus",
"updated_at": 1_775_553_285u64
}
}));
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::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.len(), 1);
assert_eq!(keys[0]["oauth_plan_type"], "plus");
}
#[tokio::test]
async fn gateway_pool_plan_free_selector_prefers_upstream_plan_type() {
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
}
})),
);
provider.provider_type = "codex".to_string();
let mut key = sample_key(
"key-codex-selector",
"provider-codex",
"openai:cli",
"oauth-placeholder",
);
key.auth_type = "oauth".to_string();
key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
&json!({
"plan_type": "free",
"account_id": "acct-codex-legacy"
})
.to_string(),
)
.expect("auth config should encrypt"),
);
key.upstream_metadata = Some(json!({
"codex": {
"plan_type": "plus",
"updated_at": 1_775_553_285u64
}
}));
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-codex/keys/resolve-selection",
Some(json!({
"quick_selectors": ["plan_free"]
})),
)
.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!(0));
assert_eq!(
payload["items"]
.as_array()
.expect("items should be array")
.len(),
0
);
}
#[tokio::test]
async fn gateway_handles_admin_pool_resolve_selection_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));

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

@@ -1,16 +1,17 @@
use std::sync::{Arc, Mutex};
use aether_data::repository::global_models::InMemoryGlobalModelReadRepository;
use aether_contracts::ExecutionPlan;
use aether_crypto::DEVELOPMENT_ENCRYPTION_KEY;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogEndpoint;
use axum::body::Body;
use axum::routing::any;
use axum::{extract::Request, Router};
use axum::{extract::Request, Json, Router};
use http::StatusCode;
use serde_json::json;
use super::super::{
build_router_with_state, sample_admin_global_model, sample_admin_provider_model, sample_key,
build_router_with_state, build_state_with_execution_runtime_override, sample_key,
sample_provider, start_server, AppState,
};
use crate::constants::{
@@ -63,22 +64,48 @@ async fn assert_admin_provider_query_route(
}
#[tokio::test]
async fn gateway_handles_admin_provider_query_models_locally_with_trusted_admin_principal() {
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 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 {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
*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![sample_provider("provider-openai", "OpenAI", 10)],
vec![provider],
vec![StoredProviderCatalogEndpoint::new(
"endpoint-openai-chat".to_string(),
"provider-openai".to_string(),
@@ -89,7 +116,7 @@ async fn gateway_handles_admin_provider_query_models_locally_with_trusted_admin_
)
.expect("endpoint should build")
.with_transport_fields(
"https://api.openai.com/v1".to_string(),
"https://api.openai.example".to_string(),
None,
None,
None,
@@ -99,67 +126,20 @@ async fn gateway_handles_admin_provider_query_models_locally_with_trusted_admin_
None,
)
.expect("endpoint transport should build")],
vec![
{
let mut key = sample_key(
"key-openai-allowed",
"provider-openai",
"openai:chat",
"sk-test",
);
key.allowed_models = Some(json!(["gpt-5"]));
key
},
sample_key(
"key-openai-all",
"provider-openai",
"openai:chat",
"sk-test-2",
),
],
vec![sample_key(
"key-openai-selected",
"provider-openai",
"openai:chat",
"sk-test",
)],
));
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"),
sample_admin_global_model("global-gpt-4.1", "gpt-4.1", "GPT 4.1"),
])
.with_admin_provider_models(vec![
{
let mut model = sample_admin_provider_model(
"provider-model-gpt-5",
"provider-openai",
"global-gpt-5",
"gpt-5",
);
model.global_model_name = Some("gpt-5".to_string());
model.global_model_display_name = Some("GPT 5".to_string());
model
},
{
let mut model = sample_admin_provider_model(
"provider-model-gpt-4.1",
"provider-openai",
"global-gpt-4.1",
"gpt-4.1",
);
model.global_model_name = Some("gpt-4.1".to_string());
model.global_model_display_name = Some("GPT 4.1".to_string());
model
},
]),
);
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::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
)
.with_global_model_repository_for_tests(global_model_repository),
),
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;
@@ -171,7 +151,7 @@ async fn gateway_handles_admin_provider_query_models_locally_with_trusted_admin_
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-openai",
"api_key_id": "key-openai-allowed"
"api_key_id": "key-openai-selected"
}))
.send()
.await
@@ -184,33 +164,167 @@ async fn gateway_handles_admin_provider_query_models_locally_with_trusted_admin_
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!(true));
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_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: Vec<_> = models
let model_ids = models
.iter()
.map(|model| {
(
model["id"].as_str().expect("id should be present"),
model["display_name"]
.as_str()
.expect("display_name should be present"),
)
})
.collect();
assert_eq!(model_ids, vec![("gpt-4.1", "GPT 4.1"), ("gpt-5", "GPT 5")]);
for model in models {
assert_eq!(model["owned_by"], "OpenAI");
assert_eq!(model["api_format"], "openai:chat");
assert_eq!(model["api_formats"], json!(["openai:chat"]));
}
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
.map(|model| model["id"].as_str().expect("id should exist"))
.collect::<Vec<_>>();
assert_eq!(model_ids, vec!["gpt-5", "gpt-4.1"]);
assert_eq!(
*execution_runtime_hits.lock().expect("mutex should lock"),
2
);
gateway_handle.abort();
upstream_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]

3
apps/aether-gateway/src/tests/control/helpers.rs
View File

@@ -351,6 +351,9 @@ pub(super) fn sample_admin_global_model(
})),
Some(json!(["streaming", "vision"])),
Some(json!({"streaming": true, "vision": false, "billing": {"currency": "USD"}})),
0,
0,
0,
Some(1_711_000_000),
Some(1_711_000_100),
)

79
apps/aether-gateway/src/tests/frontdoor/ai.rs
View File

@@ -174,6 +174,85 @@ async fn gateway_handles_public_openai_models_without_hitting_fallback_probe() {
fallback_probe_handle.abort();
}
#[tokio::test]
async fn gateway_handles_public_openai_models_with_cross_format_candidates_without_hitting_fallback_probe(
) {
let fallback_probe_hits = Arc::new(Mutex::new(0usize));
let fallback_probe_hits_clone = Arc::clone(&fallback_probe_hits);
let fallback_probe = Router::new().route(
"/{*path}",
any(move |_request: Request| {
let fallback_probe_hits_inner = Arc::clone(&fallback_probe_hits_clone);
async move {
*fallback_probe_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("proxied"))
}
}),
);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-openai-models-cross-format")),
unrestricted_models_snapshot("key-1", "user-1"),
)]));
let candidate_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
sample_models_candidate_row(
"provider-claude",
"claude",
"claude:chat",
"claude-3-7-sonnet",
10,
),
]));
let (_unused_fallback_probe_url, fallback_probe_handle) = start_server(fallback_probe).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
crate::data::GatewayDataState::with_minimal_candidate_selection_and_auth_for_tests(
candidate_repository,
auth_repository,
),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
let list_response = client
.get(format!("{gateway_url}/v1/models"))
.header("authorization", "Bearer sk-openai-models-cross-format")
.send()
.await
.expect("request should succeed");
assert_eq!(list_response.status(), StatusCode::OK);
let list_payload: serde_json::Value =
list_response.json().await.expect("json body should parse");
assert_eq!(list_payload["object"], "list");
assert_eq!(list_payload["data"][0]["id"], "claude-3-7-sonnet");
assert_eq!(list_payload["data"][0]["owned_by"], "claude");
let detail_response = client
.get(format!("{gateway_url}/v1/models/claude-3-7-sonnet"))
.header("authorization", "Bearer sk-openai-models-cross-format")
.send()
.await
.expect("request should succeed");
assert_eq!(detail_response.status(), StatusCode::OK);
let detail_payload: serde_json::Value = detail_response
.json()
.await
.expect("json body should parse");
assert_eq!(detail_payload["id"], "claude-3-7-sonnet");
assert_eq!(detail_payload["owned_by"], "claude");
assert_eq!(*fallback_probe_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
fallback_probe_handle.abort();
}
#[tokio::test]
async fn gateway_handles_public_claude_models_without_hitting_fallback_probe() {
let fallback_probe_hits = Arc::new(Mutex::new(0usize));

106
apps/aether-gateway/src/tests/frontdoor/core.rs
View File

@@ -1,5 +1,9 @@
use std::fs;
use std::time::{SystemTime, UNIX_EPOCH};
use crate::tests::{
any, build_router, start_server, Arc, Body, Mutex, Request, Router, StatusCode, READYZ_PATH,
any, attach_static_frontend, build_router, start_server, Arc, Body, Mutex, Request, Router,
StatusCode, READYZ_PATH,
};
#[tokio::test]
@@ -112,64 +116,69 @@ async fn gateway_handles_public_service_health_without_proxying_upstream() {
}
#[tokio::test]
async fn gateway_handles_public_root_without_proxying_upstream() {
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("proxied"))
}
}),
);
async fn gateway_serves_frontend_routes_and_assets_without_shadowing_public_api() {
let unique_suffix = SystemTime::now()
.duration_since(UNIX_EPOCH)
.expect("system clock should be monotonic enough for tests")
.as_nanos();
let static_dir =
std::env::temp_dir().join(format!("aether-gateway-static-test-{unique_suffix}"));
let assets_dir = static_dir.join("assets");
fs::create_dir_all(&assets_dir).expect("static assets dir should be created");
fs::write(
static_dir.join("index.html"),
"<!doctype html><html><body>Aether Frontend</body></html>",
)
.expect("index.html should be written");
fs::write(assets_dir.join("app.js"), "console.log('frontend asset');")
.expect("asset file should be written");
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router().expect("gateway should build");
let gateway =
attach_static_frontend(build_router().expect("gateway should build"), &static_dir);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
let response = reqwest::Client::new()
let response = client
.get(format!("{gateway_url}/"))
.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["status"], "running");
let content_type = response
.headers()
.get("content-type")
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string();
let body = response.text().await.expect("html body should be readable");
assert!(content_type.starts_with("text/html"));
assert!(body.contains("Aether Frontend"));
let response = client
.get(format!("{gateway_url}/guide"))
.send()
.await
.expect("spa request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let body = response.text().await.expect("spa body should be readable");
assert!(body.contains("Aether Frontend"));
let response = client
.get(format!("{gateway_url}/assets/app.js"))
.send()
.await
.expect("asset request should succeed");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
payload["message"],
"AI Proxy with Modular Architecture v4.0.0"
);
assert_eq!(payload["endpoints"]["health"], "/v1/health");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_public_site_info_without_proxying_upstream() {
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("proxied"))
}
}),
response
.text()
.await
.expect("asset body should be readable"),
"console.log('frontend asset');"
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router().expect("gateway should build");
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
let response = client
.get(format!("{gateway_url}/api/public/site-info"))
.send()
.await
@@ -179,8 +188,7 @@ async fn gateway_handles_public_site_info_without_proxying_upstream() {
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["site_name"], "Aether");
assert_eq!(payload["site_subtitle"], "AI Gateway");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
let _ = fs::remove_dir_all(&static_dir);
}

12
apps/aether-gateway/src/tests/frontdoor/public_support.rs
View File

@@ -3871,15 +3871,15 @@ async fn gateway_handles_wallet_balance_locally_without_proxying_upstream() {
#[tokio::test]
async fn gateway_handles_wallet_today_cost_locally_without_proxying_upstream() {
let now = Utc::now();
let auth_now = Utc::now();
let usage_now = chrono::DateTime::<chrono::Utc>::from_naive_utc_and_offset(
Utc::now()
auth_now
.date_naive()
.and_hms_opt(12, 0, 0)
.expect("midday should be valid"),
chrono::Utc,
);
let user = sample_auth_user(now);
let user = sample_auth_user(auth_now);
let access_token = build_test_auth_token(
"access",
serde_json::Map::from_iter([
@@ -3891,7 +3891,7 @@ async fn gateway_handles_wallet_today_cost_locally_without_proxying_upstream() {
),
("session_id".to_string(), json!("session-wallet-today-1")),
]),
now + chrono::Duration::hours(1),
auth_now + chrono::Duration::hours(1),
);
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
sample_user_usage_audit(
@@ -3916,13 +3916,13 @@ async fn gateway_handles_wallet_today_cost_locally_without_proxying_upstream() {
let (gateway_url, upstream_hits, gateway_handle, upstream_handle) =
start_auth_gateway_with_usage_state(
user,
sample_auth_wallet("user-auth-1", now),
sample_auth_wallet("user-auth-1", auth_now),
[sample_auth_session(
"user-auth-1",
"session-wallet-today-1",
"device-wallet-today-1",
"refresh-token-placeholder",
now,
auth_now,
)],
usage_repository,
)

2
apps/aether-gateway/src/tests/mod.rs
View File

@@ -25,7 +25,7 @@ pub(super) use super::async_task::VideoTaskTruthSourceMode;
pub(super) use super::constants::*;
pub(super) use super::fallback_metrics::{GatewayFallbackMetricKind, GatewayFallbackReason};
pub(super) use super::rate_limit::FrontdoorUserRpmConfig;
pub(super) use super::router::{build_router, build_router_with_state};
pub(super) use super::router::{attach_static_frontend, build_router, build_router_with_state};
pub(super) use super::state::{AppState, FrontdoorCorsConfig};
pub(super) use super::usage::UsageRuntimeConfig;

34
apps/aether-gateway/src/usage/reporting/mod.rs
View File

@@ -2,7 +2,7 @@ use std::collections::BTreeMap;
use aether_contracts::ExecutionError;
use aether_data_contracts::repository::candidates::RequestCandidateStatus;
use aether_scheduler_core::execution_error_details;
use aether_scheduler_core::{execution_error_details, SchedulerRequestCandidateStatusUpdate};
use tracing::{debug, warn};
use uuid::Uuid;
@@ -203,13 +203,15 @@ async fn handle_local_sync_report(state: &AppState, payload: &GatewaySyncReportR
record_report_request_candidate_status(
state,
payload.report_context.as_ref(),
status,
Some(payload.status_code),
error_type,
error_message,
latency_ms,
Some(terminal_unix_secs),
Some(terminal_unix_secs),
SchedulerRequestCandidateStatusUpdate {
status,
status_code: Some(payload.status_code),
error_type,
error_message,
latency_ms,
started_at_unix_secs: Some(terminal_unix_secs),
finished_at_unix_secs: Some(terminal_unix_secs),
},
)
.await;
}
@@ -223,13 +225,15 @@ async fn handle_local_stream_report(state: &AppState, payload: &GatewayStreamRep
record_report_request_candidate_status(
state,
payload.report_context.as_ref(),
RequestCandidateStatus::Success,
Some(payload.status_code),
None,
None,
latency_ms,
Some(terminal_unix_secs),
Some(terminal_unix_secs),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Success,
status_code: Some(payload.status_code),
error_type: None,
error_message: None,
latency_ms,
started_at_unix_secs: Some(terminal_unix_secs),
finished_at_unix_secs: Some(terminal_unix_secs),
},
)
.await;
}

7
apps/aether-gateway/src/video_tasks/tests/projection.rs
View File

@@ -119,8 +119,11 @@ fn rust_authoritative_service_builds_openai_content_stream_plan_from_direct_vide
let LocalVideoTaskContentAction::StreamPlan(plan) = action else {
panic!("content action should be stream plan");
};
assert_eq!(plan.method, "GET");
assert_eq!(plan.url, "https://cdn.example.com/ext-video-task-123.mp4");
assert_eq!(plan.method.as_str(), "GET");
assert_eq!(
plan.url.as_str(),
"https://cdn.example.com/ext-video-task-123.mp4"
);
assert!(plan.headers.is_empty());
}