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

Fix/api key concurrency runtime miss (#309)

Browse Source 
* test(cli): 覆盖 API key 并发等待与超时路径

* feat(scheduler): API key 并发饱和时等待可用槽位

* fix(proxy): 区分 API key 并发受限与真正的 runtime miss

* fix(outcome): runtime miss 仅归因真实执行候选

* feat(api-keys): 统一 concurrent_limit 默认值与校验辅助

* feat(admin): 独立 Key 接口支持 concurrent_limit

* feat(admin): 用户 API Key 路由支持 concurrent_limit

* feat(public): 自助 API Key 路由支持 concurrent_limit

* feat(import): 导入与存储层持久化 concurrent_limit

* feat(frontend): 同步 API Key concurrent_limit 类型定义

* feat(frontend): 独立 Key 表单支持 concurrent_limit

* feat(frontend): 管理员用户 API Key 表单支持 concurrent_limit

* feat(frontend): 自助 API Key 页面支持 concurrent_limit

* chore(fmt): 统一 runtime 归因相关 Rust 格式

* chore(fmt): 统一 admin API key 路由 Rust 格式

* chore(fmt): 统一 public 路由与相关测试 Rust 格式

* fix(test): 对齐 no-execution usage 归因断言

* test(middleware): 固定 access log tracing 用例线程模型

* fix(frontend): 提取用户 API Key payload 默认并发辅助

* fix(frontend): 保留用户 Key 的 concurrent_limit 默认值

* fix(api-keys): remove hardcoded concurrent limit default

---------

Co-authored-by: fawney19 <elky0401@gmail.com>
This commit is contained in:
RWDai
2026-04-17 14:21:43 +08:00
committed by GitHub
parent e5d3722adf
commit b8702ae124
39 changed files with 1739 additions and 128 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
apps/aether-gateway/src/ai_pipeline/planner/common.rs
View File

@@ -101,6 +101,16 @@ pub(crate) fn apply_local_candidate_terminal_plan_reason(
let skipped_candidate_count = diagnostic.skipped_candidate_count.unwrap_or(0);
diagnostic.reason = if candidate_count == 0 {
"candidate_list_empty".to_string()
} else if skipped_candidate_count >= candidate_count
&& diagnostic.skip_reasons.len() == 1
&& diagnostic
.skip_reasons
.get("api_key_concurrency_limit_reached")
.copied()
.unwrap_or(0)
> 0
{
"api_key_concurrency_limit_reached".to_string()
} else if skipped_candidate_count >= candidate_count {
"all_candidates_skipped".to_string()
} else {
@@ -222,7 +232,15 @@ mod tests {
apply_local_candidate_terminal_plan_reason(&mut diagnostic, "no_local_sync_plans");
assert_eq!(diagnostic.reason, "all_candidates_skipped");
diagnostic.skip_reasons = std::collections::BTreeMap::from([(
"api_key_concurrency_limit_reached".to_string(),
2,
)]);
apply_local_candidate_terminal_plan_reason(&mut diagnostic, "no_local_sync_plans");
assert_eq!(diagnostic.reason, "api_key_concurrency_limit_reached");
diagnostic.skipped_candidate_count = Some(1);
diagnostic.skip_reasons.clear();
apply_local_candidate_terminal_plan_reason(&mut diagnostic, "no_local_sync_plans");
assert_eq!(diagnostic.reason, "no_local_sync_plans");
}

2
apps/aether-gateway/src/ai_pipeline/planner/materialization_policy.rs
View File

@@ -47,7 +47,7 @@ pub(crate) fn build_local_candidate_persistence_policy<'a>(
LocalCandidatePersistencePolicyKind::OpenAiCliDecision => (
"gateway local openai cli decision request candidate upsert failed",
"gateway local openai cli decision failed to persist skipped candidate",
false,
true,
),
LocalCandidatePersistencePolicyKind::GeminiFilesDecision => (
"gateway local gemini files request candidate upsert failed",

90
apps/aether-gateway/src/ai_pipeline/planner/state/scheduler.rs
View File

@@ -1,6 +1,12 @@
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use std::time::Duration;
use tokio::time::Instant;
use super::{GatewayAuthApiKeySnapshot, PlannerAppState};
use crate::clock::current_unix_secs;
use crate::constants::{
API_KEY_CONCURRENCY_WAIT_POLL_INTERVAL_MS, API_KEY_CONCURRENCY_WAIT_TIMEOUT_MS,
};
use crate::scheduler::candidate::SchedulerSkippedCandidate;
use crate::GatewayError;
@@ -42,17 +48,39 @@ impl<'a> PlannerAppState<'a> {
),
GatewayError,
> {
crate::scheduler::candidate::list_selectable_candidates_with_skip_reasons(
self.app().data.as_ref(),
self.app(),
api_format,
global_model_name,
require_streaming,
required_capabilities,
auth_snapshot,
now_unix_secs,
)
.await
let wait_timeout = Duration::from_millis(API_KEY_CONCURRENCY_WAIT_TIMEOUT_MS);
let wait_interval = Duration::from_millis(API_KEY_CONCURRENCY_WAIT_POLL_INTERVAL_MS.max(1));
let wait_deadline = Instant::now() + wait_timeout;
let mut attempt_now_unix_secs = now_unix_secs;
loop {
let result = crate::scheduler::candidate::list_selectable_candidates_with_skip_reasons(
self.app().data.as_ref(),
self.app(),
api_format,
global_model_name,
require_streaming,
required_capabilities,
auth_snapshot,
attempt_now_unix_secs,
)
.await?;
if !crate::scheduler::candidate::is_exact_all_skipped_by_auth_limit(
&result.0, &result.1,
) {
return Ok(result);
}
let now = Instant::now();
if now >= wait_deadline {
return Ok(result);
}
let remaining = wait_deadline.duration_since(now);
tokio::time::sleep(wait_interval.min(remaining)).await;
attempt_now_unix_secs = current_unix_secs();
}
}
pub(crate) async fn list_selectable_candidates_for_required_capability_without_requested_model(
@@ -63,15 +91,35 @@ impl<'a> PlannerAppState<'a> {
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
now_unix_secs: u64,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
crate::scheduler::candidate::list_selectable_candidates_for_required_capability_without_requested_model(
self.app().data.as_ref(),
self.app(),
candidate_api_format,
required_capability,
require_streaming,
auth_snapshot,
now_unix_secs,
)
.await
let wait_timeout = Duration::from_millis(API_KEY_CONCURRENCY_WAIT_TIMEOUT_MS);
let wait_interval = Duration::from_millis(API_KEY_CONCURRENCY_WAIT_POLL_INTERVAL_MS.max(1));
let wait_deadline = Instant::now() + wait_timeout;
let mut attempt_now_unix_secs = now_unix_secs;
loop {
let (result, auth_limit_blocked) = crate::scheduler::candidate::list_selectable_candidates_for_required_capability_without_requested_model_with_auth_limit_signal(
self.app().data.as_ref(),
self.app(),
candidate_api_format,
required_capability,
require_streaming,
auth_snapshot,
attempt_now_unix_secs,
)
.await?;
if !auth_limit_blocked {
return Ok(result);
}
let now = Instant::now();
if now >= wait_deadline {
return Ok(result);
}
let remaining = wait_deadline.duration_since(now);
tokio::time::sleep(wait_interval.min(remaining)).await;
attempt_now_unix_secs = current_unix_secs();
}
}
}

4
apps/aether-gateway/src/constants.rs
View File

@@ -26,6 +26,10 @@ pub(crate) const EXECUTION_PATH_EXECUTION_RUNTIME_STREAM: &str = "execution_runt
pub(crate) const EXECUTION_PATH_CONTROL_EXECUTE_SYNC: &str = "control_execute_sync";
pub(crate) const EXECUTION_PATH_CONTROL_EXECUTE_STREAM: &str = "control_execute_stream";
pub(crate) const EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS: &str = "local_execution_runtime_miss";
pub(crate) const EXECUTION_PATH_LOCAL_API_KEY_CONCURRENCY_LIMITED: &str =
"local_api_key_concurrency_limited";
pub(crate) const API_KEY_CONCURRENCY_WAIT_TIMEOUT_MS: u64 = 150;
pub(crate) const API_KEY_CONCURRENCY_WAIT_POLL_INTERVAL_MS: u64 = 10;
pub(crate) const EXECUTION_PATH_LOCAL_AUTH_DENIED: &str = "local_auth_denied";
pub(crate) const EXECUTION_PATH_LOCAL_RATE_LIMITED: &str = "local_rate_limited";
pub(crate) const EXECUTION_PATH_LOCAL_ROUTE_NOT_FOUND: &str = "local_route_not_found";

150
apps/aether-gateway/src/executor/outcome.rs
View File

@@ -1,5 +1,5 @@
use std::collections::{BTreeMap, BTreeSet};
use std::time::Instant;
use std::time::{Duration, Instant};
use aether_contracts::ExecutionPlan;
use aether_data_contracts::repository::candidates::{
@@ -73,6 +73,23 @@ impl LocalExecutionRuntimeMissContext {
self.candidate_contexts.len()
}
pub(crate) fn all_candidates_skipped_for_reason(&self, reason: &str) -> bool {
let reason = reason.trim();
if reason.is_empty() || self.candidate_contexts.is_empty() {
return false;
}
self.candidate_contexts.iter().all(|candidate| {
candidate.candidate.status == RequestCandidateStatus::Skipped
&& candidate
.candidate
.skip_reason
.as_deref()
.map(str::trim)
.is_some_and(|value| value == reason)
})
}
pub(crate) fn candidate_summary(&self) -> Option<String> {
const MAX_ITEMS: usize = 5;
@@ -166,7 +183,10 @@ pub(crate) async fn build_local_execution_runtime_miss_context(
auth_api_key_id: auth_context.map(|value| value.api_key_id.clone()),
auth_username: auth_context.and_then(|value| value.username.clone()),
auth_api_key_name: auth_context.and_then(|value| value.api_key_name.clone()),
candidate_contexts: load_runtime_miss_candidate_contexts(state, request_id, decision).await,
candidate_contexts: load_runtime_miss_candidate_contexts_with_retry(
state, request_id, decision,
)
.await,
}
}
@@ -175,6 +195,7 @@ pub(crate) async fn record_failed_usage_for_exhausted_request(
exhaustion: LocalExecutionExhaustion,
started_at: &Instant,
local_execution_runtime_miss_detail: &str,
execution_path: &str,
diagnostic: Option<&LocalExecutionRuntimeMissDiagnostic>,
) {
if !state.usage_runtime.is_enabled() {
@@ -244,6 +265,7 @@ pub(crate) async fn record_failed_usage_for_exhausted_request(
apply_runtime_miss_usage_routing(
&mut data,
&mut request_metadata,
execution_path,
candidate_id.as_deref(),
candidate_index,
None,
@@ -264,6 +286,7 @@ pub(crate) async fn record_failed_usage_for_runtime_miss_request(
request_id: &str,
started_at: &Instant,
local_execution_runtime_miss_detail: &str,
execution_path: &str,
decision: Option<&GatewayControlDecision>,
diagnostic: Option<&LocalExecutionRuntimeMissDiagnostic>,
context: &LocalExecutionRuntimeMissContext,
@@ -272,7 +295,8 @@ pub(crate) async fn record_failed_usage_for_runtime_miss_request(
return;
}
let selected_candidate = select_last_runtime_miss_candidate(&context.candidate_contexts);
let selected_candidate =
select_last_runtime_miss_executed_candidate(&context.candidate_contexts);
let api_format = selected_candidate
.and_then(|value| value.client_api_format.clone())
.or_else(|| {
@@ -371,6 +395,7 @@ pub(crate) async fn record_failed_usage_for_runtime_miss_request(
apply_runtime_miss_usage_routing(
&mut data,
&mut request_metadata,
execution_path,
selected_candidate.map(|value| value.candidate.id.as_str()),
selected_candidate.map(|value| value.candidate.candidate_index),
selected_candidate.and_then(|value| value.key_name.as_deref()),
@@ -410,20 +435,34 @@ fn select_last_failed_request_candidate(
})
}
fn select_last_runtime_miss_candidate(
fn select_last_runtime_miss_executed_candidate(
candidates: &[RuntimeMissCandidateContext],
) -> Option<&RuntimeMissCandidateContext> {
candidates.iter().max_by_key(|candidate| {
(
candidate.candidate.retry_index,
candidate.candidate.candidate_index,
candidate
.candidate
.finished_at_unix_ms
.or(candidate.candidate.started_at_unix_ms)
.unwrap_or(candidate.candidate.created_at_unix_ms),
)
})
candidates
.iter()
.filter(|candidate| request_candidate_represents_provider_execution(&candidate.candidate))
.max_by_key(|candidate| {
(
candidate.candidate.retry_index,
candidate.candidate.candidate_index,
candidate
.candidate
.finished_at_unix_ms
.or(candidate.candidate.started_at_unix_ms)
.unwrap_or(candidate.candidate.created_at_unix_ms),
)
})
}
fn request_candidate_represents_provider_execution(candidate: &StoredRequestCandidate) -> bool {
matches!(
candidate.status,
RequestCandidateStatus::Pending
| RequestCandidateStatus::Streaming
| RequestCandidateStatus::Success
| RequestCandidateStatus::Failed
| RequestCandidateStatus::Cancelled
)
}
fn error_category_for_failed_status(status_code: u16) -> Option<String> {
@@ -582,6 +621,27 @@ async fn load_runtime_miss_candidate_contexts(
.collect()
}
async fn load_runtime_miss_candidate_contexts_with_retry(
state: &AppState,
request_id: &str,
decision: Option<&GatewayControlDecision>,
) -> Vec<RuntimeMissCandidateContext> {
let mut contexts = load_runtime_miss_candidate_contexts(state, request_id, decision).await;
if !contexts.is_empty() {
return contexts;
}
for _ in 0..4 {
tokio::time::sleep(Duration::from_millis(10)).await;
contexts = load_runtime_miss_candidate_contexts(state, request_id, decision).await;
if !contexts.is_empty() {
break;
}
}
contexts
}
fn collect_present_ids<'a>(ids: impl Iterator<Item = &'a str>) -> Vec<String> {
ids.filter_map(|value| {
let trimmed = value.trim();
@@ -740,6 +800,7 @@ fn infer_endpoint_kind(api_format: &str) -> Option<&str> {
fn apply_runtime_miss_usage_routing(
data: &mut UsageEventData,
request_metadata: &mut Map<String, Value>,
execution_path: &str,
candidate_id: Option<&str>,
candidate_index: Option<u32>,
key_name: Option<&str>,
@@ -761,7 +822,7 @@ fn apply_runtime_miss_usage_routing(
data.execution_path = data
.execution_path
.clone()
.or_else(|| Some(EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS.to_string()));
.or_else(|| trimmed_non_empty(Some(execution_path)));
data.local_execution_runtime_miss_reason = data
.local_execution_runtime_miss_reason
.clone()
@@ -796,9 +857,15 @@ fn trimmed_non_empty(value: Option<&str>) -> Option<String> {
#[cfg(test)]
mod tests {
use super::apply_runtime_miss_usage_routing;
use super::{
apply_runtime_miss_usage_routing, request_candidate_represents_provider_execution,
select_last_runtime_miss_executed_candidate, RuntimeMissCandidateContext,
};
use crate::constants::EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS;
use crate::state::LocalExecutionRuntimeMissDiagnostic;
use aether_data_contracts::repository::candidates::{
RequestCandidateStatus, StoredRequestCandidate,
};
use aether_usage_runtime::UsageEventData;
use serde_json::{json, Map, Value};
@@ -811,6 +878,7 @@ mod tests {
apply_runtime_miss_usage_routing(
&mut data,
&mut request_metadata,
EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS,
Some("cand-1"),
Some(2),
Some("primary"),
@@ -846,4 +914,52 @@ mod tests {
})
);
}
#[test]
fn runtime_miss_executed_candidate_selection_ignores_skipped_only_histories() {
let skipped_candidate = StoredRequestCandidate::new(
"cand-skipped".to_string(),
"req-1".to_string(),
Some("user-1".to_string()),
Some("api-key-1".to_string()),
Some("alice".to_string()),
Some("default".to_string()),
0,
0,
Some("provider-1".to_string()),
Some("endpoint-1".to_string()),
Some("provider-key-1".to_string()),
RequestCandidateStatus::Skipped,
Some("api_key_concurrency_limit_reached".to_string()),
false,
None,
None,
None,
None,
None,
None,
None,
100,
None,
None,
)
.expect("candidate should build");
assert!(!request_candidate_represents_provider_execution(
&skipped_candidate
));
let contexts = vec![RuntimeMissCandidateContext {
candidate: skipped_candidate,
provider_name: Some("openai".to_string()),
key_name: Some("prod".to_string()),
client_api_format: Some("openai:cli".to_string()),
provider_api_format: Some("openai:cli".to_string()),
global_model_name: Some("gpt-5".to_string()),
selected_provider_model_name: Some("gpt-5-upstream".to_string()),
endpoint_url: Some("https://api.openai.example/v1/responses".to_string()),
}];
assert!(select_last_runtime_miss_executed_candidate(&contexts).is_none());
}
}

16
apps/aether-gateway/src/handlers/admin/auth/api_keys/mutation_routes.rs
View File

@@ -13,6 +13,7 @@ use crate::handlers::admin::users::{
normalize_admin_optional_api_key_name, normalize_admin_user_api_formats,
normalize_admin_user_string_list,
};
use crate::handlers::shared::normalize_optional_api_key_concurrent_limit;
use crate::GatewayError;
use aether_admin::system::serialize_admin_system_users_export_wallet;
use axum::{
@@ -114,6 +115,11 @@ pub(super) async fn build_admin_create_api_key_response(
"rate_limit 必须大于等于 0",
));
}
let concurrent_limit =
match normalize_optional_api_key_concurrent_limit(payload.concurrent_limit) {
Ok(value) => value,
Err(detail) => return Ok(build_admin_api_keys_bad_request_response(detail)),
};
let (initial_balance_usd, unlimited_balance) = match normalize_standalone_initial_balance(
payload.initial_balance_usd,
payload.unlimited_balance,
@@ -154,7 +160,7 @@ pub(super) async fn build_admin_create_api_key_response(
allowed_api_formats,
allowed_models,
rate_limit: payload.rate_limit,
concurrent_limit: 5,
concurrent_limit,
force_capabilities: None,
is_active: true,
expires_at_unix_secs,
@@ -184,6 +190,7 @@ pub(super) async fn build_admin_create_api_key_response(
"is_standalone": true,
"is_active": created.is_active,
"rate_limit": created.rate_limit,
"concurrent_limit": created.concurrent_limit,
"allowed_providers": created.allowed_providers,
"allowed_api_formats": created.allowed_api_formats,
"allowed_models": created.allowed_models,
@@ -270,6 +277,11 @@ pub(super) async fn build_admin_update_api_key_response(
"rate_limit 必须大于等于 0",
));
}
let concurrent_limit =
match normalize_optional_api_key_concurrent_limit(payload.concurrent_limit) {
Ok(value) => value,
Err(detail) => return Ok(build_admin_api_keys_bad_request_response(detail)),
};
let allowed_providers = if field_presence.contains("allowed_providers") {
match normalize_admin_user_string_list(payload.allowed_providers, "allowed_providers") {
Ok(value) => Some(value),
@@ -356,6 +368,8 @@ pub(super) async fn build_admin_update_api_key_response(
name,
rate_limit_present: field_presence.contains("rate_limit"),
rate_limit: payload.rate_limit,
concurrent_limit_present: field_presence.contains("concurrent_limit"),
concurrent_limit,
allowed_providers,
allowed_api_formats,
allowed_models,

4
apps/aether-gateway/src/handlers/admin/auth/api_keys/shared.rs
View File

@@ -22,6 +22,7 @@ pub(super) struct AdminStandaloneApiKeyCreateRequest {
pub(super) allowed_api_formats: Option<Vec<String>>,
pub(super) allowed_models: Option<Vec<String>>,
pub(super) rate_limit: Option<i32>,
pub(super) concurrent_limit: Option<i32>,
pub(super) initial_balance_usd: Option<f64>,
pub(super) unlimited_balance: Option<bool>,
pub(super) expire_days: Option<i32>,
@@ -36,6 +37,7 @@ pub(super) struct AdminStandaloneApiKeyUpdateRequest {
pub(super) allowed_api_formats: Option<Vec<String>>,
pub(super) allowed_models: Option<Vec<String>>,
pub(super) rate_limit: Option<i32>,
pub(super) concurrent_limit: Option<i32>,
pub(super) initial_balance_usd: Option<f64>,
pub(super) unlimited_balance: Option<bool>,
pub(super) expire_days: Option<i32>,
@@ -144,6 +146,7 @@ pub(super) fn build_admin_api_key_list_item_payload(
"total_tokens": total_tokens,
"total_cost_usd": record.total_cost_usd,
"rate_limit": record.rate_limit,
"concurrent_limit": record.concurrent_limit,
"allowed_providers": record.allowed_providers,
"allowed_api_formats": record.allowed_api_formats,
"allowed_models": record.allowed_models,
@@ -173,6 +176,7 @@ pub(super) fn build_admin_api_key_detail_payload(
"total_tokens": total_tokens,
"total_cost_usd": record.total_cost_usd,
"rate_limit": record.rate_limit,
"concurrent_limit": record.concurrent_limit,
"allowed_providers": record.allowed_providers,
"allowed_api_formats": record.allowed_api_formats,
"allowed_models": record.allowed_models,

20
apps/aether-gateway/src/handlers/admin/request/system/import.rs
View File

@@ -1871,8 +1871,10 @@ impl<'a> AdminAppState<'a> {
let concurrent_limit = invalid_value!(imported_optional_i32(
key.get("concurrent_limit"),
"concurrent_limit"
))
.unwrap_or(5);
));
if concurrent_limit.is_some_and(|value| value < 0) {
return Ok(Err(invalid_request("concurrent_limit 必须是非负整数")));
}
let force_capabilities = imported_optional_value(key.get("force_capabilities"));
let is_active =
invalid_value!(imported_optional_bool(key.get("is_active"))).unwrap_or(true);
@@ -1913,6 +1915,11 @@ impl<'a> AdminAppState<'a> {
api_key_id: existing_key.api_key_id.clone(),
name: name.clone(),
rate_limit: Some(rate_limit),
concurrent_limit: if key.contains_key("concurrent_limit") {
concurrent_limit
} else {
None
},
},
)
.await?;
@@ -2051,8 +2058,10 @@ impl<'a> AdminAppState<'a> {
let concurrent_limit = invalid_value!(imported_optional_i32(
key.get("concurrent_limit"),
"concurrent_limit"
))
.unwrap_or(5);
));
if concurrent_limit.is_some_and(|value| value < 0) {
return Ok(Err(invalid_request("concurrent_limit 必须是非负整数")));
}
let force_capabilities = imported_optional_value(key.get("force_capabilities"));
let is_active =
invalid_value!(imported_optional_bool(key.get("is_active"))).unwrap_or(true);
@@ -2099,6 +2108,8 @@ impl<'a> AdminAppState<'a> {
name: name.clone(),
rate_limit_present: true,
rate_limit: Some(rate_limit),
concurrent_limit_present: key.contains_key("concurrent_limit"),
concurrent_limit,
allowed_providers: Some(allowed_providers.clone()),
allowed_api_formats: Some(allowed_api_formats.clone()),
allowed_models: Some(allowed_models.clone()),
@@ -2123,7 +2134,6 @@ impl<'a> AdminAppState<'a> {
|| key.contains_key("force_capabilities")
|| key.contains_key("total_requests")
|| key.contains_key("total_cost_usd")
|| key.contains_key("concurrent_limit")
{
stats.errors.push(
"现有独立余额 Key 仅覆盖基础字段；高级导入字段保持原值".to_string(),

1
apps/aether-gateway/src/handlers/admin/users/api_keys/helpers.rs
View File

@@ -43,6 +43,7 @@ pub(super) fn build_admin_user_api_key_detail_payload(
"total_requests": record.total_requests,
"total_cost_usd": record.total_cost_usd,
"rate_limit": record.rate_limit,
"concurrent_limit": record.concurrent_limit,
"expires_at": format_optional_unix_secs_iso8601(record.expires_at_unix_secs),
"last_used_at": serde_json::Value::Null,
"created_at": serde_json::Value::Null,

17
apps/aether-gateway/src/handlers/admin/users/api_keys/responses/create.rs
View File

@@ -10,6 +10,7 @@ use super::super::helpers::{
use super::super::paths::admin_user_id_from_api_keys_path;
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::handlers::shared::normalize_optional_api_key_concurrent_limit;
use crate::GatewayError;
use axum::{
body::Body,
@@ -69,7 +70,7 @@ pub(crate) async fn build_admin_create_user_api_key_response(
{
return Ok((
http::StatusCode::BAD_REQUEST,
Json(json!({ "detail": "当前仅支持 name、rate_limit、allowed_providers 字段" })),
Json(json!({ "detail": "当前仅支持 name、rate_limit、concurrent_limit、allowed_providers 字段" })),
)
.into_response());
}
@@ -103,6 +104,17 @@ pub(crate) async fn build_admin_create_user_api_key_response(
)
.into_response());
}
let concurrent_limit =
match normalize_optional_api_key_concurrent_limit(payload.concurrent_limit) {
Ok(value) => value,
Err(detail) => {
return Ok((
http::StatusCode::BAD_REQUEST,
Json(json!({ "detail": detail })),
)
.into_response());
}
};
let plaintext_key = generate_admin_user_api_key_plaintext();
let Some(key_encrypted) = state.encrypt_catalog_secret_with_fallbacks(&plaintext_key) else {
@@ -124,7 +136,7 @@ pub(crate) async fn build_admin_create_user_api_key_response(
allowed_api_formats: None,
allowed_models: None,
rate_limit,
concurrent_limit: 5,
concurrent_limit,
force_capabilities: None,
is_active: true,
expires_at_unix_secs: None,
@@ -156,6 +168,7 @@ pub(crate) async fn build_admin_create_user_api_key_response(
"name": created.name,
"key_display": masked_user_api_key_display(state, created.key_encrypted.as_deref()),
"rate_limit": created.rate_limit,
"concurrent_limit": created.concurrent_limit,
"expires_at": format_optional_unix_secs_iso8601(created.expires_at_unix_secs),
"created_at": chrono::Utc::now().to_rfc3339(),
"message": "API Key创建成功，请妥善保存完整密钥",

1
apps/aether-gateway/src/handlers/admin/users/api_keys/responses/list.rs
View File

@@ -62,6 +62,7 @@ pub(crate) async fn build_admin_list_user_api_keys_response(
"total_requests": record.total_requests,
"total_cost_usd": record.total_cost_usd,
"rate_limit": record.rate_limit,
"concurrent_limit": record.concurrent_limit,
"expires_at": format_optional_unix_secs_iso8601(record.expires_at_unix_secs),
"last_used_at": serde_json::Value::Null,
"created_at": serde_json::Value::Null,

13
apps/aether-gateway/src/handlers/admin/users/api_keys/responses/update.rs
View File

@@ -9,6 +9,7 @@ use super::super::helpers::{
use super::super::paths::admin_user_api_key_parts;
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::handlers::shared::normalize_optional_api_key_concurrent_limit;
use crate::GatewayError;
use axum::{
body::Body,
@@ -68,6 +69,17 @@ pub(crate) async fn build_admin_update_user_api_key_response(
)
.into_response());
}
let concurrent_limit =
match normalize_optional_api_key_concurrent_limit(payload.concurrent_limit) {
Ok(value) => value,
Err(detail) => {
return Ok((
http::StatusCode::BAD_REQUEST,
Json(json!({ "detail": detail })),
)
.into_response());
}
};
let Some(updated) = state
.update_user_api_key_basic(aether_data::repository::auth::UpdateUserApiKeyBasicRecord {
@@ -75,6 +87,7 @@ pub(crate) async fn build_admin_update_user_api_key_response(
api_key_id: api_key_id.clone(),
name,
rate_limit: payload.rate_limit,
concurrent_limit,
})
.await?
else {

4
apps/aether-gateway/src/handlers/admin/users/shared.rs
View File

@@ -22,6 +22,8 @@ pub(super) struct AdminCreateUserApiKeyRequest {
#[serde(default)]
pub(super) rate_limit: Option<i32>,
#[serde(default)]
pub(super) concurrent_limit: Option<i32>,
#[serde(default)]
pub(super) expire_days: Option<i32>,
#[serde(default)]
pub(super) expires_at: Option<String>,
@@ -41,6 +43,8 @@ pub(super) struct AdminUpdateUserApiKeyRequest {
pub(super) name: Option<String>,
#[serde(default)]
pub(super) rate_limit: Option<i32>,
#[serde(default)]
pub(super) concurrent_limit: Option<i32>,
}
#[derive(Debug, serde::Deserialize)]

161
apps/aether-gateway/src/handlers/proxy/mod.rs
View File

@@ -14,14 +14,14 @@ use crate::constants::{
DEPENDENCY_REASON_HEADER, EXECUTION_PATH_CONTROL_EXECUTE_STREAM,
EXECUTION_PATH_CONTROL_EXECUTE_SYNC, EXECUTION_PATH_DISTRIBUTED_OVERLOADED,
EXECUTION_PATH_EXECUTION_RUNTIME_STREAM, EXECUTION_PATH_EXECUTION_RUNTIME_SYNC,
EXECUTION_PATH_LOCAL_AI_PUBLIC, EXECUTION_PATH_LOCAL_AUTH_DENIED,
EXECUTION_PATH_LOCAL_EXECUTION_LOOP_DETECTED, EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS,
EXECUTION_PATH_LOCAL_OVERLOADED, EXECUTION_PATH_LOCAL_PROXY_PASSTHROUGH_REMOVED,
EXECUTION_PATH_LOCAL_RATE_LIMITED, EXECUTION_PATH_LOCAL_ROUTE_NOT_FOUND,
EXECUTION_PATH_PUBLIC_PROXY_PASSTHROUGH, EXECUTION_RUNTIME_LOOP_GUARD_HEADER,
FORWARDED_FOR_HEADER, FORWARDED_HOST_HEADER, FORWARDED_PROTO_HEADER, GATEWAY_HEADER,
LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER, TRACE_ID_HEADER,
TRUSTED_AUTH_ACCESS_ALLOWED_HEADER, TRUSTED_AUTH_API_KEY_ID_HEADER,
EXECUTION_PATH_LOCAL_AI_PUBLIC, EXECUTION_PATH_LOCAL_API_KEY_CONCURRENCY_LIMITED,
EXECUTION_PATH_LOCAL_AUTH_DENIED, EXECUTION_PATH_LOCAL_EXECUTION_LOOP_DETECTED,
EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS, EXECUTION_PATH_LOCAL_OVERLOADED,
EXECUTION_PATH_LOCAL_PROXY_PASSTHROUGH_REMOVED, EXECUTION_PATH_LOCAL_RATE_LIMITED,
EXECUTION_PATH_LOCAL_ROUTE_NOT_FOUND, EXECUTION_PATH_PUBLIC_PROXY_PASSTHROUGH,
EXECUTION_RUNTIME_LOOP_GUARD_HEADER, FORWARDED_FOR_HEADER, FORWARDED_HOST_HEADER,
FORWARDED_PROTO_HEADER, GATEWAY_HEADER, LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER,
TRACE_ID_HEADER, TRUSTED_AUTH_ACCESS_ALLOWED_HEADER, TRUSTED_AUTH_API_KEY_ID_HEADER,
TRUSTED_AUTH_BALANCE_HEADER, TRUSTED_AUTH_USER_ID_HEADER, TUNNEL_AFFINITY_FORWARDED_BY_HEADER,
TUNNEL_AFFINITY_OWNER_INSTANCE_HEADER,
};
@@ -75,6 +75,8 @@ const LOCAL_PROXY_PASSTHROUGH_REMOVED_DETAIL: &str =
"Route matched a removed compatibility passthrough; implement it in Rust or retire the route";
const LOCAL_EXECUTION_LOOP_DETECTED_DETAIL: &str =
"Gateway detected an execution runtime request loop back into the local frontdoor";
const AUTH_API_KEY_CONCURRENCY_LIMIT_REACHED_DETAIL: &str =
"当前 API Key 并发请求数已达上限，请稍后重试";
const EXECUTION_PATH_TUNNEL_AFFINITY_FORWARD: &str = "tunnel_affinity_forward";
fn local_execution_outcome_label(outcome: &LocalExecutionRequestOutcome) -> &'static str {
@@ -1018,23 +1020,38 @@ pub(crate) async fn proxy_request(
}
let local_execution_runtime_miss_diagnostic =
state.take_local_execution_runtime_miss_diagnostic(&trace_id);
let local_execution_runtime_miss_context =
build_local_execution_runtime_miss_context(&state, &trace_id, control_decision).await;
let auth_api_key_concurrency_limited = diagnostic_is_auth_api_key_concurrency_limited(
local_execution_runtime_miss_diagnostic.as_ref(),
) || local_execution_runtime_miss_context
.all_candidates_skipped_for_reason("api_key_concurrency_limit_reached");
let local_execution_runtime_miss_detail = local_execution_runtime_miss_detail(
control_decision,
local_execution_runtime_miss_diagnostic.as_ref(),
auth_api_key_concurrency_limited,
stream_request,
)
.unwrap_or_else(|| {
"AI public execution runtime miss did not match a Rust execution path".to_string()
});
let local_execution_failure_path = if auth_api_key_concurrency_limited {
EXECUTION_PATH_LOCAL_API_KEY_CONCURRENCY_LIMITED
} else {
EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS
};
let local_execution_failure_log = if auth_api_key_concurrency_limited {
"gateway local execution blocked by api key concurrency limit"
} else {
"gateway local execution runtime miss"
};
state.record_fallback_metric(
GatewayFallbackMetricKind::LocalExecutionRuntimeMiss,
control_decision,
None,
Some(EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS),
Some(local_execution_failure_path),
GatewayFallbackReason::LocalExecutionPathRequired,
);
let local_execution_runtime_miss_context =
build_local_execution_runtime_miss_context(&state, &trace_id, control_decision).await;
warn!(
trace_id = %trace_id,
local_execution_runtime_miss_reason = local_execution_runtime_miss_diagnostic
@@ -1094,7 +1111,7 @@ pub(crate) async fn proxy_request(
request_candidates = local_execution_runtime_miss_context
.candidate_summary()
.unwrap_or_default(),
"gateway local execution runtime miss"
local_execution_failure_log
);
if let Some(exhaustion) = local_execution_exhaustion {
record_failed_usage_for_exhausted_request(
@@ -1102,6 +1119,7 @@ pub(crate) async fn proxy_request(
exhaustion,
&started_at,
local_execution_runtime_miss_detail.as_str(),
local_execution_failure_path,
local_execution_runtime_miss_diagnostic.as_ref(),
)
.await;
@@ -1111,6 +1129,7 @@ pub(crate) async fn proxy_request(
&trace_id,
&started_at,
local_execution_runtime_miss_detail.as_str(),
local_execution_failure_path,
control_decision,
local_execution_runtime_miss_diagnostic.as_ref(),
&local_execution_runtime_miss_context,
@@ -1123,21 +1142,28 @@ pub(crate) async fn proxy_request(
http::StatusCode::SERVICE_UNAVAILABLE,
local_execution_runtime_miss_detail.as_str(),
)?;
if let Some(diagnostic) = local_execution_runtime_miss_diagnostic {
if !diagnostic.reason.trim().is_empty() {
response.headers_mut().insert(
HeaderName::from_static(LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER),
HeaderValue::from_str(diagnostic.reason.as_str())
.map_err(|err| GatewayError::Internal(err.to_string()))?,
);
}
let local_execution_runtime_miss_reason = local_execution_runtime_miss_diagnostic
.as_ref()
.map(|diagnostic| diagnostic.reason.trim())
.filter(|reason| !reason.is_empty())
.map(ToOwned::to_owned)
.or_else(|| {
auth_api_key_concurrency_limited
.then_some("api_key_concurrency_limit_reached".to_string())
});
if let Some(reason) = local_execution_runtime_miss_reason {
response.headers_mut().insert(
HeaderName::from_static(LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER),
HeaderValue::from_str(reason.as_str())
.map_err(|err| GatewayError::Internal(err.to_string()))?,
);
}
return Ok(finalize_gateway_response_with_context(
&state,
response,
&remote_addr,
&request_context,
EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS,
local_execution_failure_path,
&started_at,
request_permit.take(),
));
@@ -1163,8 +1189,15 @@ pub(crate) async fn proxy_request(
fn local_execution_runtime_miss_detail(
decision: Option<&GatewayControlDecision>,
diagnostic: Option<&LocalExecutionRuntimeMissDiagnostic>,
auth_api_key_concurrency_limited: bool,
stream_request: bool,
) -> Option<String> {
if auth_api_key_concurrency_limited
|| diagnostic_is_auth_api_key_concurrency_limited(diagnostic)
{
return Some(AUTH_API_KEY_CONCURRENCY_LIMIT_REACHED_DETAIL.to_string());
}
if let Some(detail) = local_execution_runtime_miss_model_detail(diagnostic, stream_request) {
return Some(detail);
}
@@ -1195,6 +1228,23 @@ fn local_execution_runtime_miss_model_detail(
))
}
fn diagnostic_is_auth_api_key_concurrency_limited(
diagnostic: Option<&LocalExecutionRuntimeMissDiagnostic>,
) -> bool {
let Some(diagnostic) = diagnostic else {
return false;
};
diagnostic.reason == "api_key_concurrency_limit_reached"
|| (diagnostic.reason == "all_candidates_skipped"
&& diagnostic.skip_reasons.len() == 1
&& diagnostic
.skip_reasons
.get("api_key_concurrency_limit_reached")
.copied()
.unwrap_or(0)
> 0)
}
fn local_execution_runtime_miss_route_detail(
decision: Option<&GatewayControlDecision>,
) -> Option<&'static str> {
@@ -1226,8 +1276,8 @@ fn local_execution_runtime_miss_route_detail(
#[cfg(test)]
mod tests {
use super::{
local_execution_runtime_miss_detail, GatewayControlDecision,
LocalExecutionRuntimeMissDiagnostic,
diagnostic_is_auth_api_key_concurrency_limited, local_execution_runtime_miss_detail,
GatewayControlDecision, LocalExecutionRuntimeMissDiagnostic,
};
#[test]
@@ -1245,7 +1295,8 @@ mod tests {
..LocalExecutionRuntimeMissDiagnostic::default()
};
let detail = local_execution_runtime_miss_detail(Some(&decision), Some(&diagnostic), true);
let detail =
local_execution_runtime_miss_detail(Some(&decision), Some(&diagnostic), false, true);
assert_eq!(
detail.as_deref(),
@@ -1268,13 +1319,73 @@ mod tests {
..LocalExecutionRuntimeMissDiagnostic::default()
};
let detail = local_execution_runtime_miss_detail(Some(&decision), Some(&diagnostic), false);
let detail =
local_execution_runtime_miss_detail(Some(&decision), Some(&diagnostic), false, false);
assert_eq!(
detail.as_deref(),
Some("Claude messages execution runtime miss did not match a Rust execution path")
);
}
#[test]
fn runtime_miss_detail_returns_api_key_concurrency_message_for_exact_all_skipped_limit_case() {
let decision = GatewayControlDecision::synthetic(
"/v1/responses",
Some("ai_public".to_string()),
Some("openai".to_string()),
Some("cli".to_string()),
Some("openai:cli".to_string()),
);
let diagnostic = LocalExecutionRuntimeMissDiagnostic {
reason: "all_candidates_skipped".to_string(),
skip_reasons: std::collections::BTreeMap::from([(
"api_key_concurrency_limit_reached".to_string(),
1,
)]),
requested_model: Some("gpt-5.4".to_string()),
..LocalExecutionRuntimeMissDiagnostic::default()
};
let detail =
local_execution_runtime_miss_detail(Some(&decision), Some(&diagnostic), false, false);
assert_eq!(
detail.as_deref(),
Some("当前 API Key 并发请求数已达上限，请稍后重试")
);
assert!(diagnostic_is_auth_api_key_concurrency_limited(Some(
&diagnostic
)));
}
#[test]
fn runtime_miss_detail_prefers_api_key_concurrency_message_when_classified_from_context() {
let decision = GatewayControlDecision::synthetic(
"/v1/responses",
Some("ai_public".to_string()),
Some("openai".to_string()),
Some("cli".to_string()),
Some("openai:cli".to_string()),
);
let diagnostic = LocalExecutionRuntimeMissDiagnostic {
reason: "all_candidates_skipped".to_string(),
skip_reasons: std::collections::BTreeMap::from([(
"format_conversion_disabled".to_string(),
1,
)]),
requested_model: Some("gpt-5.4".to_string()),
..LocalExecutionRuntimeMissDiagnostic::default()
};
let detail =
local_execution_runtime_miss_detail(Some(&decision), Some(&diagnostic), true, false);
assert_eq!(
detail.as_deref(),
Some("当前 API Key 并发请求数已达上限，请稍后重试")
);
}
}
#[path = "finalize.rs"]

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

@@ -10,7 +10,8 @@ use serde::Deserialize;
use serde_json::json;
use crate::handlers::shared::{
api_key_placeholder_display, generate_gateway_api_key_plaintext, masked_gateway_api_key_display,
api_key_placeholder_display, generate_gateway_api_key_plaintext,
masked_gateway_api_key_display, normalize_optional_api_key_concurrent_limit,
};
use super::{
@@ -28,6 +29,8 @@ struct UsersMeCreateApiKeyRequest {
name: String,
#[serde(default)]
rate_limit: Option<i32>,
#[serde(default)]
concurrent_limit: Option<i32>,
}
#[derive(Debug, Deserialize)]
@@ -36,6 +39,8 @@ struct UsersMeUpdateApiKeyRequest {
name: Option<String>,
#[serde(default)]
rate_limit: Option<i32>,
#[serde(default)]
concurrent_limit: Option<i32>,
}
#[derive(Debug, Deserialize)]
@@ -148,6 +153,7 @@ fn build_users_me_api_key_list_payload(
"total_requests": record.total_requests,
"total_cost_usd": record.total_cost_usd,
"rate_limit": record.rate_limit,
"concurrent_limit": record.concurrent_limit,
"allowed_providers": record.allowed_providers,
"force_capabilities": record.force_capabilities,
})
@@ -167,6 +173,7 @@ fn build_users_me_api_key_detail_payload(
"allowed_providers": record.allowed_providers,
"force_capabilities": record.force_capabilities,
"rate_limit": record.rate_limit,
"concurrent_limit": record.concurrent_limit,
"last_used_at": serde_json::Value::Null,
"expires_at": format_users_me_optional_unix_secs_iso8601(record.expires_at_unix_secs),
"created_at": serde_json::Value::Null,
@@ -515,6 +522,13 @@ pub(super) async fn handle_users_me_api_key_create(
false,
);
}
let concurrent_limit =
match normalize_optional_api_key_concurrent_limit(payload.concurrent_limit) {
Ok(value) => value,
Err(detail) => {
return build_auth_error_response(http::StatusCode::BAD_REQUEST, detail, false);
}
};
let plaintext_key = generate_users_me_api_key_plaintext();
let Some(key_encrypted) = encrypt_catalog_secret_with_fallbacks(state, &plaintext_key) else {
@@ -534,7 +548,7 @@ pub(super) async fn handle_users_me_api_key_create(
allowed_api_formats: None,
allowed_models: None,
rate_limit,
concurrent_limit: 5,
concurrent_limit,
force_capabilities: None,
is_active: true,
expires_at_unix_secs: None,
@@ -561,6 +575,7 @@ pub(super) async fn handle_users_me_api_key_create(
"key": plaintext_key,
"key_display": users_me_masked_api_key_display(state, created.key_encrypted.as_deref()),
"rate_limit": created.rate_limit,
"concurrent_limit": created.concurrent_limit,
"message": "API密钥创建成功",
}))
.into_response()
@@ -621,6 +636,13 @@ pub(super) async fn handle_users_me_api_key_update(
false,
);
}
let concurrent_limit =
match normalize_optional_api_key_concurrent_limit(payload.concurrent_limit) {
Ok(value) => value,
Err(detail) => {
return build_auth_error_response(http::StatusCode::BAD_REQUEST, detail, false);
}
};
let Some(updated) = (match state
.update_user_api_key_basic(aether_data::repository::auth::UpdateUserApiKeyBasicRecord {
@@ -628,6 +650,7 @@ pub(super) async fn handle_users_me_api_key_update(
api_key_id: snapshot.api_key_id.clone(),
name,
rate_limit,
concurrent_limit,
})
.await
{

9
apps/aether-gateway/src/handlers/shared/api_keys.rs
View File

@@ -60,6 +60,15 @@ pub(crate) fn masked_gateway_api_key_display(full_key: Option<&str>) -> String {
format!("{prefix}...{suffix}")
}
pub(crate) fn normalize_optional_api_key_concurrent_limit(
value: Option<i32>,
) -> Result<Option<i32>, String> {
if value.is_some_and(|limit| limit < 0) {
return Err("concurrent_limit 必须是非负整数".to_string());
}
Ok(value)
}
#[cfg(test)]
mod tests {
use super::{

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

@@ -15,7 +15,7 @@ pub(crate) use self::admin_proxy::{
};
pub(crate) use self::api_keys::{
api_key_placeholder_display, configured_api_key_prefix, generate_gateway_api_key_plaintext,
masked_gateway_api_key_display,
masked_gateway_api_key_display, normalize_optional_api_key_concurrent_limit,
};
pub(crate) use self::catalog::{
build_admin_provider_key_response, decrypt_catalog_secret_with_fallbacks,

14
apps/aether-gateway/src/middleware/access_log.rs
View File

@@ -212,7 +212,7 @@ mod tests {
}
}
#[tokio::test]
#[tokio::test(flavor = "current_thread")]
async fn access_log_emits_completed_events_by_default() {
let writer = SharedBuffer::default();
let subscriber = tracing_subscriber::registry().with(
@@ -271,7 +271,7 @@ mod tests {
assert_eq!(logs[0]["execution_path"], "local_route");
}
#[tokio::test]
#[tokio::test(flavor = "current_thread")]
async fn access_log_propagates_generated_trace_id_to_downstream_handler() {
let app = Router::new()
.route(
@@ -318,7 +318,7 @@ mod tests {
assert_eq!(seen_trace_id, response_trace_id);
}
#[tokio::test]
#[tokio::test(flavor = "current_thread")]
async fn access_log_shortens_long_request_ids() {
let writer = SharedBuffer::default();
let subscriber = tracing_subscriber::registry().with(
@@ -371,7 +371,7 @@ mod tests {
assert_eq!(logs[0]["request_id"], "d07e1e94");
}
#[tokio::test]
#[tokio::test(flavor = "current_thread")]
async fn access_log_emits_failed_events_by_default_for_server_errors() {
let writer = SharedBuffer::default();
let subscriber = tracing_subscriber::registry().with(
@@ -419,7 +419,7 @@ mod tests {
assert_eq!(logs[0]["execution_path"], "execution_runtime_sync");
}
#[tokio::test]
#[tokio::test(flavor = "current_thread")]
async fn access_log_treats_client_errors_as_completed_events() {
let writer = SharedBuffer::default();
let subscriber = tracing_subscriber::registry().with(
@@ -467,7 +467,7 @@ mod tests {
assert_eq!(logs[0]["execution_path"], "local_auth_denied");
}
#[tokio::test]
#[tokio::test(flavor = "current_thread")]
async fn access_log_emits_completed_events_for_streaming_responses() {
let writer = SharedBuffer::default();
let subscriber = tracing_subscriber::registry().with(
@@ -522,7 +522,7 @@ mod tests {
assert_eq!(logs[0]["execution_path"], "execution_runtime_stream");
}
#[tokio::test]
#[tokio::test(flavor = "current_thread")]
async fn access_log_downgrades_usage_active_polling_to_trace() {
let writer = SharedBuffer::default();
let subscriber = tracing_subscriber::registry().with(

44
apps/aether-gateway/src/scheduler/candidate/mod.rs
View File

@@ -69,6 +69,13 @@ pub(crate) async fn list_selectable_candidates(
.await
}
pub(crate) fn is_exact_all_skipped_by_auth_limit(
selected: &[SchedulerMinimalCandidateSelectionCandidate],
skipped: &[SchedulerSkippedCandidate],
) -> bool {
selection::is_exact_all_skipped_by_auth_limit(selected, skipped)
}
pub(crate) async fn list_selectable_candidates_with_skip_reasons(
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
runtime_state: &impl SchedulerRuntimeState,
@@ -107,9 +114,33 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
now_unix_secs: u64,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
Ok(
list_selectable_candidates_for_required_capability_without_requested_model_with_auth_limit_signal(
selection_row_source,
runtime_state,
candidate_api_format,
required_capability,
require_streaming,
auth_snapshot,
now_unix_secs,
)
.await?
.0,
)
}
pub(crate) async fn list_selectable_candidates_for_required_capability_without_requested_model_with_auth_limit_signal(
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
runtime_state: &impl SchedulerRuntimeState,
candidate_api_format: &str,
required_capability: &str,
require_streaming: bool,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
now_unix_secs: u64,
) -> Result<(Vec<SchedulerMinimalCandidateSelectionCandidate>, bool), GatewayError> {
let normalized_api_format = normalize_api_format(candidate_api_format);
if normalized_api_format.is_empty() {
return Ok(Vec::new());
return Ok((Vec::new(), false));
}
let capability_mode = required_capability_match_mode(required_capability);
@@ -136,9 +167,10 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
}
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let required_capabilities = build_required_capabilities_object(required_capability);
let mut all_attempts_blocked_by_auth_limit = !model_names.is_empty();
for global_model_name in model_names {
let mut candidates = list_selectable_candidates(
let (mut candidates, skipped_candidates) = collect_selectable_candidates_with_skip_reasons(
selection_row_source,
runtime_state,
&normalized_api_format,
@@ -149,6 +181,8 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
now_unix_secs,
)
.await?;
all_attempts_blocked_by_auth_limit &=
is_exact_all_skipped_by_auth_limit(&candidates, &skipped_candidates);
match capability_mode {
RequiredCapabilityMatchMode::Exclusive => {
let filtered = candidates
@@ -158,7 +192,7 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
})
.collect::<Vec<_>>();
if !filtered.is_empty() {
return Ok(filtered);
return Ok((filtered, false));
}
}
RequiredCapabilityMatchMode::Compatible => {
@@ -168,12 +202,12 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
candidates.sort_by_key(|candidate| {
!candidate_supports_required_capability(candidate, required_capability)
});
return Ok(candidates);
return Ok((candidates, false));
}
}
}
Ok(Vec::new())
Ok((Vec::new(), all_attempts_blocked_by_auth_limit))
}
fn required_capability_match_mode(required_capability: &str) -> RequiredCapabilityMatchMode {

15
apps/aether-gateway/src/scheduler/candidate/selection.rs
View File

@@ -31,6 +31,19 @@ pub(crate) struct SchedulerSkippedCandidate {
pub(crate) skip_reason: &'static str,
}
pub(super) const API_KEY_CONCURRENCY_LIMIT_SKIP_REASON: &str = "api_key_concurrency_limit_reached";
pub(super) fn is_exact_all_skipped_by_auth_limit(
selected: &[SchedulerMinimalCandidateSelectionCandidate],
skipped: &[SchedulerSkippedCandidate],
) -> bool {
selected.is_empty()
&& !skipped.is_empty()
&& skipped
.iter()
.all(|candidate| candidate.skip_reason == API_KEY_CONCURRENCY_LIMIT_SKIP_REASON)
}
pub(super) fn reorder_candidates_by_scheduler_health(
candidates: &mut [SchedulerMinimalCandidateSelectionCandidate],
provider_key_rpm_states: &BTreeMap<String, StoredProviderCatalogKey>,
@@ -218,7 +231,7 @@ pub(super) async fn collect_selectable_candidates_with_skip_reasons(
.into_iter()
.map(|candidate| SchedulerSkippedCandidate {
candidate,
skip_reason: "api_key_concurrency_limit_reached",
skip_reason: API_KEY_CONCURRENCY_LIMIT_SKIP_REASON,
})
.collect(),
));

87
apps/aether-gateway/src/scheduler/candidate/tests/required_capability.rs
View File

@@ -1,13 +1,21 @@
use std::sync::Arc;
use aether_data::repository::candidate_selection::InMemoryMinimalCandidateSelectionReadRepository;
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data::repository::quota::InMemoryProviderQuotaRepository;
use aether_data_contracts::repository::candidates::{
RequestCandidateStatus, StoredRequestCandidate,
};
use crate::data::GatewayDataState;
use crate::AppState;
use super::super::list_selectable_candidates_for_required_capability_without_requested_model;
use super::support::sample_row;
use super::super::{
list_selectable_candidates_for_required_capability_without_requested_model,
list_selectable_candidates_for_required_capability_without_requested_model_with_auth_limit_signal,
};
use super::support::{sample_auth_snapshot, sample_provider, sample_row};
#[tokio::test]
async fn compatible_required_capability_prefers_matching_keys_without_hard_filtering() {
@@ -114,3 +122,78 @@ async fn exclusive_required_capability_keeps_hard_filtering_only_matching_keys()
assert_eq!(selection[0].provider_id, "provider-b");
assert_eq!(selection[0].key_id, "key-b");
}
#[tokio::test]
async fn required_capability_reports_auth_limit_signal_when_every_model_is_blocked_by_api_key_concurrency(
) {
let mut candidate = sample_row();
candidate.key_capabilities = Some(serde_json::json!({"cache_1h": true}));
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
candidate,
]));
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-1", None)],
Vec::new(),
Vec::new(),
));
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
StoredRequestCandidate::new(
"cand-1".to_string(),
"req-1".to_string(),
Some("user-1".to_string()),
Some("api-key-1".to_string()),
None,
None,
0,
0,
Some("provider-1".to_string()),
Some("endpoint-1".to_string()),
Some("key-1".to_string()),
RequestCandidateStatus::Pending,
None,
false,
None,
None,
None,
None,
None,
None,
None,
95_000,
Some(95_000),
None,
)
.expect("candidate should build"),
]));
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(
GatewayDataState::with_candidate_selection_provider_catalog_quota_and_request_candidates_for_tests(
candidates,
provider_catalog,
quotas,
request_candidates,
),
);
let mut auth_snapshot = sample_auth_snapshot("api-key-1");
auth_snapshot.api_key_concurrent_limit = Some(1);
let (selection, auth_limit_blocked) =
list_selectable_candidates_for_required_capability_without_requested_model_with_auth_limit_signal(
state.data.as_ref(),
&state,
"openai:chat",
"cache_1h",
false,
Some(&auth_snapshot),
100,
)
.await
.expect("selection should succeed");
assert!(selection.is_empty());
assert!(auth_limit_blocked);
}

20
apps/aether-gateway/src/scheduler/candidate/tests/selection.rs
View File

@@ -23,7 +23,7 @@ use super::super::runtime::should_skip_provider_quota;
use super::super::selection::{
collect_selectable_candidates as collect_selectable_candidates_impl,
collect_selectable_candidates_with_skip_reasons as collect_selectable_candidates_with_skip_reasons_impl,
select_minimal_candidate as select_candidate_impl,
is_exact_all_skipped_by_auth_limit, select_minimal_candidate as select_candidate_impl,
};
use super::support::{sample_auth_snapshot, sample_key, sample_provider, sample_row};
@@ -914,6 +914,23 @@ async fn returns_none_when_auth_api_key_concurrent_limit_is_reached() {
.expect("selection should succeed");
assert!(selected.is_none());
let (selected_candidates, skipped_candidates) =
collect_selectable_candidates_with_skip_reasons(
state.data.as_ref(),
&state,
"openai:chat",
"gpt-4.1",
false,
Some(&auth_snapshot),
100,
)
.await
.expect("selection should succeed");
assert!(is_exact_all_skipped_by_auth_limit(
&selected_candidates,
&skipped_candidates,
));
}
#[tokio::test]
@@ -1138,6 +1155,7 @@ async fn exposes_runtime_skipped_candidates_with_skip_reasons() {
assert_eq!(skipped.len(), 1);
assert_eq!(skipped[0].candidate.provider_id, "provider-a");
assert_eq!(skipped[0].skip_reason, "key_circuit_open");
assert!(!is_exact_all_skipped_by_auth_limit(&selected, &skipped));
}
#[tokio::test]

736
apps/aether-gateway/src/tests/ai_execute/sync/cli.rs
View File

@@ -3,6 +3,9 @@ use super::{
to_bytes, Arc, Body, Json, Mutex, Request, Router, StatusCode,
EXECUTION_PATH_EXECUTION_RUNTIME_SYNC, EXECUTION_PATH_HEADER, TRACE_ID_HEADER,
};
use crate::constants::{
EXECUTION_PATH_LOCAL_API_KEY_CONCURRENCY_LIMITED, LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER,
};
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
use aether_data::repository::auth::{
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot,
@@ -14,7 +17,8 @@ use aether_data_contracts::repository::candidate_selection::{
StoredMinimalCandidateSelectionRow, StoredProviderModelMapping,
};
use aether_data_contracts::repository::candidates::{
RequestCandidateReadRepository, RequestCandidateStatus,
RequestCandidateReadRepository, RequestCandidateStatus, RequestCandidateWriteRepository,
UpsertRequestCandidateRecord,
};
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
@@ -535,6 +539,736 @@ async fn gateway_executes_openai_cli_sync_via_local_decision_gate_with_local_syn
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_waits_for_api_key_concurrency_slot_then_executes_openai_cli_sync() {
fn hash_api_key(value: &str) -> String {
let mut hasher = Sha256::new();
hasher.update(value.as_bytes());
format!("{:x}", hasher.finalize())
}
fn sample_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
StoredAuthApiKeySnapshot::new(
user_id.to_string(),
"alice".to_string(),
Some("alice@example.com".to_string()),
"user".to_string(),
"local".to_string(),
true,
false,
Some(serde_json::json!(["openai"])),
Some(serde_json::json!(["openai:cli"])),
Some(serde_json::json!(["gpt-5"])),
api_key_id.to_string(),
Some("default".to_string()),
true,
false,
false,
Some(60),
Some(5),
Some(4_102_444_800),
Some(serde_json::json!(["openai"])),
Some(serde_json::json!(["openai:cli"])),
Some(serde_json::json!(["gpt-5"])),
)
.expect("auth snapshot should build")
}
fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: "provider-openai-cli-local-limit-1".to_string(),
provider_name: "openai".to_string(),
provider_type: "custom".to_string(),
provider_priority: 10,
provider_is_active: true,
endpoint_id: "endpoint-openai-cli-local-limit-1".to_string(),
endpoint_api_format: "openai:cli".to_string(),
endpoint_api_family: Some("openai".to_string()),
endpoint_kind: Some("cli".to_string()),
endpoint_is_active: true,
key_id: "key-openai-cli-local-limit-1".to_string(),
key_name: "prod".to_string(),
key_auth_type: "api_key".to_string(),
key_is_active: true,
key_api_formats: Some(vec!["openai:cli".to_string()]),
key_allowed_models: None,
key_capabilities: None,
key_internal_priority: 5,
key_global_priority_by_format: Some(serde_json::json!({"openai:cli": 1})),
model_id: "model-openai-cli-local-limit-1".to_string(),
global_model_id: "global-model-openai-cli-local-limit-1".to_string(),
global_model_name: "gpt-5".to_string(),
global_model_mappings: None,
global_model_supports_streaming: Some(true),
model_provider_model_name: "gpt-5-upstream".to_string(),
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
name: "gpt-5-upstream".to_string(),
priority: 1,
api_formats: Some(vec!["openai:cli".to_string()]),
}]),
model_supports_streaming: Some(true),
model_is_active: true,
model_is_available: true,
}
}
fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
"provider-openai-cli-local-limit-1".to_string(),
"openai".to_string(),
Some("https://example.com".to_string()),
"custom".to_string(),
)
.expect("provider should build")
.with_transport_fields(
true,
false,
false,
None,
Some(2),
None,
Some(20.0),
None,
None,
)
}
fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
"endpoint-openai-cli-local-limit-1".to_string(),
"provider-openai-cli-local-limit-1".to_string(),
"openai:cli".to_string(),
Some("openai".to_string()),
Some("cli".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
"https://api.openai.example/custom/v1/responses".to_string(),
None,
None,
Some(2),
Some("/custom/v1/responses".to_string()),
None,
None,
None,
)
.expect("endpoint transport should build")
}
fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
StoredProviderCatalogKey::new(
"key-openai-cli-local-limit-1".to_string(),
"provider-openai-cli-local-limit-1".to_string(),
"prod".to_string(),
"api_key".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(serde_json::json!(["openai:cli"])),
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
"sk-upstream-openai-cli-limit",
)
.expect("api key should encrypt"),
None,
None,
Some(serde_json::json!({"openai:cli": 1})),
None,
None,
None,
None,
)
.expect("key transport should build")
}
let execution_runtime_hits = Arc::new(Mutex::new(0usize));
let execution_runtime_hits_clone = Arc::clone(&execution_runtime_hits);
let public_hits = Arc::new(Mutex::new(0usize));
let public_hits_clone = Arc::clone(&public_hits);
let now_unix_ms = chrono::Utc::now().timestamp_millis().max(0);
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
aether_data_contracts::repository::candidates::StoredRequestCandidate::new(
"cand-pending-openai-cli-local-limit-1".to_string(),
"req-inflight-openai-cli-local-limit-1".to_string(),
Some("user-openai-cli-local-limit-123".to_string()),
Some("key-openai-cli-local-limit-123".to_string()),
Some("alice".to_string()),
Some("default".to_string()),
0,
0,
Some("provider-openai-cli-local-limit-1".to_string()),
Some("endpoint-openai-cli-local-limit-1".to_string()),
Some("key-openai-cli-local-limit-1".to_string()),
RequestCandidateStatus::Pending,
None,
false,
None,
None,
None,
None,
None,
None,
None,
now_unix_ms,
Some(now_unix_ms),
None,
)
.expect("pending candidate should build"),
]));
let upstream = Router::new()
.route(
"/api/internal/gateway/resolve",
any(|_request: Request| async move {
Json(json!({
"action": "proxy_public",
"route_class": "ai_public",
"route_family": "openai",
"route_kind": "cli",
"auth_endpoint_signature": "openai:cli",
"execution_runtime_candidate": true,
"auth_context": {
"user_id": "user-openai-cli-local-limit-123",
"api_key_id": "key-openai-cli-local-limit-123",
"access_allowed": true
},
"public_path": "/v1/responses"
}))
}),
)
.route(
"/api/internal/gateway/decision-sync",
any(|_request: Request| async move { Json(json!({"action": "proxy_public"})) }),
)
.route(
"/api/internal/gateway/plan-sync",
any(|_request: Request| async move { Json(json!({"action": "proxy_public"})) }),
)
.route(
"/v1/responses",
any(move |_request: Request| {
let public_hits_inner = Arc::clone(&public_hits_clone);
async move {
*public_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::IM_A_TEAPOT, Body::from("public-route-hit"))
}
}),
);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |_request: Request| {
let execution_runtime_hits_inner = Arc::clone(&execution_runtime_hits_clone);
async move {
*execution_runtime_hits_inner
.lock()
.expect("mutex should lock") += 1;
Json(json!({
"request_id": "trace-openai-cli-local-limit-123",
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"id": "resp-cli-local-limit-123",
"object": "response",
"model": "gpt-5-upstream",
"output": [],
"usage": {
"input_tokens": 1,
"output_tokens": 2,
"total_tokens": 3
}
}
},
"telemetry": {
"elapsed_ms": 21
}
}))
}
}),
);
let mut auth_snapshot = sample_auth_snapshot(
"key-openai-cli-local-limit-123",
"user-openai-cli-local-limit-123",
);
auth_snapshot.api_key_concurrent_limit = Some(1);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-client-openai-cli-local-limit")),
auth_snapshot,
)]));
let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
sample_candidate_row(),
]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider_catalog_provider()],
vec![sample_provider_catalog_endpoint()],
vec![sample_provider_catalog_key()],
));
let (_upstream_url, upstream_handle) = start_server(upstream).await;
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
auth_repository,
candidate_selection_repository,
provider_catalog_repository,
Arc::clone(&request_candidate_repository),
DEVELOPMENT_ENCRYPTION_KEY,
),
);
let gateway = build_router_with_state(gateway_state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let request_candidate_repository_for_release = Arc::clone(&request_candidate_repository);
let release_inflight_candidate = tokio::spawn(async move {
tokio::time::sleep(std::time::Duration::from_millis(25)).await;
let pending = request_candidate_repository_for_release
.list_by_request_id("req-inflight-openai-cli-local-limit-1")
.await
.expect("inflight candidates should read")
.into_iter()
.find(|candidate| candidate.id == "cand-pending-openai-cli-local-limit-1")
.expect("seeded inflight candidate should exist");
request_candidate_repository_for_release
.upsert(UpsertRequestCandidateRecord {
id: pending.id,
request_id: pending.request_id,
user_id: pending.user_id,
api_key_id: pending.api_key_id,
username: pending.username,
api_key_name: pending.api_key_name,
candidate_index: pending.candidate_index,
retry_index: pending.retry_index,
provider_id: pending.provider_id,
endpoint_id: pending.endpoint_id,
key_id: pending.key_id,
status: RequestCandidateStatus::Success,
skip_reason: None,
is_cached: Some(false),
status_code: Some(200),
error_type: None,
error_message: None,
latency_ms: Some(1),
concurrent_requests: pending.concurrent_requests,
extra_data: pending.extra_data,
required_capabilities: pending.required_capabilities,
created_at_unix_ms: Some(pending.created_at_unix_ms),
started_at_unix_ms: pending.started_at_unix_ms,
finished_at_unix_ms: Some(pending.created_at_unix_ms.saturating_add(25)),
})
.await
.expect("inflight candidate should update");
});
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/responses"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(
http::header::AUTHORIZATION,
"Bearer sk-client-openai-cli-local-limit",
)
.header(TRACE_ID_HEADER, "trace-openai-cli-local-limit-123")
.body("{\"model\":\"gpt-5\",\"input\":\"hello\",\"store\":false}")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response
.headers()
.get(EXECUTION_PATH_HEADER)
.and_then(|value| value.to_str().ok()),
Some(EXECUTION_PATH_EXECUTION_RUNTIME_SYNC)
);
assert_eq!(
response
.headers()
.get(LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER)
.and_then(|value| value.to_str().ok()),
None
);
let payload: serde_json::Value = response.json().await.expect("body should parse");
assert_eq!(payload["model"], "gpt-5-upstream");
let stored_candidates = request_candidate_repository
.list_by_request_id("trace-openai-cli-local-limit-123")
.await
.expect("request candidate trace should read");
assert_eq!(stored_candidates.len(), 1);
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success);
assert_eq!(stored_candidates[0].skip_reason.as_deref(), None);
assert_eq!(
*execution_runtime_hits.lock().expect("mutex should lock"),
1
);
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
release_inflight_candidate
.await
.expect("release task should complete");
gateway_handle.abort();
execution_runtime_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_returns_concurrency_limited_after_wait_budget_expires_for_openai_cli_sync() {
fn hash_api_key(value: &str) -> String {
let mut hasher = Sha256::new();
hasher.update(value.as_bytes());
format!("{:x}", hasher.finalize())
}
fn sample_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
StoredAuthApiKeySnapshot::new(
user_id.to_string(),
"alice".to_string(),
Some("alice@example.com".to_string()),
"user".to_string(),
"local".to_string(),
true,
false,
Some(serde_json::json!(["openai"])),
Some(serde_json::json!(["openai:cli"])),
Some(serde_json::json!(["gpt-5"])),
api_key_id.to_string(),
Some("default".to_string()),
true,
false,
false,
Some(60),
Some(5),
Some(4_102_444_800),
Some(serde_json::json!(["openai"])),
Some(serde_json::json!(["openai:cli"])),
Some(serde_json::json!(["gpt-5"])),
)
.expect("auth snapshot should build")
}
fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: "provider-openai-cli-local-timeout-1".to_string(),
provider_name: "openai".to_string(),
provider_type: "custom".to_string(),
provider_priority: 10,
provider_is_active: true,
endpoint_id: "endpoint-openai-cli-local-timeout-1".to_string(),
endpoint_api_format: "openai:cli".to_string(),
endpoint_api_family: Some("openai".to_string()),
endpoint_kind: Some("cli".to_string()),
endpoint_is_active: true,
key_id: "key-openai-cli-local-timeout-1".to_string(),
key_name: "prod".to_string(),
key_auth_type: "api_key".to_string(),
key_is_active: true,
key_api_formats: Some(vec!["openai:cli".to_string()]),
key_allowed_models: None,
key_capabilities: None,
key_internal_priority: 5,
key_global_priority_by_format: Some(serde_json::json!({"openai:cli": 1})),
model_id: "model-openai-cli-local-timeout-1".to_string(),
global_model_id: "global-model-openai-cli-local-timeout-1".to_string(),
global_model_name: "gpt-5".to_string(),
global_model_mappings: None,
global_model_supports_streaming: Some(true),
model_provider_model_name: "gpt-5-upstream".to_string(),
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
name: "gpt-5-upstream".to_string(),
priority: 1,
api_formats: Some(vec!["openai:cli".to_string()]),
}]),
model_supports_streaming: Some(true),
model_is_active: true,
model_is_available: true,
}
}
fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
"provider-openai-cli-local-timeout-1".to_string(),
"openai".to_string(),
Some("https://example.com".to_string()),
"custom".to_string(),
)
.expect("provider should build")
.with_transport_fields(
true,
false,
false,
None,
Some(2),
None,
Some(20.0),
None,
None,
)
}
fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
"endpoint-openai-cli-local-timeout-1".to_string(),
"provider-openai-cli-local-timeout-1".to_string(),
"openai:cli".to_string(),
Some("openai".to_string()),
Some("cli".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
"https://api.openai.example/custom/v1/responses".to_string(),
None,
None,
Some(2),
Some("/custom/v1/responses".to_string()),
None,
None,
None,
)
.expect("endpoint transport should build")
}
fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
StoredProviderCatalogKey::new(
"key-openai-cli-local-timeout-1".to_string(),
"provider-openai-cli-local-timeout-1".to_string(),
"prod".to_string(),
"api_key".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(serde_json::json!(["openai:cli"])),
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
"sk-upstream-openai-cli-timeout",
)
.expect("api key should encrypt"),
None,
None,
Some(serde_json::json!({"openai:cli": 1})),
None,
None,
None,
None,
)
.expect("key transport should build")
}
let execution_runtime_hits = Arc::new(Mutex::new(0usize));
let execution_runtime_hits_clone = Arc::clone(&execution_runtime_hits);
let public_hits = Arc::new(Mutex::new(0usize));
let public_hits_clone = Arc::clone(&public_hits);
let now_unix_ms = chrono::Utc::now().timestamp_millis().max(0);
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
aether_data_contracts::repository::candidates::StoredRequestCandidate::new(
"cand-pending-openai-cli-local-timeout-1".to_string(),
"req-inflight-openai-cli-local-timeout-1".to_string(),
Some("user-openai-cli-local-timeout-123".to_string()),
Some("key-openai-cli-local-timeout-123".to_string()),
Some("alice".to_string()),
Some("default".to_string()),
0,
0,
Some("provider-openai-cli-local-timeout-1".to_string()),
Some("endpoint-openai-cli-local-timeout-1".to_string()),
Some("key-openai-cli-local-timeout-1".to_string()),
RequestCandidateStatus::Pending,
None,
false,
None,
None,
None,
None,
None,
None,
None,
now_unix_ms,
Some(now_unix_ms),
None,
)
.expect("pending candidate should build"),
]));
let upstream = Router::new()
.route(
"/api/internal/gateway/resolve",
any(|_request: Request| async move {
Json(json!({
"action": "proxy_public",
"route_class": "ai_public",
"route_family": "openai",
"route_kind": "cli",
"auth_endpoint_signature": "openai:cli",
"execution_runtime_candidate": true,
"auth_context": {
"user_id": "user-openai-cli-local-timeout-123",
"api_key_id": "key-openai-cli-local-timeout-123",
"access_allowed": true
},
"public_path": "/v1/responses"
}))
}),
)
.route(
"/api/internal/gateway/decision-sync",
any(|_request: Request| async move { Json(json!({"action": "proxy_public"})) }),
)
.route(
"/api/internal/gateway/plan-sync",
any(|_request: Request| async move { Json(json!({"action": "proxy_public"})) }),
)
.route(
"/v1/responses",
any(move |_request: Request| {
let public_hits_inner = Arc::clone(&public_hits_clone);
async move {
*public_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::IM_A_TEAPOT, Body::from("public-route-hit"))
}
}),
);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |_request: Request| {
let execution_runtime_hits_inner = Arc::clone(&execution_runtime_hits_clone);
async move {
*execution_runtime_hits_inner
.lock()
.expect("mutex should lock") += 1;
Json(json!({
"request_id": "trace-openai-cli-local-timeout-123",
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"id": "resp-cli-local-timeout-123",
"object": "response",
"model": "gpt-5-upstream",
"output": [],
"usage": {
"input_tokens": 1,
"output_tokens": 2,
"total_tokens": 3
}
}
},
"telemetry": {
"elapsed_ms": 21
}
}))
}
}),
);
let mut auth_snapshot = sample_auth_snapshot(
"key-openai-cli-local-timeout-123",
"user-openai-cli-local-timeout-123",
);
auth_snapshot.api_key_concurrent_limit = Some(1);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-client-openai-cli-local-timeout")),
auth_snapshot,
)]));
let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
sample_candidate_row(),
]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider_catalog_provider()],
vec![sample_provider_catalog_endpoint()],
vec![sample_provider_catalog_key()],
));
let (_upstream_url, upstream_handle) = start_server(upstream).await;
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
auth_repository,
candidate_selection_repository,
provider_catalog_repository,
Arc::clone(&request_candidate_repository),
DEVELOPMENT_ENCRYPTION_KEY,
),
);
let gateway = build_router_with_state(gateway_state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let started_at = std::time::Instant::now();
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/responses"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(
http::header::AUTHORIZATION,
"Bearer sk-client-openai-cli-local-timeout",
)
.header(TRACE_ID_HEADER, "trace-openai-cli-local-timeout-123")
.body("{\"model\":\"gpt-5\",\"input\":\"hello\",\"store\":false}")
.send()
.await
.expect("request should complete");
assert!(
started_at.elapsed() >= std::time::Duration::from_millis(100),
"request should wait for the bounded concurrency window before failing"
);
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
assert_eq!(
response
.headers()
.get(EXECUTION_PATH_HEADER)
.and_then(|value| value.to_str().ok()),
Some(EXECUTION_PATH_LOCAL_API_KEY_CONCURRENCY_LIMITED)
);
assert_eq!(
response
.headers()
.get(LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER)
.and_then(|value| value.to_str().ok()),
Some("api_key_concurrency_limit_reached")
);
let payload: serde_json::Value = response.json().await.expect("body should parse");
assert_eq!(
payload["error"]["message"],
serde_json::Value::String("当前 API Key 并发请求数已达上限，请稍后重试".to_string())
);
let stored_candidates = request_candidate_repository
.list_by_request_id("trace-openai-cli-local-timeout-123")
.await
.expect("request candidate trace should read");
assert_eq!(stored_candidates.len(), 1);
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Skipped);
assert_eq!(
stored_candidates[0].skip_reason.as_deref(),
Some("api_key_concurrency_limit_reached")
);
assert_eq!(
*execution_runtime_hits.lock().expect("mutex should lock"),
0
);
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
execution_runtime_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_returns_openai_cli_error_for_local_sync_failure() {
fn hash_api_key(value: &str) -> String {

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

@@ -447,6 +447,7 @@ async fn gateway_handles_admin_api_keys_create_locally_with_trusted_admin_princi
assert_eq!(payload["name"], json!("standalone-key"));
assert_eq!(payload["is_standalone"], json!(true));
assert_eq!(payload["rate_limit"], serde_json::Value::Null);
assert_eq!(payload["concurrent_limit"], serde_json::Value::Null);
assert_eq!(payload["allowed_providers"], json!(["openai"]));
assert_eq!(payload["allowed_api_formats"], json!(["openai:chat"]));
assert_eq!(payload["allowed_models"], json!(["gpt-4.1"]));
@@ -515,6 +516,7 @@ async fn gateway_handles_admin_api_keys_update_locally_with_trusted_admin_princi
.json(&json!({
"name": "renamed-key",
"rate_limit": null,
"concurrent_limit": 12,
"allowed_providers": ["gemini"],
"allowed_api_formats": ["gemini:chat"],
"allowed_models": ["gemini-2.5-pro"],
@@ -531,6 +533,7 @@ async fn gateway_handles_admin_api_keys_update_locally_with_trusted_admin_princi
assert_eq!(payload["id"], json!("key-123"));
assert_eq!(payload["name"], json!("renamed-key"));
assert_eq!(payload["rate_limit"], serde_json::Value::Null);
assert_eq!(payload["concurrent_limit"], json!(12));
assert_eq!(payload["allowed_providers"], json!(["gemini"]));
assert_eq!(payload["allowed_api_formats"], json!(["gemini:chat"]));
assert_eq!(payload["allowed_models"], json!(["gemini-2.5-pro"]));

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

@@ -629,6 +629,7 @@ async fn gateway_handles_admin_user_api_key_routes_locally_with_trusted_admin_pr
.expect("json body should parse");
assert_eq!(create_payload["name"], "new-key");
assert_eq!(create_payload["rate_limit"], 90);
assert_eq!(create_payload["concurrent_limit"], serde_json::Value::Null);
assert_eq!(
create_payload["message"],
"API Key创建成功，请妥善保存完整密钥"
@@ -651,6 +652,7 @@ async fn gateway_handles_admin_user_api_key_routes_locally_with_trusted_admin_pr
.json(&json!({
"name": "renamed",
"rate_limit": 120,
"concurrent_limit": 9,
}))
.send()
.await
@@ -664,6 +666,7 @@ async fn gateway_handles_admin_user_api_key_routes_locally_with_trusted_admin_pr
assert_eq!(update_payload["name"], "renamed");
assert_eq!(update_payload["is_locked"], false);
assert_eq!(update_payload["rate_limit"], 120);
assert_eq!(update_payload["concurrent_limit"], 9);
assert_eq!(update_payload["message"], "API Key更新成功");
let lock_response = client

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

@@ -6206,6 +6206,7 @@ async fn gateway_handles_users_me_api_key_writes_locally_without_proxying_upstre
.to_string();
assert_eq!(create_payload["name"], "writer-key");
assert_eq!(create_payload["rate_limit"], 120);
assert_eq!(create_payload["concurrent_limit"], serde_json::Value::Null);
assert_eq!(create_payload["message"], "API密钥创建成功");
assert!(create_payload["key"]
.as_str()
@@ -6219,7 +6220,8 @@ async fn gateway_handles_users_me_api_key_writes_locally_without_proxying_upstre
.header("user-agent", "AetherTest/1.0")
.json(&json!({
"name": "writer-key-renamed",
"rate_limit": 30
"rate_limit": 30,
"concurrent_limit": 4
}))
.send()
.await
@@ -6231,6 +6233,7 @@ async fn gateway_handles_users_me_api_key_writes_locally_without_proxying_upstre
.expect("json body should parse");
assert_eq!(update_payload["name"], "writer-key-renamed");
assert_eq!(update_payload["rate_limit"], 30);
assert_eq!(update_payload["concurrent_limit"], 4);
assert_eq!(update_payload["message"], "API密钥已更新");
let toggle_response = client
@@ -6308,6 +6311,7 @@ async fn gateway_handles_users_me_api_key_writes_locally_without_proxying_upstre
detail_payload["allowed_providers"],
json!(["provider-openai"])
);
assert_eq!(detail_payload["concurrent_limit"], 4);
assert_eq!(detail_payload["force_capabilities"], json!({}));
let delete_response = client

11
apps/aether-gateway/src/tests/usage/local.rs
View File

@@ -1333,14 +1333,14 @@ async fn gateway_records_failed_usage_when_all_local_claude_cli_candidates_are_s
stored_usage.user_id.as_deref(),
Some("user-claude-cli-usage-local-miss-1")
);
assert_eq!(stored_usage.provider_name, "RightCode");
assert_eq!(stored_usage.provider_name, "claude");
assert_eq!(stored_usage.model, "gpt-5.4");
assert_eq!(stored_usage.api_format.as_deref(), Some("claude:cli"));
assert_eq!(
stored_usage.endpoint_api_format.as_deref(),
Some("openai:cli")
Some("claude:cli")
);
assert_eq!(stored_usage.routing_key_name(), Some("codex"));
assert_eq!(stored_usage.routing_key_name(), None);
assert_eq!(stored_usage.routing_planner_kind(), Some("claude_cli_sync"));
assert_eq!(stored_usage.routing_route_family(), Some("claude"));
assert_eq!(stored_usage.routing_route_kind(), Some("cli"));
@@ -1375,10 +1375,7 @@ async fn gateway_records_failed_usage_when_all_local_claude_cli_candidates_are_s
stored_candidates[0].skip_reason.as_deref(),
Some("format_conversion_disabled")
);
assert_eq!(
stored_usage.routing_candidate_id(),
Some(stored_candidates[0].id.as_str())
);
assert_eq!(stored_usage.routing_candidate_id(), None);
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();