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

Merge pull request #767 from zhefox/fix/provider-key-concurrency-cache-affinity

Browse Source 
fix(gateway): improve provider pool concurrency, quotas, and affinity
This commit is contained in:
ZheFox
2026-09-01 08:11:46 +08:00
committed by GitHub
parent 9372d6cfa5 9631b229b3
commit 9b819169d5
26 changed files with 1176 additions and 154 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/control/auth/gate.rs
+34 -5
View File
@@ -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(
apps/aether-gateway/src/dispatch/pool_scheduler.rs
+50 -11
View File
@@ -1826,17 +1826,17 @@ fn pool_key_candidate_order_for_group(
})
.collect::<Vec<_>>();
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::<Vec<_>>();
.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(
apps/aether-gateway/src/execution_runtime/stream/execution.rs
+36 -16
View File
@@ -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(
apps/aether-gateway/src/execution_runtime/sync/execution.rs
+29 -9
View File
@@ -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,
apps/aether-gateway/src/handlers/admin/provider/pool/config.rs
+8
View File
@@ -167,6 +167,14 @@ fn parse_pool_score_rules(pool_advanced: &Map<String, Value>) -> PoolMemberScore
fn normalize_pool_preset_mode(preset: &str, raw_mode: Option<&Value>) -> Option<String> {
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",
apps/aether-gateway/src/handlers/admin/provider/pool_admin/payloads.rs
+1
View File
@@ -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),
apps/aether-gateway/src/handlers/proxy/websocket/responses/admission.rs
+12 -9
View File
@@ -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,
apps/aether-gateway/src/maintenance/runtime/oauth_token_refresh.rs
+48 -7
View File
@@ -256,12 +256,16 @@ fn auth_config_has_refresh_token(auth_config: Option<&str>) -> bool {
let Ok(value) = serde_json::from_str::<Value>(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() {
apps/aether-gateway/src/provider_key_auth.rs
+14 -5
View File
@@ -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",
apps/aether-gateway/src/provider_pool_demand.rs
+176 -11
View File
@@ -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<RuntimeSemaphorePermit>,
released: bool,
}
pub(crate) enum ProviderPoolInFlightAdmission {
Acquired(Option<ProviderPoolInFlightGuard>),
Saturated { limit: usize },
}
enum ProviderPoolInFlightGuardKind {
Disabled,
Local {
provider_id: String,
counter: Arc<AtomicUsize>,
@@ -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<ProviderPoolInFlightGuard> {
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<RuntimeState>,
provider_id: &str,
request_id: &str,
candidate_id: Option<&str>,
key_id: &str,
concurrent_limit: Option<usize>,
) -> Result<Option<ProviderPoolInFlightGuard>, 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<ProviderPoolInFlightAdmission, GatewayError> {
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());
apps/aether-gateway/src/state/oauth.rs
+89 -4
View File
@@ -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<InMemoryProviderCatalogReadRepository>, 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<InMemoryProviderCatalogReadRepository>, 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!({
apps/aether-gateway/src/tests/control/admin/pool.rs
+4
View File
@@ -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!(
apps/aether-gateway/src/wallet_runtime/access.rs
+25 -1
View File
@@ -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);
crates/aether-oauth/src/provider/providers/antigravity.rs
+53 -16
View File
@@ -38,7 +38,17 @@ impl ProviderOAuthAdapter for AntigravityProviderOAuthAdapter {
state: &str,
code_challenge: Option<&str>,
) -> Result<crate::core::OAuthAuthorizeResponse, crate::core::OAuthError> {
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<OAuthHttpResponse, crate::core::OAuthError> {
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<OAuthHttpResponse, crate::core::OAuthError> {
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::<BTreeMap<_, _>>();
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(),
crates/aether-oauth/src/provider/providers/generic.rs
+3 -4
View File
@@ -374,10 +374,9 @@ impl ProviderOAuthAdapter for GenericProviderOAuthAdapter {
ctx: &ProviderOAuthTransportContext,
account: &ProviderOAuthAccount,
) -> Result<ProviderOAuthTokenSet, OAuthError> {
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"))?;
crates/aether-provider/pool/src/presets.rs
+44 -8
View File
@@ -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 轮号"}
])
);
}
}
crates/aether-provider/transport/src/generic_oauth/mod.rs
+47 -4
View File
@@ -513,10 +513,10 @@ fn generic_provider_type(provider_type: &str) -> Option<&'static str> {
}
fn refresh_token_from_auth_config(auth_config: &Value) -> Option<String> {
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<String> {
@@ -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();
crates/aether-runtime/state/src/lib.rs
+74 -2
View File
@@ -717,7 +717,23 @@ impl RuntimeState {
limit: usize,
config: RuntimeSemaphoreConfig,
) -> Result<RuntimeSemaphore, RuntimeSemaphoreError> {
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<RuntimeSemaphore, RuntimeSemaphoreError> {
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<Self, RuntimeSemaphoreError> {
@@ -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<std::sync::atomic::AtomicBool>,
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());
frontend/src/api/endpoints/pool.ts
+1
View File
@@ -152,6 +152,7 @@ export interface PoolKeyDetail {
rate_multipliers?: Record<string, number> | null
internal_priority?: number
rpm_limit?: number | null
concurrent_limit?: number | null
cache_ttl_minutes?: number
max_probe_interval_minutes?: number
note?: string | null
frontend/src/features/pool/components/PoolKeyQuotaPanel.vue
+57 -8
View File
@@ -20,6 +20,7 @@
>
{{ accountQuotaText }}
</div>
<ResetCredits />
</div>
<div
v-else-if="accountQuotaText || fallbackText"
@@ -47,6 +48,7 @@
>
{{ accountQuotaText }}
</div>
<ResetCredits />
</div>
<span
v-else-if="accountQuotaText || fallbackText"
@@ -72,23 +74,61 @@ export interface PoolQuotaProgressDisplayItem {
meterText: string
barClass: string
meterClass: string
numericOnly?: boolean
}
withDefaults(defineProps<{
const props = withDefaults(defineProps<{
items: PoolQuotaProgressDisplayItem[]
accountQuotaText?: string | null
fallbackText?: string | null
textClass?: string
variant?: 'desktop' | 'mobile'
resetCreditText?: string | null
resetCreditItems?: string[]
canConsumeResetCredit?: boolean
consumingResetCredit?: boolean
}>(), {
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),
]),
]))
frontend/src/features/pool/components/PoolSchedulingDialog.vue
+35 -8
View File
@@ -48,6 +48,27 @@
>
{{ activeDistributionDesc }}
</p>
<div
v-if="activeDistributionItem?.modeOptions.length"
data-testid="pool-cache-affinity-secondary-mode"
class="mt-2 flex w-fit flex-wrap gap-1 rounded-lg bg-muted/50 p-1"
>
<button
v-for="modeOpt in activeDistributionItem.modeOptions"
:key="modeOpt.value"
type="button"
:data-mode="modeOpt.value"
class="rounded-md px-2.5 py-1 text-xs font-medium transition-all"
:class="[
activeDistributionItem.mode === modeOpt.value
? 'bg-primary text-primary-foreground shadow-sm'
: 'text-muted-foreground hover:bg-background/70 hover:text-foreground'
]"
@click="setPresetModeByPreset(activeDistributionItem.preset, modeOpt.value)"
>
{{ modeOpt.label }}
</button>
</div>
</div>
</div>
@@ -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(() => {
frontend/src/features/pool/components/__tests__/PoolKeyDisplayPanels.spec.ts
+24
View File
@@ -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)
frontend/src/features/pool/components/__tests__/PoolSchedulingDialog.cache-affinity.spec.ts
+125
View File
@@ -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<void> {
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<HTMLButtonElement>('[data-mode="lru"]')?.click()
await nextTick()
const saveButton = [...document.body.querySelectorAll<HTMLButtonElement>('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',
}],
},
})
})
})
frontend/src/features/pool/utils/__tests__/poolKeyBatchSettings.spec.ts
+12
View File
@@ -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()
frontend/src/features/pool/utils/poolQuotaRefresh.ts
+21 -11
View File
@@ -5,24 +5,34 @@ export function mergePoolKeyQuotaSnapshots(
keys: PoolKeyDetail[],
results: RefreshQuotaResult['results'],
): PoolKeyDetail[] {
const quotaByKeyId = new Map<string, NonNullable<RefreshQuotaResult['results'][number]['quota_snapshot']>>()
const resultByKeyId = new Map<string, RefreshQuotaResult['results'][number]>()
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,
},
} : {}),
}
})
}
frontend/src/views/admin/PoolManagement.vue
+154 -15
View File
@@ -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)"
/>
</TableCell>
<TableCell class="py-3 px-2 align-middle">
@@ -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)"
/>
<div class="flex items-center gap-0.5">
@@ -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<PoolManagementSortOrder>(restoredViewState.sortOrder)
const hasPoolKeyFilters = computed(() => searchQuery.value.trim().length > 0 || statusFilter.value !== 'all')
const MANUAL_QUOTA_REFRESH_COOLDOWN_SECONDS = 5 * 60
const refreshingOAuthKeyId = ref<string | null>(null)
const consumingCodexResetCreditKeyId = ref<string | null>(null)
const resettingCycleKeyId = ref<string | null>(null)
const savingProxyKeyId = ref<string | null>(null)
const proxyDesktopPopoverOpenKeyId = ref<string | null>(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<Record<string, QuotaProgressDisplayItem
label: getQuotaProgressLabel(item.label),
remainingPercent: item.remainingPercent,
resetText: getQuotaProgressResetDisplayText(item),
meterText: getQuotaProgressMeterDisplayText(item),
meterText: item.numericOnly
? formatQuotaValue(item.remainingPercent)
: getQuotaProgressMeterDisplayText(item),
barClass: getQuotaRemainingBarColorByRemaining(item.remainingPercent),
meterClass: getQuotaRemainingClassByRemaining(item.remainingPercent),
numericOnly: item.numericOnly,
}))
}
return map
@@ -2204,6 +2237,107 @@ function applyQuotaRefreshResultToCurrentPage(result: Awaited<ReturnType<typeof
keyPage.value.keys = mergePoolKeyQuotaSnapshots(keyPage.value.keys, result.results)
}
function getCodexResetCredits(key: PoolKeyDetail) {
return getQuotaSnapshotProviderType(key) === 'codex'
? key.status_snapshot?.quota?.reset_credits ?? null
: null
}
function getCodexCredentialGeneration(key: PoolKeyDetail): string | null | undefined {
const codex = key.upstream_metadata?.codex
return codex && typeof codex === 'object'
? codex.credential_generation?.trim() || null
: undefined
}
function getPendingCodexResetCreditIdempotencyKey(key: PoolKeyDetail): string | null {
const codex = key.upstream_metadata?.codex
const serverReservation = getCodexResetCreditReservationIdempotencyKey(codex)
if (serverReservation) return serverReservation
const generation = getCodexCredentialGeneration(key)
return generation === undefined
? null
: readPendingCodexResetCreditIdempotencyKey(key.key_id, generation)
}
function getCodexResetCreditCountText(key: PoolKeyDetail): string | null {
const count = getCodexResetCreditAvailableCount(getCodexResetCredits(key))
return count === null && !getPendingCodexResetCreditIdempotencyKey(key)
? null
: formatCodexResetCreditCount(count)
}
function getCodexResetCreditItemTexts(key: PoolKeyDetail): string[] {
return getVisibleCodexResetCreditItems(getCodexResetCredits(key), undefined, 3)
.map(item => `${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<void> {
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<string, unknown> } }).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') {