From 9631b229b3ac07d5ecccb7990268717a72d086a2 Mon Sep 17 00:00:00 2001 From: ZheFox <77232781+zhefox@users.noreply.github.com> Date: Tue, 1 Sep 2026 08:06:58 +0800 Subject: [PATCH] fix(gateway): add provider key concurrency and cache affinity modes --- apps/aether-gateway/src/control/auth/gate.rs | 39 +++- .../src/dispatch/pool_scheduler.rs | 61 ++++-- .../src/execution_runtime/stream/execution.rs | 52 +++-- .../src/execution_runtime/sync/execution.rs | 38 +++- .../handlers/admin/provider/pool/config.rs | 8 + .../admin/provider/pool_admin/payloads.rs | 1 + .../proxy/websocket/responses/admission.rs | 21 +- .../runtime/oauth_token_refresh.rs | 55 +++++- apps/aether-gateway/src/provider_key_auth.rs | 19 +- .../src/provider_pool_demand.rs | 187 ++++++++++++++++-- apps/aether-gateway/src/state/oauth.rs | 93 ++++++++- .../src/tests/control/admin/pool.rs | 4 + .../src/wallet_runtime/access.rs | 26 ++- .../src/provider/providers/antigravity.rs | 69 +++++-- .../src/provider/providers/generic.rs | 7 +- crates/aether-provider/pool/src/presets.rs | 52 ++++- .../transport/src/generic_oauth/mod.rs | 51 ++++- crates/aether-runtime/state/src/lib.rs | 76 ++++++- frontend/src/api/endpoints/pool.ts | 1 + .../pool/components/PoolKeyQuotaPanel.vue | 65 +++++- .../pool/components/PoolSchedulingDialog.vue | 43 +++- .../__tests__/PoolKeyDisplayPanels.spec.ts | 24 +++ ...oolSchedulingDialog.cache-affinity.spec.ts | 125 ++++++++++++ .../__tests__/poolKeyBatchSettings.spec.ts | 12 ++ .../features/pool/utils/poolQuotaRefresh.ts | 32 +-- frontend/src/views/admin/PoolManagement.vue | 169 ++++++++++++++-- 26 files changed, 1176 insertions(+), 154 deletions(-) create mode 100644 frontend/src/features/pool/components/__tests__/PoolSchedulingDialog.cache-affinity.spec.ts diff --git a/apps/aether-gateway/src/control/auth/gate.rs b/apps/aether-gateway/src/control/auth/gate.rs index 208c13cf3..0b5ce4d33 100644 --- a/apps/aether-gateway/src/control/auth/gate.rs +++ b/apps/aether-gateway/src/control/auth/gate.rs @@ -205,8 +205,8 @@ async fn available_balance_capacity_usd( .as_ref() .is_some_and(|wallet| wallet.limit_mode.eq_ignore_ascii_case("unlimited")); Ok(match quota.as_ref() { - Some(quota) if !quota.allow_wallet_overage => Some(quota.remaining_usd.max(0.0)), Some(_) if wallet_is_unlimited => None, + Some(quota) if !quota.allow_wallet_overage => Some(quota.remaining_usd.max(0.0)), Some(quota) => Some(quota.remaining_usd.max(0.0) + wallet_available_usd.unwrap_or(0.0)), None if wallet_is_unlimited => None, None => wallet_available_usd, @@ -832,9 +832,10 @@ mod tests { use serde_json::json; use super::{ - execution_plan_balance_capacity_rejection, execution_plan_cost_upper_bound_cache_key, - max_output_tokens_from_request, openai_request_input_is_self_contained, - output_choice_count_upper_bound, request_model_local_rejection, GatewayLocalAuthRejection, + available_balance_capacity_usd, execution_plan_balance_capacity_rejection, + execution_plan_cost_upper_bound_cache_key, max_output_tokens_from_request, + openai_request_input_is_self_contained, output_choice_count_upper_bound, + request_model_local_rejection, GatewayLocalAuthRejection, }; use crate::control::{GatewayControlAuthContext, GatewayControlDecision}; use crate::data::GatewayDataState; @@ -939,6 +940,14 @@ mod tests { fn state_with_quota_and_wallet( quota: UserDailyQuotaAvailabilityRecord, context: StoredBillingModelContext, + ) -> AppState { + state_with_quota_context_and_wallet(quota, context, sample_wallet("user-1", 30.0)) + } + + fn state_with_quota_context_and_wallet( + quota: UserDailyQuotaAvailabilityRecord, + context: StoredBillingModelContext, + wallet: StoredWalletSnapshot, ) -> AppState { let candidate_repository = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![ @@ -952,7 +961,7 @@ mod tests { AppState::new() .expect("state should build") .with_data_state_for_tests(data) - .with_auth_wallets_for_tests(vec![sample_wallet("user-1", 30.0)]) + .with_auth_wallets_for_tests(vec![wallet]) } fn state_with_model_mapping() -> AppState { @@ -1350,6 +1359,26 @@ mod tests { } } + #[tokio::test] + async fn unlimited_wallet_capacity_ignores_exhausted_non_overage_quota() { + let context = billing_context_with_pricing(None, None, None, None); + let mut wallet = sample_wallet("user-1", 0.0); + wallet.limit_mode = "unlimited".to_string(); + let state = + state_with_quota_context_and_wallet(quota_availability(0.0, false), context, wallet); + let decision = decision_with_allowed_models(vec!["gpt-5".to_string()]); + let auth_context = decision + .auth_context + .as_ref() + .expect("decision should include auth context"); + + let capacity = available_balance_capacity_usd(&state, auth_context) + .await + .expect("capacity should resolve"); + + assert_eq!(capacity, None); + } + #[tokio::test] async fn positive_balance_does_not_allow_historical_invalid_processing_pricing() { let context = billing_context_with_pricing( diff --git a/apps/aether-gateway/src/dispatch/pool_scheduler.rs b/apps/aether-gateway/src/dispatch/pool_scheduler.rs index 8fe81fce5..d79af387b 100644 --- a/apps/aether-gateway/src/dispatch/pool_scheduler.rs +++ b/apps/aether-gateway/src/dispatch/pool_scheduler.rs @@ -1826,17 +1826,17 @@ fn pool_key_candidate_order_for_group( }) .collect::>(); let active_presets = ProviderPoolService::with_builtin_adapters() - .normalize_scheduling_presets(group.transport.provider.provider_type.as_str(), &presets) - .into_iter() - .map(|preset| preset.preset) - .collect::>(); + .normalize_scheduling_presets(group.transport.provider.provider_type.as_str(), &presets); if let Some(distribution_mode) = active_presets .iter() - .find(|preset| pool_distribution_mode_preset(preset.as_str())) - .map(String::as_str) + .find(|preset| pool_distribution_mode_preset(preset.preset.as_str())) { - return match distribution_mode { - "cache_affinity" => StoredPoolKeyCandidateOrder::CacheAffinity, + return match distribution_mode.preset.as_str() { + "cache_affinity" => match distribution_mode.mode.as_deref() { + Some("lru") => StoredPoolKeyCandidateOrder::Lru, + Some("single_account") => StoredPoolKeyCandidateOrder::SingleAccount, + _ => StoredPoolKeyCandidateOrder::CacheAffinity, + }, "load_balance" => StoredPoolKeyCandidateOrder::LoadBalance { seed: pool_sort_seed(), }, @@ -2151,7 +2151,7 @@ mod tests { } #[test] - fn pool_scheduler_promotes_sticky_hit_before_other_sorted_keys() { + fn pool_scheduler_promotes_sticky_hit_before_lru_secondary_order() { let key_a = sample_eligible_candidate( "provider-pool", "endpoint-1", @@ -2159,7 +2159,11 @@ mod tests { 10, Some(json!({ "pool_advanced": { - "scheduling_presets": [{"preset": "cache_affinity", "enabled": true}] + "scheduling_presets": [{ + "preset": "cache_affinity", + "enabled": true, + "mode": "lru" + }] } })), ); @@ -2170,7 +2174,11 @@ mod tests { 10, Some(json!({ "pool_advanced": { - "scheduling_presets": [{"preset": "cache_affinity", "enabled": true}] + "scheduling_presets": [{ + "preset": "cache_affinity", + "enabled": true, + "mode": "lru" + }] } })), ); @@ -2204,6 +2212,37 @@ mod tests { ); } + #[test] + fn cache_affinity_secondary_modes_select_distinct_candidate_orders() { + for (mode, expected) in [ + ("single_account", StoredPoolKeyCandidateOrder::SingleAccount), + ("lru", StoredPoolKeyCandidateOrder::Lru), + ] { + let group = sample_eligible_candidate( + "provider-pool", + "endpoint-1", + "key-a", + 10, + Some(json!({ + "pool_advanced": { + "scheduling_presets": [{ + "preset": "cache_affinity", + "enabled": true, + "mode": mode + }] + } + })), + ); + let config = pool_config_for_candidate(&group).expect("pool config should parse"); + + assert!(admin_provider_pool_cache_affinity_enabled(&config)); + assert_eq!( + pool_key_candidate_order_for_group(&group, Some(&config)), + expected + ); + } + } + #[test] fn pool_scheduler_ignores_sticky_hit_without_cache_affinity() { let key_a = sample_eligible_candidate( diff --git a/apps/aether-gateway/src/execution_runtime/stream/execution.rs b/apps/aether-gateway/src/execution_runtime/stream/execution.rs index e7c90815a..80830455e 100644 --- a/apps/aether-gateway/src/execution_runtime/stream/execution.rs +++ b/apps/aether-gateway/src/execution_runtime/stream/execution.rs @@ -126,7 +126,7 @@ use crate::orchestration::{ LocalOAuthSuccessEffect, LocalPoolErrorEffect, }; use crate::provider_pool_demand::{ - acquire_provider_pool_in_flight_guard, ProviderPoolInFlightGuard, + acquire_provider_pool_execution_guard, ProviderPoolInFlightAdmission, ProviderPoolInFlightGuard, }; use crate::request_candidate_runtime::{ ensure_execution_request_candidate_slot, persist_local_request_candidate_status_record, @@ -3772,6 +3772,41 @@ async fn execute_execution_runtime_stream_inner( plan_kind, report_context.as_ref(), ); + let candidate_started_unix_secs = current_request_candidate_unix_ms(); + let provider_in_flight_started_at = Instant::now(); + let mut provider_pool_in_flight_guard = + match acquire_provider_pool_execution_guard(state, &plan).await? { + ProviderPoolInFlightAdmission::Acquired(guard) => guard, + ProviderPoolInFlightAdmission::Saturated { limit } => { + if let Some(retry_scope) = retry_scope_out.as_deref_mut() { + *retry_scope = AiAttemptRetryScope::Candidate; + } + if let Some(snapshot) = request_candidate_status_snapshot.as_ref() { + record_local_request_candidate_status_snapshot( + state, + snapshot, + SchedulerRequestCandidateStatusUpdate { + status: RequestCandidateStatus::Skipped, + status_code: Some(http::StatusCode::TOO_MANY_REQUESTS.as_u16()), + error_type: Some("provider_key_concurrency_limit_reached".to_string()), + error_message: Some(format!( + "provider key concurrency limit reached: {limit}" + )), + latency_ms: Some(0), + started_at_unix_ms: Some(candidate_started_unix_secs), + finished_at_unix_ms: Some(candidate_started_unix_secs), + }, + ) + .await; + } + return Ok(None); + } + }; + observe_gateway_stage_trace_ms( + &mut stage_trace, + "stream_provider_in_flight", + provider_in_flight_started_at.elapsed().as_millis() as u64, + ); // Inline passthrough records its lifecycle seed after upstream headers are // available. Avoid constructing a throwaway seed on the common path. let mut lifecycle_seed = (!defer_stream_pending_for_direct_inline) @@ -3781,7 +3816,6 @@ async fn execute_execution_runtime_stream_inner( record_stream_pending_lifecycle(state, seed, &mut stage_trace).await; lifecycle_pending_recorded = true; } - let candidate_started_unix_secs = current_request_candidate_unix_ms(); if let Some(snapshot) = request_candidate_status_snapshot.clone() { record_local_request_candidate_status_snapshot( state, @@ -3810,20 +3844,6 @@ async fn execute_execution_runtime_stream_inner( .and_then(|context| context.candidate_index) .map(|value| value.to_string()) .unwrap_or_else(|| "-".to_string()); - let provider_in_flight_started_at = Instant::now(); - let mut provider_pool_in_flight_guard = acquire_provider_pool_in_flight_guard( - state.runtime_state.clone(), - &plan.provider_id, - plan.request_id.as_str(), - plan.candidate_id.as_deref(), - key_id.as_str(), - ) - .await; - observe_gateway_stage_trace_ms( - &mut stage_trace, - "stream_provider_in_flight", - provider_in_flight_started_at.elapsed().as_millis() as u64, - ); match maybe_execute_grok_stream(&plan, report_context.as_ref()).await { Ok(Some(grok_stream)) => { return execute_stream_from_frame_stream_with_retry_scope( diff --git a/apps/aether-gateway/src/execution_runtime/sync/execution.rs b/apps/aether-gateway/src/execution_runtime/sync/execution.rs index 5d6266c7b..84d67a25a 100644 --- a/apps/aether-gateway/src/execution_runtime/sync/execution.rs +++ b/apps/aether-gateway/src/execution_runtime/sync/execution.rs @@ -78,7 +78,9 @@ use crate::orchestration::{ LocalExecutionEffectContext, LocalHealthFailureEffect, LocalHealthSuccessEffect, LocalOAuthInvalidationEffect, LocalOAuthSuccessEffect, LocalPoolErrorEffect, }; -use crate::provider_pool_demand::acquire_provider_pool_in_flight_guard; +use crate::provider_pool_demand::{ + acquire_provider_pool_execution_guard, ProviderPoolInFlightAdmission, +}; use crate::request_candidate_runtime::{ ensure_execution_request_candidate_slot, record_local_request_candidate_extra_data, record_local_request_candidate_status, record_local_request_candidate_status_snapshot, @@ -1995,6 +1997,32 @@ async fn execute_execution_runtime_sync_impl( .unwrap_or_else(|| "-".to_string()); let candidate_started_at = Instant::now(); let candidate_started_unix_secs = current_request_candidate_unix_ms(); + let _provider_pool_in_flight_guard = match acquire_provider_pool_execution_guard(state, &plan) + .await? + { + ProviderPoolInFlightAdmission::Acquired(guard) => guard, + ProviderPoolInFlightAdmission::Saturated { limit } => { + if let Some(retry_scope) = retry_scope_out.as_deref_mut() { + *retry_scope = AiAttemptRetryScope::Candidate; + } + record_local_request_candidate_status( + state, + &plan, + report_context.as_ref(), + SchedulerRequestCandidateStatusUpdate { + status: RequestCandidateStatus::Skipped, + status_code: Some(StatusCode::TOO_MANY_REQUESTS.as_u16()), + error_type: Some("provider_key_concurrency_limit_reached".to_string()), + error_message: Some(format!("provider key concurrency limit reached: {limit}")), + latency_ms: Some(0), + started_at_unix_ms: Some(candidate_started_unix_secs), + finished_at_unix_ms: Some(candidate_started_unix_secs), + }, + ) + .await; + return Ok(None); + } + }; let lifecycle_seed = build_lifecycle_usage_seed(&plan, report_context.as_ref()); let usage_data = state.usage_lifecycle_data_state().as_ref().clone(); state @@ -2024,14 +2052,6 @@ async fn execute_execution_runtime_sync_impl( candidate_started_at, ); let result = (async { - let _provider_pool_in_flight_guard = acquire_provider_pool_in_flight_guard( - state.runtime_state.clone(), - &plan.provider_id, - plan_request_id.as_str(), - plan_candidate_id.as_deref(), - key_id.as_str(), - ) - .await; record_sync_execution_active( state, &plan, diff --git a/apps/aether-gateway/src/handlers/admin/provider/pool/config.rs b/apps/aether-gateway/src/handlers/admin/provider/pool/config.rs index e9fe4c008..216f00405 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/pool/config.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/pool/config.rs @@ -167,6 +167,14 @@ fn parse_pool_score_rules(pool_advanced: &Map) -> PoolMemberScore fn normalize_pool_preset_mode(preset: &str, raw_mode: Option<&Value>) -> Option { match preset { + "cache_affinity" => Some( + raw_mode + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| matches!(*value, "single_account" | "lru")) + .unwrap_or("single_account") + .to_string(), + ), "free_team_first" | "free_first" | "team_first" | "plus_first" | "pro_first" => { let default_mode = match preset { "free_team_first" => "both", diff --git a/apps/aether-gateway/src/handlers/admin/provider/pool_admin/payloads.rs b/apps/aether-gateway/src/handlers/admin/provider/pool_admin/payloads.rs index 21531664e..98076f0df 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/pool_admin/payloads.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/pool_admin/payloads.rs @@ -1342,6 +1342,7 @@ pub(super) fn build_admin_pool_key_payload( json!(key.internal_priority), ); payload.insert("rpm_limit".to_string(), json!(key.rpm_limit)); + payload.insert("concurrent_limit".to_string(), json!(key.concurrent_limit)); payload.insert( "cache_ttl_minutes".to_string(), json!(key.cache_ttl_minutes), diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/responses/admission.rs b/apps/aether-gateway/src/handlers/proxy/websocket/responses/admission.rs index 8192f5e90..31f1e4e73 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/responses/admission.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/responses/admission.rs @@ -11,7 +11,7 @@ use aether_contracts::ExecutionPlan; use crate::execution_runtime::acquire_upstream_execution_gate; use crate::provider_pool_demand::{ - acquire_provider_pool_in_flight_guard, ProviderPoolInFlightGuard, + acquire_provider_pool_execution_guard, ProviderPoolInFlightAdmission, ProviderPoolInFlightGuard, }; use crate::upstream_admission::UpstreamTargetAdmissionPermit; use crate::{AppState, GatewayError}; @@ -41,14 +41,17 @@ impl ResponsesWebSocketTurnAdmission { return Err(error); } }; - let provider_pool = acquire_provider_pool_in_flight_guard( - state.runtime_state.clone(), - &plan.provider_id, - &plan.request_id, - plan.candidate_id.as_deref(), - &plan.key_id, - ) - .await; + let provider_pool = match acquire_provider_pool_execution_guard(state, plan).await? { + ProviderPoolInFlightAdmission::Acquired(guard) => guard, + ProviderPoolInFlightAdmission::Saturated { limit } => { + drop(upstream_target); + drop(upstream_execution); + return Err(GatewayError::Client { + status: http::StatusCode::TOO_MANY_REQUESTS, + message: format!("上游账号并发已达上限 ({limit})"), + }); + } + }; Ok(Self { upstream_execution, diff --git a/apps/aether-gateway/src/maintenance/runtime/oauth_token_refresh.rs b/apps/aether-gateway/src/maintenance/runtime/oauth_token_refresh.rs index 9cb608c12..7ec147a7e 100644 --- a/apps/aether-gateway/src/maintenance/runtime/oauth_token_refresh.rs +++ b/apps/aether-gateway/src/maintenance/runtime/oauth_token_refresh.rs @@ -256,12 +256,16 @@ fn auth_config_has_refresh_token(auth_config: Option<&str>) -> bool { let Ok(value) = serde_json::from_str::(auth_config) else { return false; }; - value - .as_object() - .and_then(|object| object.get("refresh_token")) - .and_then(Value::as_str) - .map(str::trim) - .is_some_and(|value| !value.is_empty()) + let Some(object) = value.as_object() else { + return false; + }; + ["refresh_token", "refreshToken"].iter().any(|field| { + object + .get(*field) + .and_then(Value::as_str) + .map(str::trim) + .is_some_and(|value| !value.is_empty()) + }) } fn now_unix_secs() -> u64 { @@ -273,7 +277,44 @@ fn now_unix_secs() -> u64 { #[cfg(test)] mod tests { - use super::agent_identity_needs_task_recovery; + use aether_data_contracts::repository::provider_catalog::{ + StoredProviderCatalogKey, StoredProviderCatalogProvider, + }; + + use super::{ + agent_identity_needs_task_recovery, auth_config_has_refresh_token, oauth_refresh_candidate, + }; + + #[test] + fn legacy_antigravity_refresh_token_is_refreshable() { + assert!(auth_config_has_refresh_token(Some( + r#"{"refreshToken":"legacy-refresh-token"}"#, + ))); + } + + #[test] + fn expiring_antigravity_oauth_key_is_refresh_candidate() { + let provider = StoredProviderCatalogProvider::new( + "provider-antigravity".to_string(), + "Antigravity".to_string(), + None, + "antigravity".to_string(), + ) + .expect("provider should build"); + let mut key = StoredProviderCatalogKey::new( + "key-antigravity".to_string(), + provider.id.clone(), + "Antigravity OAuth".to_string(), + "oauth".to_string(), + None, + true, + ) + .expect("key should build"); + key.encrypted_auth_config = Some("encrypted-auth-config".to_string()); + key.expires_at_unix_secs = Some(120); + + assert!(oauth_refresh_candidate(&provider, &key, 120)); + } #[test] fn pending_agent_identity_without_task_is_recoverable() { diff --git a/apps/aether-gateway/src/provider_key_auth.rs b/apps/aether-gateway/src/provider_key_auth.rs index 8bc5ae2dc..62379e62f 100644 --- a/apps/aether-gateway/src/provider_key_auth.rs +++ b/apps/aether-gateway/src/provider_key_auth.rs @@ -90,11 +90,15 @@ pub(crate) fn provider_key_can_refresh_oauth( ) -> bool { auth_semantics.can_refresh_oauth() && (provider_key_auth_config_is_agent_identity(provider_type, auth_config) - || auth_config - .and_then(|config| config.get("refresh_token")) - .and_then(Value::as_str) - .map(str::trim) - .is_some_and(|value| !value.is_empty())) + || auth_config.is_some_and(|config| { + ["refresh_token", "refreshToken"].iter().any(|field| { + config + .get(*field) + .and_then(Value::as_str) + .map(str::trim) + .is_some_and(|value| !value.is_empty()) + }) + })) } pub(crate) fn provider_key_can_export_oauth( @@ -425,6 +429,11 @@ mod tests { "codex", json!({ "refresh_token": "refresh-token" }).as_object() )); + assert!(provider_key_can_refresh_oauth( + provider_key_auth_semantics(&sample_key("oauth"), "antigravity"), + "antigravity", + json!({ "refreshToken": "legacy-refresh-token" }).as_object() + )); assert!(provider_key_can_refresh_oauth( semantics, "codex", diff --git a/apps/aether-gateway/src/provider_pool_demand.rs b/apps/aether-gateway/src/provider_pool_demand.rs index 819755995..6506cb56e 100644 --- a/apps/aether-gateway/src/provider_pool_demand.rs +++ b/apps/aether-gateway/src/provider_pool_demand.rs @@ -4,14 +4,20 @@ use std::sync::{ }; use std::time::{Duration, SystemTime, UNIX_EPOCH}; -use aether_runtime_state::RuntimeState; +use aether_contracts::ExecutionPlan; +use aether_runtime_state::{ + RuntimeSemaphoreConfig, RuntimeSemaphoreError, RuntimeSemaphorePermit, RuntimeState, +}; use dashmap::DashMap; use serde::{Deserialize, Serialize}; use tokio::task::JoinHandle; use tracing::debug; use uuid::Uuid; +use crate::{AppState, GatewayError}; + const PROVIDER_POOL_IN_FLIGHT_TOKENS_PREFIX: &str = "ap:provider_pool:in_flight"; +const PROVIDER_KEY_CONCURRENCY_GATE: &str = "provider_key"; const PROVIDER_POOL_DEMAND_SNAPSHOT_PREFIX: &str = "ap:provider_pool:demand"; const PROVIDER_POOL_BURST_PENDING_PREFIX: &str = "ap:quota_probe:burst_pending"; const PROVIDER_POOL_IN_FLIGHT_TOKEN_TTL_MS: u64 = 120_000; @@ -42,10 +48,17 @@ pub(crate) struct ProviderPoolDemandSnapshot { pub(crate) struct ProviderPoolInFlightGuard { kind: ProviderPoolInFlightGuardKind, + provider_key_permit: Option, released: bool, } +pub(crate) enum ProviderPoolInFlightAdmission { + Acquired(Option), + Saturated { limit: usize }, +} + enum ProviderPoolInFlightGuardKind { + Disabled, Local { provider_id: String, counter: Arc, @@ -69,6 +82,9 @@ impl ProviderPoolInFlightGuard { return; } match &mut self.kind { + ProviderPoolInFlightGuardKind::Disabled => { + self.released = true; + } ProviderPoolInFlightGuardKind::Local { provider_id, counter, @@ -101,6 +117,16 @@ impl ProviderPoolInFlightGuard { } } } + if self.released { + if let Some(provider_key_permit) = self.provider_key_permit.take() { + if let Err(err) = provider_key_permit.release().await { + debug!( + error = ?err, + "gateway provider pool demand: failed to release provider key permit; scheduling drop fallback" + ); + } + } + } } } @@ -111,6 +137,7 @@ impl Drop for ProviderPoolInFlightGuard { } self.released = true; match &mut self.kind { + ProviderPoolInFlightGuardKind::Disabled => {} ProviderPoolInFlightGuardKind::Local { provider_id, counter, @@ -290,32 +317,83 @@ pub(crate) async fn acquire_provider_pool_in_flight_guard( candidate_id: Option<&str>, key_id: &str, ) -> Option { + acquire_provider_pool_in_flight_guard_with_key_limit( + runtime, + provider_id, + request_id, + candidate_id, + key_id, + None, + ) + .await + .ok() + .flatten() +} + +pub(crate) async fn acquire_provider_pool_in_flight_guard_with_key_limit( + runtime: Arc, + provider_id: &str, + request_id: &str, + candidate_id: Option<&str>, + key_id: &str, + concurrent_limit: Option, +) -> Result, RuntimeSemaphoreError> { + let provider_key_permit = match concurrent_limit.filter(|limit| *limit > 0) { + Some(limit) => Some( + runtime + .keyed_semaphore( + PROVIDER_KEY_CONCURRENCY_GATE, + key_id, + limit, + RuntimeSemaphoreConfig::default(), + )? + .try_acquire() + .await?, + ), + None => None, + }; let provider_id = provider_id.trim(); if provider_id.is_empty() { - return None; + return Ok( + provider_key_permit.map(|provider_key_permit| ProviderPoolInFlightGuard { + kind: ProviderPoolInFlightGuardKind::Disabled, + provider_key_permit: Some(provider_key_permit), + released: false, + }), + ); } match provider_pool_in_flight_mode() { - ProviderPoolInFlightMode::Off => return None, + ProviderPoolInFlightMode::Off => { + return Ok( + provider_key_permit.map(|provider_key_permit| ProviderPoolInFlightGuard { + kind: ProviderPoolInFlightGuardKind::Disabled, + provider_key_permit: Some(provider_key_permit), + released: false, + }), + ); + } ProviderPoolInFlightMode::Local => { let counter = increment_local_provider_in_flight(provider_id); - return Some(ProviderPoolInFlightGuard { + return Ok(Some(ProviderPoolInFlightGuard { kind: ProviderPoolInFlightGuardKind::Local { provider_id: provider_id.to_string(), counter, }, + provider_key_permit, released: false, - }); + })); } ProviderPoolInFlightMode::Runtime if runtime.is_memory() => { let counter = increment_local_provider_in_flight(provider_id); - return Some(ProviderPoolInFlightGuard { + return Ok(Some(ProviderPoolInFlightGuard { kind: ProviderPoolInFlightGuardKind::Local { provider_id: provider_id.to_string(), counter, }, + provider_key_permit, released: false, - }); + })); } ProviderPoolInFlightMode::Runtime => {} } @@ -335,7 +413,13 @@ pub(crate) async fn acquire_provider_pool_in_flight_guard( error = ?err, "gateway provider pool demand: failed to acquire in-flight token" ); - return None; + return Ok( + provider_key_permit.map(|provider_key_permit| ProviderPoolInFlightGuard { + kind: ProviderPoolInFlightGuardKind::Disabled, + provider_key_permit: Some(provider_key_permit), + released: false, + }), + ); } Err(_) => { debug!( @@ -343,7 +427,13 @@ pub(crate) async fn acquire_provider_pool_in_flight_guard( timeout_ms = provider_pool_in_flight_acquire_timeout().as_millis() as u64, "gateway provider pool demand: skipped in-flight token after acquire timeout" ); - return None; + return Ok( + provider_key_permit.map(|provider_key_permit| ProviderPoolInFlightGuard { + kind: ProviderPoolInFlightGuardKind::Disabled, + provider_key_permit: Some(provider_key_permit), + released: false, + }), + ); } } @@ -355,7 +445,7 @@ pub(crate) async fn acquire_provider_pool_in_flight_guard( stop_renewal.clone(), ); - Some(ProviderPoolInFlightGuard { + Ok(Some(ProviderPoolInFlightGuard { kind: ProviderPoolInFlightGuardKind::Runtime { runtime, tokens_key, @@ -363,8 +453,39 @@ pub(crate) async fn acquire_provider_pool_in_flight_guard( stop_renewal, renew_handle: Some(renew_handle), }, + provider_key_permit, released: false, - }) + })) +} + +pub(crate) async fn acquire_provider_pool_execution_guard( + state: &AppState, + plan: &ExecutionPlan, +) -> Result { + let concurrent_limit = state + .read_provider_catalog_keys_by_ids(std::slice::from_ref(&plan.key_id)) + .await? + .into_iter() + .find(|key| key.id == plan.key_id) + .and_then(|key| key.concurrent_limit) + .filter(|limit| *limit > 0) + .and_then(|limit| usize::try_from(limit).ok()); + match acquire_provider_pool_in_flight_guard_with_key_limit( + state.runtime_state.clone(), + &plan.provider_id, + &plan.request_id, + plan.candidate_id.as_deref(), + &plan.key_id, + concurrent_limit, + ) + .await + { + Ok(guard) => Ok(ProviderPoolInFlightAdmission::Acquired(guard)), + Err(RuntimeSemaphoreError::Saturated { limit, .. }) => { + Ok(ProviderPoolInFlightAdmission::Saturated { limit }) + } + Err(error) => Err(GatewayError::Internal(error.to_string())), + } } pub(crate) async fn provider_pool_live_in_flight_count( @@ -586,6 +707,50 @@ mod tests { ); } + #[tokio::test] + async fn provider_key_limit_rejects_concurrent_guard_until_release() { + let runtime = Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default())); + let first = acquire_provider_pool_in_flight_guard_with_key_limit( + runtime.clone(), + "provider-limit", + "request-1", + Some("candidate-1"), + "key-limit", + Some(1), + ) + .await + .expect("first admission should resolve") + .expect("first guard should be acquired"); + + let second = acquire_provider_pool_in_flight_guard_with_key_limit( + runtime.clone(), + "provider-limit", + "request-2", + Some("candidate-2"), + "key-limit", + Some(1), + ) + .await; + assert!(matches!( + second, + Err(RuntimeSemaphoreError::Saturated { limit: 1, .. }) + )); + + first.release().await; + let replacement = acquire_provider_pool_in_flight_guard_with_key_limit( + runtime, + "provider-limit", + "request-3", + Some("candidate-3"), + "key-limit", + Some(1), + ) + .await + .expect("replacement admission should resolve") + .expect("replacement guard should acquire after release"); + drop(replacement); + } + #[tokio::test] async fn demand_snapshot_uses_instant_in_flight_for_fast_rise_and_ema_for_fall() { let runtime = RuntimeState::memory(MemoryRuntimeStateConfig::default()); diff --git a/apps/aether-gateway/src/state/oauth.rs b/apps/aether-gateway/src/state/oauth.rs index 1f53e3a48..f541e59b9 100644 --- a/apps/aether-gateway/src/state/oauth.rs +++ b/apps/aether-gateway/src/state/oauth.rs @@ -33,7 +33,9 @@ use std::sync::atomic::Ordering; use std::sync::Arc; use std::time::{Duration, SystemTime, UNIX_EPOCH}; -use aether_crypto::encrypt_python_fernet_plaintext; +use aether_crypto::{ + decrypt_python_fernet_ciphertext, encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY, +}; const LOCAL_OAUTH_HTTP_TIMEOUT_MS: u64 = 30_000; const REMOTE_OAUTH_REFRESH_WAIT_TIMEOUT: Duration = Duration::from_secs(35); @@ -3395,7 +3397,10 @@ mod tests { use std::sync::Arc; use std::time::{Duration, Instant}; - use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY}; + use aether_crypto::{ + decrypt_python_fernet_ciphertext, encrypt_python_fernet_plaintext, + DEVELOPMENT_ENCRYPTION_KEY, + }; use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository; use aether_data_contracts::repository::provider_catalog::{ ProviderCatalogKeyAdminCasUpdate, ProviderCatalogKeyListQuery, @@ -3478,9 +3483,17 @@ mod tests { fn codex_oauth_state( auth_config: &serde_json::Value, access_token: &str, + ) -> (AppState, Arc, String) { + provider_oauth_state("codex", auth_config, access_token) + } + + fn provider_oauth_state( + provider_type: &str, + auth_config: &serde_json::Value, + access_token: &str, ) -> (AppState, Arc, String) { let mut provider = sample_provider(); - provider.provider_type = "codex".to_string(); + provider.provider_type = provider_type.to_string(); let encrypted_auth_config = encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, &auth_config.to_string()) .expect("auth config should encrypt"); @@ -3490,7 +3503,7 @@ mod tests { let key = StoredProviderCatalogKey::new( "key-1".to_string(), "provider-1".to_string(), - "Codex OAuth".to_string(), + format!("{provider_type} OAuth"), "oauth".to_string(), None, true, @@ -4564,6 +4577,78 @@ mod tests { ); } + #[tokio::test] + async fn antigravity_refresh_entry_persists_tokens_and_expiry() { + let initial_config = json!({ + "provider_type": "antigravity", + "refreshToken": "legacy-refresh-token", + "expires_at": 1, + }); + let (state, repository, _) = + provider_oauth_state("antigravity", &initial_config, "stale-access-token"); + let transport = state + .read_provider_transport_snapshot("provider-1", "endpoint-1", "key-1") + .await + .expect("transport should load") + .expect("transport should exist"); + let expected_credential_fence = state + .capture_provider_transport_credential_fence(&transport) + .await + .expect("credential fence should load") + .expect("credential fence should match"); + let expires_at = 4_102_555_900; + let refreshed_entry = crate::provider_transport::CachedOAuthEntry { + provider_type: "antigravity".to_string(), + auth_header_name: "authorization".to_string(), + auth_header_value: "Bearer fresh-access-token".to_string(), + expires_at_unix_secs: Some(expires_at), + metadata: Some(json!({ + "provider_type": "antigravity", + "refresh_token": "legacy-refresh-token", + "expires_at": expires_at, + })), + source_fingerprint: None, + }; + + state + .persist_local_oauth_refresh_entry( + &transport, + &refreshed_entry, + Some(&expected_credential_fence), + ) + .await + .expect("Antigravity refresh should persist"); + + let stored = repository + .list_keys_by_ids(&["key-1".to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should remain"); + let access_token = decrypt_python_fernet_ciphertext( + DEVELOPMENT_ENCRYPTION_KEY, + stored + .encrypted_api_key + .as_deref() + .expect("access token should persist"), + ) + .expect("access token should decrypt"); + let auth_config = decrypt_python_fernet_ciphertext( + DEVELOPMENT_ENCRYPTION_KEY, + stored + .encrypted_auth_config + .as_deref() + .expect("auth config should persist"), + ) + .expect("auth config should decrypt"); + let auth_config: serde_json::Value = + serde_json::from_str(&auth_config).expect("auth config should parse"); + + assert_eq!(access_token, "fresh-access-token"); + assert_eq!(auth_config["refresh_token"], json!("legacy-refresh-token")); + assert_eq!(stored.expires_at_unix_secs, Some(expires_at)); + } + #[tokio::test] async fn agent_auth_config_fence_rejects_metadata_only_rewrite() { let initial_config = json!({ diff --git a/apps/aether-gateway/src/tests/control/admin/pool.rs b/apps/aether-gateway/src/tests/control/admin/pool.rs index acc3b8a76..d35f5fb0b 100644 --- a/apps/aether-gateway/src/tests/control/admin/pool.rs +++ b/apps/aether-gateway/src/tests/control/admin/pool.rs @@ -629,6 +629,7 @@ async fn gateway_pool_list_includes_usage_totals_and_nullable_lru_score() { "sk-usage", ); key.name = "usage key".to_string(); + key.concurrent_limit = Some(5); key.request_count = Some(1566); key.total_tokens = 187_327_321; key.total_cost_usd = 93.1319297; @@ -661,6 +662,7 @@ async fn gateway_pool_list_includes_usage_totals_and_nullable_lru_score() { .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]["concurrent_limit"], json!(5)); 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")); @@ -3643,6 +3645,7 @@ async fn gateway_batch_updates_shared_pool_key_configuration() { "api_formats": ["openai:responses"], "internal_priority": 7, "rpm_limit": null, + "concurrent_limit": 6, "auto_fetch_models": false, "allowed_models": ["gpt-5.6-sol", "gpt-5.6-luna"], "locked_models": [], @@ -3668,6 +3671,7 @@ async fn gateway_batch_updates_shared_pool_key_configuration() { assert_eq!(key.allow_auth_channel_mismatch_formats, Some(json!([]))); assert_eq!(key.internal_priority, 7); assert_eq!(key.rpm_limit, None); + assert_eq!(key.concurrent_limit, Some(6)); assert_eq!(key.learned_rpm_limit, None); assert!(!key.auto_fetch_models); assert_eq!( diff --git a/apps/aether-gateway/src/wallet_runtime/access.rs b/apps/aether-gateway/src/wallet_runtime/access.rs index 1d314c1d6..b68824db7 100644 --- a/apps/aether-gateway/src/wallet_runtime/access.rs +++ b/apps/aether-gateway/src/wallet_runtime/access.rs @@ -55,6 +55,9 @@ async fn resolve_wallet_auth_gate_with_cache( None => WalletAccessDecision::wallet_unavailable(None), }; if !auth_snapshot.api_key_is_standalone { + let wallet_is_unlimited = wallet + .as_ref() + .is_some_and(|wallet| wallet.limit_mode.eq_ignore_ascii_case("unlimited")); let quota = if use_cache { state .find_user_daily_quota_availability_for_auth(&auth_snapshot.user_id) @@ -71,7 +74,11 @@ async fn resolve_wallet_auth_gate_with_cache( quota.remaining_usd, )))); } - if decision.failure.is_none() && !quota.allow_wallet_overage && !has_remaining_quota { + if !wallet_is_unlimited + && decision.failure.is_none() + && !quota.allow_wallet_overage + && !has_remaining_quota + { return Ok(Some(WalletAccessDecision::balance_denied(Some(0.0)))); } } @@ -264,6 +271,23 @@ mod tests { assert_eq!(decision.remaining, Some(4.0)); } + #[tokio::test] + async fn unlimited_wallet_ignores_exhausted_non_overage_quota() { + let mut wallet = empty_user_wallet(); + wallet.limit_mode = "unlimited".to_string(); + let state = state_with_wallet_and_quota(wallet, Some(quota_availability(10.0, 0.0, false))); + let auth_snapshot = ordinary_user_api_key_snapshot(); + + let decision = resolve_wallet_auth_gate(&state, &auth_snapshot) + .await + .expect("wallet gate should resolve") + .expect("wallet gate should return a decision"); + + assert!(decision.allowed); + assert_eq!(decision.failure, None); + assert_eq!(decision.remaining, None); + } + #[tokio::test] async fn disabled_auth_capacity_cache_still_gates_wallet_reads() { let mut state = state_with_wallet_and_quota(empty_user_wallet(), None); diff --git a/crates/aether-oauth/src/provider/providers/antigravity.rs b/crates/aether-oauth/src/provider/providers/antigravity.rs index 7d2f33fc9..9bc86c053 100644 --- a/crates/aether-oauth/src/provider/providers/antigravity.rs +++ b/crates/aether-oauth/src/provider/providers/antigravity.rs @@ -38,7 +38,17 @@ impl ProviderOAuthAdapter for AntigravityProviderOAuthAdapter { state: &str, code_challenge: Option<&str>, ) -> Result { - self.inner.build_authorize_url(ctx, state, code_challenge) + let mut response = self.inner.build_authorize_url(ctx, state, code_challenge)?; + let mut url = url::Url::parse(&response.authorize_url).map_err(|_| { + crate::core::OAuthError::invalid_request("authorize_url must be absolute") + })?; + { + let mut query = url.query_pairs_mut(); + query.append_pair("access_type", "offline"); + query.append_pair("prompt", "consent"); + } + response.authorize_url = url.to_string(); + Ok(response) } async fn exchange_code( @@ -112,20 +122,8 @@ mod tests { struct UnusedExecutor; - #[async_trait] - impl OAuthHttpExecutor for UnusedExecutor { - async fn execute( - &self, - _request: OAuthHttpRequest, - ) -> Result { - unreachable!("metadata probe should not execute network requests") - } - } - - #[tokio::test] - async fn antigravity_probe_marks_forbidden_metadata_invalid() { - let adapter = AntigravityProviderOAuthAdapter::default(); - let ctx = ProviderOAuthTransportContext { + fn transport_context() -> ProviderOAuthTransportContext { + ProviderOAuthTransportContext { provider_id: String::new(), provider_type: "antigravity".to_string(), endpoint_id: None, @@ -137,7 +135,46 @@ mod tests { endpoint_config: None, key_config: None, network: crate::network::OAuthNetworkContext::provider_operation(None), - }; + } + } + + #[async_trait] + impl OAuthHttpExecutor for UnusedExecutor { + async fn execute( + &self, + _request: OAuthHttpRequest, + ) -> Result { + unreachable!("metadata probe should not execute network requests") + } + } + + #[test] + fn antigravity_authorize_requests_offline_refresh_token() { + let adapter = AntigravityProviderOAuthAdapter::default(); + let response = adapter + .build_authorize_url(&transport_context(), "state-1", Some("challenge-1")) + .expect("authorize url should build"); + let url = url::Url::parse(&response.authorize_url).expect("authorize url should parse"); + let query = url.query_pairs().collect::>(); + + assert_eq!( + query.get("access_type").map(|value| value.as_ref()), + Some("offline") + ); + assert_eq!( + query.get("prompt").map(|value| value.as_ref()), + Some("consent") + ); + assert_eq!( + query.get("code_challenge").map(|value| value.as_ref()), + Some("challenge-1") + ); + } + + #[tokio::test] + async fn antigravity_probe_marks_forbidden_metadata_invalid() { + let adapter = AntigravityProviderOAuthAdapter::default(); + let ctx = transport_context(); let account = ProviderOAuthAccount { provider_type: "antigravity".to_string(), access_token: "access-token".to_string(), diff --git a/crates/aether-oauth/src/provider/providers/generic.rs b/crates/aether-oauth/src/provider/providers/generic.rs index 9e27dca46..f81c10fe9 100644 --- a/crates/aether-oauth/src/provider/providers/generic.rs +++ b/crates/aether-oauth/src/provider/providers/generic.rs @@ -374,10 +374,9 @@ impl ProviderOAuthAdapter for GenericProviderOAuthAdapter { ctx: &ProviderOAuthTransportContext, account: &ProviderOAuthAccount, ) -> Result { - let refresh_token = account - .auth_config - .get("refresh_token") - .and_then(Value::as_str) + let refresh_token = ["refresh_token", "refreshToken"] + .iter() + .find_map(|field| account.auth_config.get(*field).and_then(Value::as_str)) .map(str::trim) .filter(|value| !value.is_empty()) .ok_or_else(|| OAuthError::invalid_request("auth_config missing refresh_token"))?; diff --git a/crates/aether-provider/pool/src/presets.rs b/crates/aether-provider/pool/src/presets.rs index a3302d428..ac23051a0 100644 --- a/crates/aether-provider/pool/src/presets.rs +++ b/crates/aether-provider/pool/src/presets.rs @@ -94,14 +94,7 @@ pub fn build_admin_pool_scheduling_presets_payload() -> Value { "依据 LRU 时间戳（最近未使用优先）", &service, ), - provider_pool_preset_payload( - "cache_affinity", - "缓存亲和", - "优先复用最近使用过的 Key，利用 Prompt Caching", - None, - "依据 LRU 时间戳（最近使用优先，与 LRU 轮转相反）", - &service, - ), + cache_affinity_preset_payload(&service), provider_pool_preset_payload( "cost_first", "成本优先", @@ -220,6 +213,23 @@ fn legacy_free_team_first_preset_payload(service: &ProviderPoolService) -> Value payload } +fn cache_affinity_preset_payload(service: &ProviderPoolService) -> Value { + let mut payload = provider_pool_preset_payload( + "cache_affinity", + "缓存亲和", + "同一用户持续复用 Key，首次分配可集中或轮转", + None, + "先复用用户粘性 Key，未命中时按所选二级模式分配", + service, + ); + payload["modes"] = json!([ + {"value": "single_account", "label": "单号优先"}, + {"value": "lru", "label": "LRU 轮号"} + ]); + payload["default_mode"] = json!("single_account"); + payload +} + fn provider_pool_preset_payload( name: &'static str, label: &'static str, @@ -274,3 +284,29 @@ fn provider_pool_preset_mutex_group(preset: &str) -> Option<&'static str> { _ => None, } } + +#[cfg(test)] +mod tests { + use super::build_admin_pool_scheduling_presets_payload; + use serde_json::json; + + #[test] + fn cache_affinity_preset_exposes_secondary_distribution_modes() { + let payload = build_admin_pool_scheduling_presets_payload(); + let cache_affinity = payload + .as_array() + .expect("preset payload should be an array") + .iter() + .find(|preset| preset["name"] == "cache_affinity") + .expect("cache affinity preset should exist"); + + assert_eq!(cache_affinity["default_mode"], json!("single_account")); + assert_eq!( + cache_affinity["modes"], + json!([ + {"value": "single_account", "label": "单号优先"}, + {"value": "lru", "label": "LRU 轮号"} + ]) + ); + } +} diff --git a/crates/aether-provider/transport/src/generic_oauth/mod.rs b/crates/aether-provider/transport/src/generic_oauth/mod.rs index 0a24f7ff5..99c7fdcc5 100644 --- a/crates/aether-provider/transport/src/generic_oauth/mod.rs +++ b/crates/aether-provider/transport/src/generic_oauth/mod.rs @@ -513,10 +513,10 @@ fn generic_provider_type(provider_type: &str) -> Option<&'static str> { } fn refresh_token_from_auth_config(auth_config: &Value) -> Option { - auth_config - .as_object() - .and_then(|object| object.get("refresh_token")) - .and_then(non_empty_string) + let object = auth_config.as_object()?; + ["refresh_token", "refreshToken"] + .iter() + .find_map(|field| object.get(*field).and_then(non_empty_string)) } fn access_token_from_auth_config(auth_config: &Value) -> Option { @@ -892,6 +892,49 @@ mod tests { ); } + #[tokio::test] + async fn antigravity_expired_legacy_credential_refreshes_and_normalizes_refresh_token() { + let mut transport = sample_transport(); + transport.provider.name = "Antigravity".to_string(); + transport.provider.provider_type = "antigravity".to_string(); + transport.key.decrypted_api_key = "stale-access-token".to_string(); + transport.key.expires_at_unix_secs = Some(1); + transport.key.decrypted_auth_config = Some( + json!({ + "provider_type": "antigravity", + "refreshToken": "stable-refresh-token", + "expires_at": 1, + }) + .to_string(), + ); + let hits = Arc::new(AtomicUsize::new(0)); + let executor = StaticTokenExecutor { + hits: Arc::clone(&hits), + }; + let adapter = GenericOAuthRefreshAdapter::default() + .with_token_url_for_tests("antigravity", "https://oauth.example/token"); + + assert!(adapter.supports(&transport)); + assert!(adapter.should_refresh(&transport, None)); + + let refreshed = adapter + .refresh(&executor, &transport, None) + .await + .expect("antigravity refresh should succeed") + .expect("antigravity refresh should return a cache entry"); + + assert_eq!(hits.load(Ordering::SeqCst), 1); + assert_eq!(refreshed.provider_type, "antigravity"); + assert_eq!(refreshed.auth_header_value, "Bearer fresh-access-token"); + assert_eq!( + refreshed + .metadata + .as_ref() + .map(|metadata| &metadata["refresh_token"]), + Some(&json!("stable-refresh-token")) + ); + } + #[tokio::test] async fn fenced_force_reuses_successor_when_refresh_token_does_not_rotate() { let mut transport = sample_transport(); diff --git a/crates/aether-runtime/state/src/lib.rs b/crates/aether-runtime/state/src/lib.rs index 5c5915b5a..9a2fef67e 100644 --- a/crates/aether-runtime/state/src/lib.rs +++ b/crates/aether-runtime/state/src/lib.rs @@ -717,7 +717,23 @@ impl RuntimeState { limit: usize, config: RuntimeSemaphoreConfig, ) -> Result { - RuntimeSemaphore::new(self.clone(), gate, limit, config) + RuntimeSemaphore::new(self.clone(), gate, None, limit, config) + } + + pub fn keyed_semaphore( + &self, + gate: &'static str, + resource_key: &str, + limit: usize, + config: RuntimeSemaphoreConfig, + ) -> Result { + let resource_key = resource_key.trim(); + if resource_key.is_empty() { + return Err(RuntimeSemaphoreError::InvalidConfiguration( + "runtime semaphore resource key cannot be empty".to_string(), + )); + } + RuntimeSemaphore::new(self.clone(), gate, Some(resource_key), limit, config) } } @@ -1202,6 +1218,7 @@ impl RuntimeSemaphore { fn new( runtime: RuntimeState, gate: &'static str, + resource_key: Option<&str>, limit: usize, config: RuntimeSemaphoreConfig, ) -> Result { @@ -1222,7 +1239,9 @@ impl RuntimeSemaphore { } Ok(Self { state: Arc::new(RuntimeSemaphoreState { - key: format!("admission:{gate}"), + key: resource_key + .map(|resource_key| format!("admission:{gate}:{resource_key}")) + .unwrap_or_else(|| format!("admission:{gate}")), runtime, gate, limit, @@ -1256,6 +1275,7 @@ pub struct RuntimeSemaphorePermit { token: String, renew_task: JoinHandle<()>, healthy: Arc, + released: bool, } impl aether_runtime::AdmissionPermitHealth for RuntimeSemaphorePermit { @@ -1264,8 +1284,22 @@ impl aether_runtime::AdmissionPermitHealth for RuntimeSemaphorePermit { } } +impl RuntimeSemaphorePermit { + pub async fn release(mut self) -> Result<(), RuntimeSemaphoreError> { + self.renew_task.abort(); + let result = self.state.release(&self.token).await; + if result.is_ok() { + self.released = true; + } + result + } +} + impl Drop for RuntimeSemaphorePermit { fn drop(&mut self) { + if self.released { + return; + } self.renew_task.abort(); let state = Arc::clone(&self.state); let token = self.token.clone(); @@ -1331,6 +1365,7 @@ impl RuntimeSemaphoreState { token, renew_task, healthy, + released: false, }) } @@ -1733,6 +1768,43 @@ mod tests { assert_eq!(gate.snapshot().await.expect("snapshot").in_flight, 0); } + #[tokio::test] + async fn memory_keyed_semaphores_isolate_resource_capacity() { + let runtime = RuntimeState::memory(MemoryRuntimeStateConfig::default()); + let first = runtime + .keyed_semaphore( + "provider_key", + "key-a", + 1, + RuntimeSemaphoreConfig::default(), + ) + .expect("first gate should build"); + let second = runtime + .keyed_semaphore( + "provider_key", + "key-b", + 1, + RuntimeSemaphoreConfig::default(), + ) + .expect("second gate should build"); + let permit = first.try_acquire().await.expect("first permit"); + + assert!(matches!( + first + .try_acquire() + .await + .expect_err("same key should saturate"), + RuntimeSemaphoreError::Saturated { .. } + )); + let second_permit = second + .try_acquire() + .await + .expect("different key should retain independent capacity"); + + drop(second_permit); + drop(permit); + } + #[tokio::test(flavor = "current_thread")] async fn memory_semaphore_marks_permit_unhealthy_after_lease_loss() { let runtime = RuntimeState::memory(MemoryRuntimeStateConfig::default()); diff --git a/frontend/src/api/endpoints/pool.ts b/frontend/src/api/endpoints/pool.ts index 82621013e..9557c6328 100644 --- a/frontend/src/api/endpoints/pool.ts +++ b/frontend/src/api/endpoints/pool.ts @@ -152,6 +152,7 @@ export interface PoolKeyDetail { rate_multipliers?: Record | null internal_priority?: number rpm_limit?: number | null + concurrent_limit?: number | null cache_ttl_minutes?: number max_probe_interval_minutes?: number note?: string | null diff --git a/frontend/src/features/pool/components/PoolKeyQuotaPanel.vue b/frontend/src/features/pool/components/PoolKeyQuotaPanel.vue index e0bc29899..829a25aad 100644 --- a/frontend/src/features/pool/components/PoolKeyQuotaPanel.vue +++ b/frontend/src/features/pool/components/PoolKeyQuotaPanel.vue @@ -20,6 +20,7 @@ > {{ accountQuotaText }} +
{{ accountQuotaText }}
+ (), { accountQuotaText: null, fallbackText: null, textClass: '', variant: 'desktop', + resetCreditText: null, + resetCreditItems: () => [], + canConsumeResetCredit: false, + consumingResetCredit: false, }) +const emit = defineEmits<{ + 'consume-reset-credit': [] +}>() + const { legacyT } = useI18n() +const ResetCredits = defineComponent({ + name: 'PoolQuotaResetCredits', + setup() { + return () => props.resetCreditText ? h('div', { + 'data-testid': 'pool-quota-reset-credits', + class: 'mt-2 border-t border-border/50 pt-1.5 text-[10px] leading-4 text-muted-foreground', + }, [ + h('div', { class: 'flex flex-wrap items-center gap-x-1' }, [ + props.canConsumeResetCredit + ? h('button', { + type: 'button', + disabled: props.consumingResetCredit, + class: 'font-medium text-primary hover:underline disabled:pointer-events-none disabled:opacity-60', + onClick: () => emit('consume-reset-credit'), + }, props.consumingResetCredit ? legacyT('重置中...') : legacyT('点击以进行重置')) + : null, + h('span', props.resetCreditText), + ]), + props.resetCreditItems.length + ? h('div', { class: 'truncate tabular-nums', title: props.resetCreditItems.join(' · ') }, props.resetCreditItems.join(' · ')) + : null, + ]) : null + }, +}) + const QuotaProgressRows = defineComponent({ name: 'QuotaProgressRows', props: { @@ -122,15 +162,24 @@ const QuotaProgressRows = defineComponent({ : null, ]), h('div', { class: 'flex items-center gap-1.5' }, [ - h('div', { class: 'relative flex-1 h-1.5 rounded-full bg-border overflow-hidden' }, [ - h('div', { - class: ['absolute left-0 top-0 h-full rounded-full transition-all duration-300', item.barClass], - style: { width: `${item.remainingPercent}%` }, - }), - ]), + item.numericOnly + ? null + : h('div', { + 'data-testid': 'pool-quota-progress-track', + class: 'relative flex-1 h-1.5 rounded-full bg-border overflow-hidden', + }, [ + h('div', { + class: ['absolute left-0 top-0 h-full rounded-full transition-all duration-300', item.barClass], + style: { width: `${item.remainingPercent}%` }, + }), + ]), h('span', { 'data-testid': 'pool-quota-meter-text', - class: ['shrink-0 text-[10px] font-medium tabular-nums leading-none', item.meterClass], + class: [ + 'shrink-0 text-[10px] font-medium tabular-nums leading-none', + item.numericOnly ? 'ml-auto' : '', + item.meterClass, + ], }, item.meterText), ]), ])) diff --git a/frontend/src/features/pool/components/PoolSchedulingDialog.vue b/frontend/src/features/pool/components/PoolSchedulingDialog.vue index 6edfee14a..1c1d7a2b3 100644 --- a/frontend/src/features/pool/components/PoolSchedulingDialog.vue +++ b/frontend/src/features/pool/components/PoolSchedulingDialog.vue @@ -48,6 +48,27 @@ > {{ activeDistributionDesc }}

+
+ +
@@ -267,13 +288,16 @@ const FALLBACK_PRESET_DEFS: PoolPresetMeta[] = [ { name: 'cache_affinity', label: '缓存亲和', - description: '优先复用最近使用过的 Key，利用 Prompt Caching', + description: '同一用户持续复用 Key，首次分配可集中或轮转', mutex_group: DISTRIBUTION_GROUP, - evidence_hint: '依据 LRU 时间戳（最近使用优先，与 LRU 轮转相反）', + evidence_hint: '先复用用户粘性 Key，未命中时按所选二级模式分配', providers: [], default_enabled: true, - modes: null, - default_mode: null, + modes: [ + { value: 'single_account', label: '单号优先' }, + { value: 'lru', label: 'LRU 轮号' }, + ], + default_mode: 'single_account', }, { name: 'lru', @@ -733,14 +757,17 @@ const activeDistributionPreset = computed(() => { return found?.item.preset ?? null }) -const activeDistributionDesc = computed(() => { +const activeDistributionItem = computed(() => { const found = distributionItems.value.find(({ item }) => item.enabled && item.applicable) - return found?.item.desc ?? null + return found?.item ?? null +}) + +const activeDistributionDesc = computed(() => { + return activeDistributionItem.value?.desc ?? null }) const activeDistributionLabel = computed(() => { - const found = distributionItems.value.find(({ item }) => item.enabled && item.applicable) - return found?.item.label ?? null + return activeDistributionItem.value?.label ?? null }) const strategyItems = computed(() => { diff --git a/frontend/src/features/pool/components/__tests__/PoolKeyDisplayPanels.spec.ts b/frontend/src/features/pool/components/__tests__/PoolKeyDisplayPanels.spec.ts index 2b43f6900..cb70b1729 100644 --- a/frontend/src/features/pool/components/__tests__/PoolKeyDisplayPanels.spec.ts +++ b/frontend/src/features/pool/components/__tests__/PoolKeyDisplayPanels.spec.ts @@ -90,6 +90,30 @@ describe('pool key display panels', () => { root.remove() }) + it('renders Antigravity quota as numeric values without progress tracks', () => { + const root = document.createElement('div') + document.body.appendChild(root) + const app = createApp(PoolKeyQuotaPanel, { + items: [{ + label: 'Gemini 3.1 Pro (High)', + remainingPercent: 42, + resetText: '1h 后重置', + meterText: '42', + barClass: 'bg-amber-500', + meterClass: 'text-amber-600', + numericOnly: true, + }], + }) + app.use(createI18n()) + app.mount(root) + + expect(root.querySelector('[data-testid="pool-quota-meter-text"]')?.textContent).toBe('42') + expect(root.querySelector('[data-testid="pool-quota-progress-track"]')).toBeNull() + + app.unmount() + root.remove() + }) + it('renders single-cycle stats as plain text', () => { const root = document.createElement('div') document.body.appendChild(root) diff --git a/frontend/src/features/pool/components/__tests__/PoolSchedulingDialog.cache-affinity.spec.ts b/frontend/src/features/pool/components/__tests__/PoolSchedulingDialog.cache-affinity.spec.ts new file mode 100644 index 000000000..7848fcd72 --- /dev/null +++ b/frontend/src/features/pool/components/__tests__/PoolSchedulingDialog.cache-affinity.spec.ts @@ -0,0 +1,125 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { createApp, defineComponent, h, nextTick, ref, type App } from 'vue' + +import PoolSchedulingDialog from '../PoolSchedulingDialog.vue' + +const endpointMocks = vi.hoisted(() => ({ + getPoolSchedulingPresets: vi.fn(), + getProvider: vi.fn(), + updateProvider: vi.fn(), +})) + +vi.mock('@/api/endpoints', () => ({ + getProvider: endpointMocks.getProvider, + updateProvider: endpointMocks.updateProvider, +})) + +vi.mock('@/api/endpoints/pool', () => ({ + getPoolSchedulingPresets: endpointMocks.getPoolSchedulingPresets, +})) + +vi.mock('@/composables/useToast', () => ({ + useToast: () => ({ + success: vi.fn(), + error: vi.fn(), + }), +})) + +const mountedApps: Array<{ app: App, root: HTMLElement }> = [] +const cacheAffinityConfig = { + scheduling_presets: [{ + preset: 'cache_affinity', + enabled: true, + mode: 'single_account', + }], +} + +async function settle(): Promise { + for (let index = 0; index < 4; index += 1) { + await Promise.resolve() + await nextTick() + } +} + +function mountDialog(): void { + const root = document.createElement('div') + document.body.appendChild(root) + const TestHost = defineComponent({ + setup() { + const open = ref(false) + void nextTick(() => { open.value = true }) + return () => h(PoolSchedulingDialog, { + modelValue: open.value, + providerId: 'provider-1', + providerType: 'openai', + currentConfig: cacheAffinityConfig, + 'onUpdate:modelValue': (value: boolean) => { open.value = value }, + }) + }, + }) + const app = createApp(TestHost) + app.mount(root) + mountedApps.push({ app, root }) +} + +beforeEach(() => { + endpointMocks.getPoolSchedulingPresets.mockReset() + endpointMocks.getProvider.mockReset() + endpointMocks.updateProvider.mockReset() + endpointMocks.getPoolSchedulingPresets.mockResolvedValue([{ + name: 'cache_affinity', + label: '缓存亲和', + description: '同一用户持续复用 Key，首次分配可集中或轮转', + providers: [], + default_enabled: true, + modes: [ + { value: 'single_account', label: '单号优先' }, + { value: 'lru', label: 'LRU 轮号' }, + ], + default_mode: 'single_account', + mutex_group: 'distribution_mode', + }]) + endpointMocks.getProvider.mockResolvedValue({ + id: 'provider-1', + pool_advanced: cacheAffinityConfig, + }) + endpointMocks.updateProvider.mockResolvedValue({ id: 'provider-1' }) +}) + +afterEach(() => { + for (const { app, root } of mountedApps.splice(0)) { + app.unmount() + root.remove() + } + document.body.innerHTML = '' +}) + +describe('PoolSchedulingDialog cache affinity modes', () => { + it('shows and saves the LRU secondary mode', async () => { + mountDialog() + await settle() + + const modeControl = document.body.querySelector( + '[data-testid="pool-cache-affinity-secondary-mode"]', + ) + expect(modeControl?.textContent).toContain('单号优先') + expect(modeControl?.textContent).toContain('LRU 轮号') + + modeControl?.querySelector('[data-mode="lru"]')?.click() + await nextTick() + const saveButton = [...document.body.querySelectorAll('button')] + .find(button => button.textContent?.trim() === '保存') + saveButton?.click() + await settle() + + expect(endpointMocks.updateProvider).toHaveBeenCalledWith('provider-1', { + pool_advanced: { + scheduling_presets: [{ + preset: 'cache_affinity', + enabled: true, + mode: 'lru', + }], + }, + }) + }) +}) \ No newline at end of file diff --git a/frontend/src/features/pool/utils/__tests__/poolKeyBatchSettings.spec.ts b/frontend/src/features/pool/utils/__tests__/poolKeyBatchSettings.spec.ts index 1a7389b0d..fac763879 100644 --- a/frontend/src/features/pool/utils/__tests__/poolKeyBatchSettings.spec.ts +++ b/frontend/src/features/pool/utils/__tests__/poolKeyBatchSettings.spec.ts @@ -35,6 +35,18 @@ describe('pool key batch settings', () => { }) }) + it('serializes a positive concurrency limit as a number', () => { + const selection = createPoolKeyBatchSettingSelection() + const draft = createPoolKeyBatchSettingsDraft() + selection.concurrent_limit = true + draft.concurrent_limit = 6 + + expect(validatePoolKeyBatchSettings(selection, draft)).toEqual([]) + expect(buildPoolKeySettingsPatch(selection, draft)).toEqual({ + concurrent_limit: 6, + }) + }) + it('requires a selected field and a proxy node for set mode', () => { const selection = createPoolKeyBatchSettingSelection() const draft = createPoolKeyBatchSettingsDraft() diff --git a/frontend/src/features/pool/utils/poolQuotaRefresh.ts b/frontend/src/features/pool/utils/poolQuotaRefresh.ts index 175084e2a..80a87f9fa 100644 --- a/frontend/src/features/pool/utils/poolQuotaRefresh.ts +++ b/frontend/src/features/pool/utils/poolQuotaRefresh.ts @@ -5,24 +5,34 @@ export function mergePoolKeyQuotaSnapshots( keys: PoolKeyDetail[], results: RefreshQuotaResult['results'], ): PoolKeyDetail[] { - const quotaByKeyId = new Map>() + const resultByKeyId = new Map() for (const result of results) { - if (result.quota_snapshot) { - quotaByKeyId.set(result.key_id, result.quota_snapshot) + if (result.quota_snapshot || result.metadata) { + resultByKeyId.set(result.key_id, result) } } - if (quotaByKeyId.size === 0) return keys + if (resultByKeyId.size === 0) return keys return keys.map((key) => { - const quotaSnapshot = quotaByKeyId.get(key.key_id) - if (!quotaSnapshot) return key + const result = resultByKeyId.get(key.key_id) + if (!result) return key + const quotaSnapshot = result.quota_snapshot + const providerType = String(quotaSnapshot?.provider_type || key.provider_type || '').trim().toLowerCase() return { ...key, - quota_updated_at: quotaSnapshot.updated_at ?? quotaSnapshot.observed_at ?? key.quota_updated_at ?? null, - status_snapshot: { - ...(key.status_snapshot ?? {}), - quota: quotaSnapshot, - }, + ...(quotaSnapshot ? { + quota_updated_at: quotaSnapshot.updated_at ?? quotaSnapshot.observed_at ?? key.quota_updated_at ?? null, + status_snapshot: { + ...(key.status_snapshot ?? {}), + quota: quotaSnapshot, + }, + } : {}), + ...(result.metadata && providerType ? { + upstream_metadata: { + ...(key.upstream_metadata ?? {}), + [providerType]: result.metadata, + }, + } : {}), } }) } diff --git a/frontend/src/views/admin/PoolManagement.vue b/frontend/src/views/admin/PoolManagement.vue index 29b4aad2b..de4d72634 100644 --- a/frontend/src/views/admin/PoolManagement.vue +++ b/frontend/src/views/admin/PoolManagement.vue @@ -331,6 +331,11 @@ :account-quota-text="keyUiStateMap[key.key_id]?.accountQuotaText" :fallback-text="keyUiStateMap[key.key_id]?.quotaFallbackText" :text-class="keyUiStateMap[key.key_id]?.quotaTextClass || ''" + :reset-credit-text="getCodexResetCreditCountText(key)" + :reset-credit-items="getCodexResetCreditItemTexts(key)" + :can-consume-reset-credit="canConsumeCodexResetCredit(key)" + :consuming-reset-credit="consumingCodexResetCreditKeyId === key.key_id" + @consume-reset-credit="handleConsumeCodexResetCredit(key)" /> @@ -709,7 +714,12 @@ :account-quota-text="keyUiStateMap[key.key_id]?.accountQuotaText" :fallback-text="keyUiStateMap[key.key_id]?.quotaFallbackText" :text-class="keyUiStateMap[key.key_id]?.quotaTextClass || ''" + :reset-credit-text="getCodexResetCreditCountText(key)" + :reset-credit-items="getCodexResetCreditItemTexts(key)" + :can-consume-reset-credit="canConsumeCodexResetCredit(key)" + :consuming-reset-credit="consumingCodexResetCreditKeyId === key.key_id" variant="mobile" + @consume-reset-credit="handleConsumeCodexResetCredit(key)" />
@@ -1074,6 +1084,7 @@ import { deleteEndpointKey, updateProviderKey, refreshProviderQuota, + consumeCodexResetCredit, resetProviderKeyCycleStats, } from '@/api/endpoints/keys' import { refreshProviderOAuth } from '@/api/endpoints/provider_oauth' @@ -1132,6 +1143,22 @@ import { } from '@/features/pool/utils/poolStatsDisplay' import { resetCodexCycleUsageWindows } from '@/features/pool/utils/poolCycleStats' import { mergePoolKeyQuotaSnapshots } from '@/features/pool/utils/poolQuotaRefresh' +import { + dedupeAntigravityQuotaItemsByLabel, + resolveAntigravityQuotaLabel, +} from '@/features/providers/utils/antigravityQuota' +import { + clearPendingCodexResetCreditIdempotencyKey, + clearPendingCodexResetCreditIdempotencyKeyForOutcome, + createCodexResetCreditIdempotencyKey, + formatCodexResetCreditCount, + formatCodexResetCreditExpiresAt, + getCodexResetCreditAvailableCount, + getCodexResetCreditReservationIdempotencyKey, + getVisibleCodexResetCreditItems, + readPendingCodexResetCreditIdempotencyKey, + rememberPendingCodexResetCreditIdempotencyKey, +} from '@/features/providers/components/codex-reset-credit-display' import { getCodexQuotaWindowPresentation } from '@/utils/codexQuotaWindow' import { getOAuthOrgBadge } from '@/utils/oauthIdentity' import { formatOAuthPlanType, getOAuthPlanTypeClass } from '@/utils/oauthPlanType' @@ -1853,6 +1880,7 @@ const sortOrder = ref(restoredViewState.sortOrder) const hasPoolKeyFilters = computed(() => searchQuery.value.trim().length > 0 || statusFilter.value !== 'all') const MANUAL_QUOTA_REFRESH_COOLDOWN_SECONDS = 5 * 60 const refreshingOAuthKeyId = ref(null) +const consumingCodexResetCreditKeyId = ref(null) const resettingCycleKeyId = ref(null) const savingProxyKeyId = ref(null) const proxyDesktopPopoverOpenKeyId = ref(null) @@ -1978,6 +2006,7 @@ watch( interface QuotaProgressItem { label: string remainingPercent: number + numericOnly?: boolean sortOrder?: number detail?: string resetAtSeconds?: number | null @@ -1993,6 +2022,7 @@ interface QuotaProgressDisplayItem { meterText: string barClass: string meterClass: string + numericOnly?: boolean } type PoolKeyUiState = { @@ -2033,9 +2063,12 @@ const quotaProgressDisplayMap = computed `${item.displayKey} ${formatCodexResetCreditExpiresAt(item.expiresAt)}`) +} + +function canConsumeCodexResetCredit(key: PoolKeyDetail): boolean { + return getQuotaSnapshotProviderType(key) === 'codex' + && getCodexCredentialGeneration(key) !== undefined + && (getPendingCodexResetCreditIdempotencyKey(key) !== null + || (getCodexResetCreditAvailableCount(getCodexResetCredits(key)) ?? 0) > 0) + && consumingCodexResetCreditKeyId.value === null +} + +async function handleConsumeCodexResetCredit(key: PoolKeyDetail): Promise { + if (!canConsumeCodexResetCredit(key)) return + const generation = getCodexCredentialGeneration(key) + if (generation === undefined) return + const pendingIdempotencyKey = getPendingCodexResetCreditIdempotencyKey(key) + const confirmed = await confirm({ + title: '确认使用 Codex 重置机会', + message: pendingIdempotencyKey + ? '将继续确认上次尚未完成的 Codex 重置请求。' + : '将消耗 1 次 Codex 重置机会，完成后自动刷新账号额度。', + confirmText: '确认重置', + cancelText: '取消', + variant: 'warning', + }) + if (!confirmed) return + + consumingCodexResetCreditKeyId.value = key.key_id + try { + const idempotencyKey = pendingIdempotencyKey + || readPendingCodexResetCreditIdempotencyKey(key.key_id, generation) + || createCodexResetCreditIdempotencyKey() + rememberPendingCodexResetCreditIdempotencyKey(key.key_id, idempotencyKey, generation) + const result = await consumeCodexResetCredit(key.key_id, { + idempotency_key: idempotencyKey, + expected_credential_generation: generation, + }) + clearPendingCodexResetCreditIdempotencyKeyForOutcome(key.key_id, result.outcome) + keyPage.value.keys = mergePoolKeyQuotaSnapshots(keyPage.value.keys, [{ + key_id: result.key_id, + key_name: key.key_name, + status: result.refresh_status === 'success' ? 'success' : result.status as 'success', + metadata: result.metadata, + quota_snapshot: result.quota_snapshot, + }]) + if (result.outcome === 'reset' || result.outcome === 'already_redeemed') { + success('Codex 重置机会已使用，账号额度已刷新') + } else { + showWarning(result.message || '重置请求已处理，请查看最新额度') + } + } catch (err: unknown) { + const responseData = typeof err === 'object' && err !== null && 'response' in err + ? (err as { response?: { data?: Record } }).response?.data + : undefined + if (responseData?.outcome === 'credential_changed') { + clearPendingCodexResetCreditIdempotencyKey(key.key_id) + } else if (typeof responseData?.active_idempotency_key === 'string') { + rememberPendingCodexResetCreditIdempotencyKey( + key.key_id, + responseData.active_idempotency_key, + generation, + ) + } + showError(parseApiError(err, 'Codex 重置机会使用失败')) + } finally { + consumingCodexResetCreditKeyId.value = null + } +} + function normalizeQuotaUpdatedAt(raw: number | null | undefined): number | null { const value = Number(raw ?? 0) if (!Number.isFinite(value) || value <= 0) return null @@ -2447,6 +2581,7 @@ function toEndpointApiKey(key: PoolKeyDetail): EndpointAPIKey { rate_multipliers: key.rate_multipliers ?? null, internal_priority: key.internal_priority ?? 50, rpm_limit: key.rpm_limit ?? null, + concurrent_limit: key.concurrent_limit ?? null, allowed_models: key.allowed_models ?? null, capabilities: key.capabilities ?? null, cache_ttl_minutes: key.cache_ttl_minutes ?? 5, @@ -3730,20 +3865,24 @@ function buildQuotaProgressItemsFromSnapshot(key: PoolKeyDetail): QuotaProgressI if (providerType === 'antigravity') { const windows = getQuotaSnapshotWindowsByScope(quota, 'model') if (windows.length === 0) return [] - - const remainingPercents = windows - .map(getQuotaWindowRemainingPercent) - .filter((value): value is number => value != null) - if (remainingPercents.length === 0) return [] - - return [{ - label: '最低', - remainingPercent: Math.min(...remainingPercents), - detail: `${windows.length} 模型`, - resetAtSeconds: null, - resetSeconds: null, - updatedAtSeconds: getQuotaSnapshotUpdatedAtSeconds(quota), - }] + const opaqueDisplayIndex = { value: 1 } + return dedupeAntigravityQuotaItemsByLabel(windows + .map((window): (QuotaProgressItem & { model: string, resetSeconds: number | null }) | null => { + const remainingPercent = getQuotaWindowRemainingPercent(window) + if (remainingPercent == null) return null + const model = String(window.model || window.code || '').trim().replace(/^model:/i, '') + return { + model, + label: resolveAntigravityQuotaLabel(model, window.label, opaqueDisplayIndex), + remainingPercent, + numericOnly: true, + resetAtSeconds: normalizeUnixSeconds(window.reset_at ?? quota.reset_at ?? null), + resetSeconds: normalizeRemainingSeconds(window.reset_seconds ?? quota.reset_seconds ?? null), + updatedAtSeconds: getQuotaSnapshotUpdatedAtSeconds(quota), + allowDynamicReset: true, + } + }) + .filter((item): item is QuotaProgressItem & { model: string, resetSeconds: number | null } => item != null)) } if (providerType === 'gemini_cli') {