From dbbe7b22abf2d8aa99c7f0739f9d58522999b2b0 Mon Sep 17 00:00:00 2001 From: zhefox Date: Wed, 2 Sep 2026 15:23:23 +0800 Subject: [PATCH 01/13] fix(pool): isolate dynamic model quota buckets and 429 scheduling --- .../src/ai_serving/finalize/tests_sync.rs | 9 +- .../aether-gateway/src/ai_serving/pure/mod.rs | 1 + .../src/dispatch/pool_scheduler.rs | 134 +++++- .../src/execution_runtime/grok.rs | 9 +- .../admin/provider/pool/runtime/writes.rs | 54 ++- .../src/handlers/proxy/finalize.rs | 12 + apps/aether-gateway/src/handlers/proxy/mod.rs | 5 + .../src/handlers/proxy/websocket/live/http.rs | 5 + .../src/handlers/shared/catalog.rs | 13 +- .../maintenance/runtime/pool_quota_probe.rs | 52 +++ .../src/orchestration/effects.rs | 37 ++ .../src/scheduler/candidate/runtime.rs | 33 +- .../ai_execute/finalize_local_cli/compact.rs | 3 +- .../ai_execute/finalize_local_cli/direct.rs | 3 +- crates/aether-admin/src/provider/pool.rs | 24 + crates/aether-admin/src/provider/quota.rs | 153 +++++++ crates/aether-ai/formats/src/api.rs | 1 + .../src/formats/conversion/response.rs | 4 +- .../formats/src/formats/openai/chat/stream.rs | 12 +- .../src/formats/openai/request_contract.rs | 40 ++ .../src/formats/openai/responses/mod.rs | 102 ++++- .../src/formats/shared/sync_products.rs | 27 +- crates/aether-ai/formats/src/lib.rs | 4 +- .../formats/src/protocol/canonical.rs | 7 +- .../src/lifecycle/backfill/postgres.rs | 2 + .../control/src/public/request_context.rs | 2 + crates/aether-pool-core/src/scheduler.rs | 36 +- crates/aether-provider/pool/src/lib.rs | 284 +++++++++++- crates/aether-provider/pool/src/provider.rs | 13 +- .../pool/src/providers/antigravity.rs | 17 + .../pool/src/providers/chatgpt_web.rs | 8 +- .../pool/src/providers/codex.rs | 14 +- .../pool/src/providers/grok.rs | 8 +- .../pool/src/providers/kiro.rs | 9 +- .../pool/src/providers/windsurf.rs | 8 +- crates/aether-provider/pool/src/quota.rs | 422 +++++++++++++++++- crates/aether-provider/pool/src/service.rs | 2 + 37 files changed, 1491 insertions(+), 78 deletions(-) diff --git a/apps/aether-gateway/src/ai_serving/finalize/tests_sync.rs b/apps/aether-gateway/src/ai_serving/finalize/tests_sync.rs index 8b346e5a6..6c4acaac5 100644 --- a/apps/aether-gateway/src/ai_serving/finalize/tests_sync.rs +++ b/apps/aether-gateway/src/ai_serving/finalize/tests_sync.rs @@ -1,5 +1,6 @@ use std::collections::BTreeMap; +use aether_ai_formats::openai_responses_message_item_id; use axum::body::to_bytes; use base64::Engine as _; use serde_json::json; @@ -192,7 +193,7 @@ fn aggregates_openai_responses_stream_completed_event_to_final_response() { "output_text": "Hello", "output": [{ "type": "message", - "id": "resp_123_msg", + "id": openai_responses_message_item_id("resp_123", 0), "role": "assistant", "status": "completed", "content": [{ @@ -977,7 +978,7 @@ fn converts_gemini_cli_response_to_openai_responses_response() { "output_text": "Hello Gemini CLI", "output": [{ "type": "message", - "id": "resp_cli_123_msg", + "id": openai_responses_message_item_id("resp_cli_123", 0), "role": "assistant", "status": "completed", "content": [{ @@ -1046,7 +1047,7 @@ fn converts_gemini_cli_function_call_to_openai_responses_function_call() { "output": [ { "type": "message", - "id": "resp_cli_tool_123_msg", + "id": openai_responses_message_item_id("resp_cli_tool_123", 0), "role": "assistant", "status": "completed", "content": [{ @@ -1252,7 +1253,7 @@ fn local_finalize_handles_openai_responses_openai_family_sync_response_even_when "model": "gpt-5", "output": [{ "type": "message", - "id": "resp_cli_family_123_msg", + "id": openai_responses_message_item_id("resp_cli_family_123", 0), "role": "assistant", "status": "completed", "content": [{ diff --git a/apps/aether-gateway/src/ai_serving/pure/mod.rs b/apps/aether-gateway/src/ai_serving/pure/mod.rs index 5a406a737..1808b7887 100644 --- a/apps/aether-gateway/src/ai_serving/pure/mod.rs +++ b/apps/aether-gateway/src/ai_serving/pure/mod.rs @@ -177,6 +177,7 @@ pub(crate) use aether_ai_formats::{ api_format_defaults_to_client_error_failover, api_format_defaults_to_non_stream, api_format_permission_covers, codex_responses_lite_tool_is_client_executed, intersect_api_format_allowed_lists, is_embedding_api_format, is_rerank_api_format, + normalize_openai_responses_message_item_ids, openai_responses_message_item_id, openai_responses_request_operation, openai_responses_synthetic_reasoning_item_id, strip_incompatible_openai_responses_reasoning_items, strip_incompatible_openai_responses_reasoning_items_with_policy, ApiOperation, ClientSurface, diff --git a/apps/aether-gateway/src/dispatch/pool_scheduler.rs b/apps/aether-gateway/src/dispatch/pool_scheduler.rs index 8fe81fce5..3048ba331 100644 --- a/apps/aether-gateway/src/dispatch/pool_scheduler.rs +++ b/apps/aether-gateway/src/dispatch/pool_scheduler.rs @@ -105,6 +105,7 @@ async fn schedule_pool_page_candidates( candidates: Vec, sticky_session_token: Option<&str>, effective_pool_config: Option<&AdminProviderPoolConfig>, + provider_model_name: Option<&str>, ) -> ( Vec, Vec, @@ -128,7 +129,8 @@ async fn schedule_pool_page_candidates( entry.1.insert(candidate.candidate.key_id.clone()); } - let key_context_by_id = read_pool_catalog_key_contexts_by_id(state, &candidates).await; + let key_context_by_id = + read_pool_catalog_key_contexts_by_id(state, &candidates, provider_model_name).await; let mut runtime_by_provider = BTreeMap::new(); let mut pool_config_by_provider = BTreeMap::new(); @@ -316,7 +318,9 @@ fn active_probe_member_is_unschedulable_for_request( }) { return true; } - key_context.is_some_and(|context| context.account_blocked || context.quota_exhausted) + key_context.is_some_and(|context| { + context.account_blocked || context.quota_exhausted || context.quota_hard_blocked + }) } async fn expand_pool_group_candidate( @@ -1034,6 +1038,7 @@ impl<'a> PoolKeyCursor<'a> { candidates, self.sticky_session_token.as_deref(), self.effective_pool_config.as_ref(), + Some(self.group.candidate.selected_provider_model_name.as_str()), ) .await; self.record_skipped_candidates(&skipped); @@ -1406,6 +1411,7 @@ fn pool_candidate_from_catalog_key( async fn read_pool_catalog_key_contexts_by_id( state: PlannerAppState<'_>, candidates: &[EligibleLocalExecutionCandidate], + provider_model_name: Option<&str>, ) -> BTreeMap { let mut key_ids = Vec::new(); let mut provider_type_by_key_id = BTreeMap::::new(); @@ -1437,13 +1443,30 @@ async fn read_pool_catalog_key_contexts_by_id( key_count = key_ids.len(), "gateway pool scheduler: failed to read catalog key metadata" ); - return BTreeMap::new(); + // Do not fail open when the quota metadata read is unavailable. A + // missing context must never turn an exhausted account into an + // eligible candidate and produce another upstream 429. The caller + // treats this marker as a pool quota skip and the next request will + // retry the metadata read. + return key_ids + .into_iter() + .map(|key_id| { + ( + key_id, + PoolCatalogKeyContext { + quota_hard_blocked: true, + ..PoolCatalogKeyContext::default() + }, + ) + }) + .collect(); } }; let provider_pool_service = ProviderPoolService::with_builtin_adapters(); - keys.into_iter() + let mut contexts = keys + .into_iter() .map(|key| { let provider_type = provider_type_by_key_id .get(&key.id) @@ -1451,10 +1474,28 @@ async fn read_pool_catalog_key_contexts_by_id( .unwrap_or_default(); ( key.id.clone(), - build_pool_catalog_key_context(state, &provider_pool_service, &key, provider_type), + build_pool_catalog_key_context( + state, + &provider_pool_service, + &key, + provider_type, + provider_model_name, + ), ) }) - .collect() + .collect::>(); + // A key can disappear between the candidate-row and catalog reads. Keep + // the snapshot non-empty and fail closed for those IDs so the caller does + // not interpret an incomplete read as "all accounts are healthy". + for key_id in key_ids { + contexts + .entry(key_id) + .or_insert_with(|| PoolCatalogKeyContext { + quota_hard_blocked: true, + ..PoolCatalogKeyContext::default() + }); + } + contexts } fn build_pool_catalog_key_context( @@ -1462,6 +1503,7 @@ fn build_pool_catalog_key_context( provider_pool_service: &ProviderPoolService, key: &StoredProviderCatalogKey, provider_type: &str, + provider_model_name: Option<&str>, ) -> PoolCatalogKeyContext { let (health_score, _, _, _, _) = provider_key_health_summary(key); let health_score = key @@ -1480,8 +1522,12 @@ fn build_pool_catalog_key_context( .filter(|value| value.is_finite() && *value >= 0.0); let auth_config = parse_catalog_auth_config_json(state.app(), key); - let mut signals = - provider_pool_service.member_signals(provider_type, key, auth_config.as_ref()); + let mut signals = provider_pool_service.member_signals( + provider_type, + key, + auth_config.as_ref(), + provider_model_name, + ); signals.account_blocked |= admin_provider_pool_pure::admin_pool_key_is_known_banned(key); signals.account_blocked |= pool_key_requires_reauth_for_scheduling(key, current_unix_ms().saturating_div(1000)); @@ -1694,7 +1740,21 @@ fn run_local_execution_pool_scheduler_with_runtime_map( let key_context = key_context_by_id .get(&candidate.candidate.key_id) .cloned() - .unwrap_or_default(); + .unwrap_or_else(|| { + // An explicitly non-empty metadata snapshot should contain + // every catalog key in this page. If one disappeared between + // reads, fail closed for that key instead of sending traffic + // with an unknown quota state. Empty maps are retained for + // callers/tests that intentionally provide no runtime context. + if key_context_by_id.is_empty() { + PoolCatalogKeyContext::default() + } else { + PoolCatalogKeyContext { + quota_hard_blocked: true, + ..PoolCatalogKeyContext::default() + } + } + }); let admin_pool_config = effective_pool_config_by_provider .get(&candidate.candidate.provider_id) .cloned() @@ -1968,7 +2028,7 @@ mod tests { use aether_data_contracts::repository::provider_catalog::{ StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider, }; - use aether_pool_core::PoolSchedulingPreset; + use aether_pool_core::{PoolSchedulingPreset, POOL_ACCOUNT_EXHAUSTED_SKIP_REASON}; use aether_provider_pool::ProviderPoolService; use aether_provider_transport::snapshot::{ GatewayProviderTransportEndpoint, GatewayProviderTransportKey, @@ -2080,6 +2140,55 @@ mod tests { ); } + #[test] + fn pool_scheduler_skips_quota_exhausted_key_when_flag_is_false() { + let ready = sample_eligible_candidate( + "provider-pool", + "endpoint-1", + "key-ready", + 10, + Some(json!({ "pool_advanced": {} })), + ); + let exhausted = sample_eligible_candidate( + "provider-pool", + "endpoint-1", + "key-exhausted", + 10, + Some(json!({ "pool_advanced": { "skip_exhausted_accounts": false } })), + ); + let key_context_by_id = BTreeMap::from([ + ("key-ready".to_string(), PoolCatalogKeyContext::default()), + ( + "key-exhausted".to_string(), + PoolCatalogKeyContext { + quota_exhausted: true, + ..PoolCatalogKeyContext::default() + }, + ), + ]); + + let (scheduled, skipped) = apply_local_execution_pool_scheduler_with_runtime_map( + vec![ready, exhausted], + &BTreeMap::new(), + &key_context_by_id, + ); + + assert_eq!( + scheduled + .iter() + .map(|item| item.candidate.key_id.as_str()) + .collect::>(), + vec!["key-ready"] + ); + assert_eq!( + skipped + .iter() + .map(|item| (item.candidate.key_id.as_str(), item.skip_reason)) + .collect::>(), + vec![("key-exhausted", POOL_ACCOUNT_EXHAUSTED_SKIP_REASON)] + ); + } + #[test] fn pool_scheduler_attaches_group_and_pool_metadata_to_ranked_candidates() { let pool_first = sample_eligible_candidate( @@ -4426,6 +4535,7 @@ mod tests { &ProviderPoolService::with_builtin_adapters(), &key, "codex", + None, ); assert_eq!(context.plan_tier.as_deref(), Some("team")); @@ -4471,6 +4581,7 @@ mod tests { &ProviderPoolService::with_builtin_adapters(), &key, "codex", + None, ); assert!(!context.quota_exhausted); @@ -4491,6 +4602,7 @@ mod tests { &ProviderPoolService::with_builtin_adapters(), &key, "codex", + None, ); assert!(context.quota_exhausted); @@ -4521,6 +4633,7 @@ mod tests { &ProviderPoolService::with_builtin_adapters(), &key, "antigravity", + None, ); assert!(context.quota_exhausted); @@ -4542,6 +4655,7 @@ mod tests { &ProviderPoolService::with_builtin_adapters(), &key, "codex", + None, ); assert!(context.account_blocked); diff --git a/apps/aether-gateway/src/execution_runtime/grok.rs b/apps/aether-gateway/src/execution_runtime/grok.rs index af27d5750..5106b7539 100644 --- a/apps/aether-gateway/src/execution_runtime/grok.rs +++ b/apps/aether-gateway/src/execution_runtime/grok.rs @@ -26,7 +26,9 @@ use crate::ai_serving::api::{ CanonicalContentPart, CanonicalStreamEvent, CanonicalStreamFrame, ClaudeClientEmitter, OpenAIChatClientEmitter, OpenAIResponsesClientEmitter, StreamingCanonicalUsage, }; -use crate::ai_serving::openai_responses_synthetic_reasoning_item_id; +use crate::ai_serving::{ + openai_responses_message_item_id, openai_responses_synthetic_reasoning_item_id, +}; use crate::clock::current_unix_secs; use crate::execution_runtime::ndjson::encode_stream_frame_ndjson; use crate::execution_runtime::transport::{ @@ -2725,7 +2727,7 @@ fn openai_responses_body( }; if !message_text.trim().is_empty() { output.push(json!({ - "id": format!("{response_id}_msg"), + "id": openai_responses_message_item_id(response_id.as_str(), output.len()), "type": "message", "role": "assistant", "content": [{"type": "output_text", "text": message_text, "annotations": []}], @@ -3774,6 +3776,9 @@ mod tests { ); assert_eq!(body["output"][0]["type"], serde_json::json!("reasoning")); assert_eq!(body["output"][1]["type"], serde_json::json!("message")); + assert!(body["output"][1]["id"] + .as_str() + .is_some_and(|id| id.starts_with("msg_"))); } #[test] diff --git a/apps/aether-gateway/src/handlers/admin/provider/pool/runtime/writes.rs b/apps/aether-gateway/src/handlers/admin/provider/pool/runtime/writes.rs index 94e7a3696..d25e675f8 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/pool/runtime/writes.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/pool/runtime/writes.rs @@ -502,14 +502,12 @@ pub(crate) async fn record_admin_provider_pool_error( .or_else(|| parse_google_quota_cooldown_seconds(error_body)), pool_config, ); - set_pool_cooldown( - runtime, - provider_id, - key_id, - "rate_limited_429", - ttl_seconds, - ) - .await; + let reason = if error_body_indicates_quota_exhaustion(error_body) { + "quota_exhausted_429" + } else { + "rate_limited_429" + }; + set_pool_cooldown(runtime, provider_id, key_id, reason, ttl_seconds).await; return; } @@ -548,6 +546,27 @@ pub(crate) async fn record_admin_provider_pool_error( } } +fn error_body_indicates_quota_exhaustion(error_body: Option<&str>) -> bool { + let body = error_body.unwrap_or_default().to_ascii_lowercase(); + [ + "quota exhausted", + "quota_exhausted", + "quota exceeded", + "quota_exceeded", + "insufficient_quota", + "resource exhausted", + "resource has been exhausted", + "resource_exhausted", + "usage_limit_reached", + "limit_reached", + "quota limit reached", + "credits exhausted", + "insufficient credits", + ] + .iter() + .any(|marker| body.contains(marker)) +} + pub(crate) async fn record_admin_provider_pool_stream_timeout( runtime: &RuntimeState, provider_id: &str, @@ -589,9 +608,10 @@ pub(crate) async fn record_admin_provider_pool_stream_timeout( #[cfg(test)] mod tests { use super::{ - admin_provider_pool_key_terminal_error_reason, parse_google_quota_cooldown_seconds_at, - record_admin_provider_pool_error, record_admin_provider_pool_stream_timeout, - record_admin_provider_pool_success, resolve_transient_cooldown_ttl, + admin_provider_pool_key_terminal_error_reason, error_body_indicates_quota_exhaustion, + parse_google_quota_cooldown_seconds_at, record_admin_provider_pool_error, + record_admin_provider_pool_stream_timeout, record_admin_provider_pool_success, + resolve_transient_cooldown_ttl, }; use crate::handlers::admin::provider::pool::runtime::reads::read_admin_provider_pool_runtime_state; use crate::handlers::admin::provider::shared::support::{ @@ -803,6 +823,16 @@ mod tests { assert_eq!(resolve_transient_cooldown_ttl(500, None, &pool_config), 0); } + #[test] + fn detects_quota_exhaustion_markers_in_429_bodies() { + assert!(error_body_indicates_quota_exhaustion(Some( + r#"{"error":{"status":"RESOURCE_EXHAUSTED","message":"quota exceeded"}}"#, + ))); + assert!(!error_body_indicates_quota_exhaustion(Some( + r#"{"error":{"message":"temporary rate limit"}}"#, + ))); + } + #[tokio::test] async fn success_feedback_writes_sticky_lru_cost_and_latency() { let Some(redis) = start_managed_redis_or_skip().await else { @@ -1110,7 +1140,7 @@ mod tests { .cooldown_reason_by_key .get("key-google-429") .map(String::as_str), - Some("rate_limited_429") + Some("quota_exhausted_429") ); assert!(runtime .cooldown_ttl_by_key diff --git a/apps/aether-gateway/src/handlers/proxy/finalize.rs b/apps/aether-gateway/src/handlers/proxy/finalize.rs index 9044d9db5..3ee89bb57 100644 --- a/apps/aether-gateway/src/handlers/proxy/finalize.rs +++ b/apps/aether-gateway/src/handlers/proxy/finalize.rs @@ -50,6 +50,7 @@ pub(super) fn finalize_gateway_response( mut response: Response, trace_id: &str, remote_addr: &std::net::SocketAddr, + client_ip: std::net::IpAddr, method: &http::Method, path_and_query: &str, control_decision: Option<&GatewayControlDecision>, @@ -117,6 +118,7 @@ pub(super) fn finalize_gateway_response( trace_id = %trace_id, request_id, remote_addr = %remote_addr, + client_ip = %client_ip, method = %method, path = %sanitized_path_and_query, user_id, @@ -137,6 +139,7 @@ pub(super) fn finalize_gateway_response( trace_id = %trace_id, request_id, remote_addr = %remote_addr, + client_ip = %client_ip, method = %method, path = %sanitized_path_and_query, user_id, @@ -157,6 +160,7 @@ pub(super) fn finalize_gateway_response( trace_id = %trace_id, request_id, remote_addr = %remote_addr, + client_ip = %client_ip, method = %method, path = %sanitized_path_and_query, user_id, @@ -247,11 +251,17 @@ pub(super) fn finalize_gateway_response_with_context( started_at: &Instant, request_permit: Option, ) -> Response { + let client_ip = request_context + .client_ip + .as_deref() + .and_then(|value| value.parse().ok()) + .unwrap_or_else(|| remote_addr.ip()); finalize_gateway_response( state, response, &request_context.trace_id, remote_addr, + client_ip, &request_context.request_method, &request_context.request_path_and_query(), request_context.control_decision.as_ref(), @@ -320,6 +330,7 @@ mod tests { request_query_string: None, request_content_type: Some("application/json".to_string()), host_header: None, + client_ip: None, control_decision: None, }; let mut headers = HeaderMap::new(); @@ -373,6 +384,7 @@ mod tests { response, "trace-finalize", &remote_addr, + remote_addr.ip(), &Method::GET, "/v1beta/models/gemini-3-flash-preview:generateContent?key=secret&alt=sse", Some(&control_decision), diff --git a/apps/aether-gateway/src/handlers/proxy/mod.rs b/apps/aether-gateway/src/handlers/proxy/mod.rs index 0fe8373dc..ab45a226a 100644 --- a/apps/aether-gateway/src/handlers/proxy/mod.rs +++ b/apps/aether-gateway/src/handlers/proxy/mod.rs @@ -993,6 +993,7 @@ async fn proxy_request_inner( response, &trace_id, &remote_addr, + client_ip, request.method(), request .uri() @@ -1029,6 +1030,7 @@ async fn proxy_request_inner( response, &trace_id, &remote_addr, + client_ip, request.method(), request .uri() @@ -1069,6 +1071,7 @@ async fn proxy_request_inner( response, &trace_id, &remote_addr, + client_ip, request.method(), request .uri() @@ -1128,6 +1131,7 @@ async fn proxy_request_inner( response, &trace_id, &remote_addr, + client_ip, &parts.method, parts .uri @@ -1149,6 +1153,7 @@ async fn proxy_request_inner( &trace_id, ) .await?; + request_context.client_ip = Some(client_ip.to_string()); maybe_promote_management_token_admin_principal( &state, client_ip, diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/live/http.rs b/apps/aether-gateway/src/handlers/proxy/websocket/live/http.rs index 2d8e46cb6..a55e5b3db 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/live/http.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/live/http.rs @@ -1001,6 +1001,7 @@ mod tests { request_query_string: None, request_content_type: None, host_header: None, + client_ip: None, control_decision: None, }; let (parts, _) = http::Request::builder() @@ -1036,6 +1037,7 @@ mod tests { request_query_string: None, request_content_type: Some("application/sdp".to_string()), host_header: None, + client_ip: None, control_decision: Some(decision), }; let (parts, _) = http::Request::builder() @@ -1074,6 +1076,7 @@ mod tests { request_query_string: Some("intent=quicksilver&architecture=avas".to_string()), request_content_type: Some("multipart/form-data".to_string()), host_header: None, + client_ip: None, control_decision: Some(decision), }; let (parts, _) = http::Request::builder() @@ -1128,6 +1131,7 @@ mod tests { request_query_string: None, request_content_type: None, host_header: None, + client_ip: None, control_decision: Some(decision), }; let (parts, _) = http::Request::builder() @@ -1194,6 +1198,7 @@ mod tests { request_query_string: None, request_content_type: Some("multipart/form-data".to_string()), host_header: None, + client_ip: None, control_decision: Some(decision), }; let (content_type, body) = build_live_multipart( diff --git a/apps/aether-gateway/src/handlers/shared/catalog.rs b/apps/aether-gateway/src/handlers/shared/catalog.rs index 394c04fff..81778657a 100644 --- a/apps/aether-gateway/src/handlers/shared/catalog.rs +++ b/apps/aether-gateway/src/handlers/shared/catalog.rs @@ -975,7 +975,7 @@ fn build_codex_quota_status_snapshot( .and_then(admin_provider_quota_pure::coerce_json_bool); let reset_credits = build_codex_reset_credits_status_snapshot(metadata, observed_at_unix_secs); - let windows = [ + let mut windows = [ codex_quota_window_snapshot(metadata, "primary", "weekly", "周", observed_at_unix_secs), codex_quota_window_snapshot(metadata, "secondary", "5h", "5H", observed_at_unix_secs), codex_quota_window_snapshot( @@ -996,6 +996,17 @@ fn build_codex_quota_status_snapshot( .into_iter() .flatten() .collect::>(); + if let Some(additional_windows) = metadata + .get("additional_quota_windows") + .and_then(Value::as_array) + { + windows.extend( + additional_windows + .iter() + .filter(|window| window.is_object()) + .cloned(), + ); + } if windows.is_empty() && plan_type.is_none() diff --git a/apps/aether-gateway/src/maintenance/runtime/pool_quota_probe.rs b/apps/aether-gateway/src/maintenance/runtime/pool_quota_probe.rs index e0f9d918d..694c57d06 100644 --- a/apps/aether-gateway/src/maintenance/runtime/pool_quota_probe.rs +++ b/apps/aether-gateway/src/maintenance/runtime/pool_quota_probe.rs @@ -1213,6 +1213,18 @@ fn probe_result_hard_state(item: &Value) -> Option { .unwrap_or_default() .trim() .to_ascii_lowercase(); + // Providers frequently encode an exhausted account as HTTP 429 with a + // provider-specific status/message (for example RESOURCE_EXHAUSTED or + // `quota exceeded`) rather than the normalized `quota_exhausted` status. + // Preserve that distinction in the score so the member stays out of the + // scheduler until a successful quota probe observes a reset. + let serialized_item = item.to_string().to_ascii_lowercase(); + let status_code = item.get("status_code").and_then(Value::as_u64); + if status == "quota_exhausted" + || (status_code == Some(429) && contains_quota_exhaustion_marker(&serialized_item)) + { + return Some(PoolMemberHardState::QuotaExhausted); + } match status.as_str() { "auth_invalid" | "forbidden" => Some(PoolMemberHardState::AuthInvalid), "workspace_deactivated" => Some(PoolMemberHardState::Banned), @@ -1226,6 +1238,26 @@ fn probe_result_hard_state(item: &Value) -> Option { } } +fn contains_quota_exhaustion_marker(value: &str) -> bool { + [ + "quota exhausted", + "quota_exhausted", + "quota exceeded", + "quota_exceeded", + "insufficient_quota", + "resource exhausted", + "resource has been exhausted", + "resource_exhausted", + "usage_limit_reached", + "limit_reached", + "quota limit reached", + "credits exhausted", + "insufficient credits", + ] + .iter() + .any(|marker| value.contains(marker)) +} + async fn perform_pool_quota_probe_for_provider( state: &AppState, admin_state: &AdminAppState<'_>, @@ -1771,6 +1803,26 @@ mod tests { assert_eq!(selected, vec!["never".to_string(), "old".to_string()]); } + #[test] + fn quota_markers_in_429_probe_results_are_hard_exhaustion() { + assert_eq!( + probe_result_hard_state(&json!({ + "status": "rate_limited", + "status_code": 429, + "message": "RESOURCE_EXHAUSTED: quota exceeded" + })), + Some(PoolMemberHardState::QuotaExhausted) + ); + assert_eq!( + probe_result_hard_state(&json!({ + "status": "rate_limited", + "status_code": 429, + "message": "temporary rate limit" + })), + Some(PoolMemberHardState::Cooldown) + ); + } + fn score( member_id: &str, hard_state: PoolMemberHardState, diff --git a/apps/aether-gateway/src/orchestration/effects.rs b/apps/aether-gateway/src/orchestration/effects.rs index d439a421b..278d22227 100644 --- a/apps/aether-gateway/src/orchestration/effects.rs +++ b/apps/aether-gateway/src/orchestration/effects.rs @@ -2060,6 +2060,15 @@ fn pool_score_hard_state_for_status( return Some(pool_score_hard_state_for_terminal_error_reason(&reason)); } + // A number of providers report account quota exhaustion as HTTP 429 rather + // than 402. Keep those members out of the score-based pool fallback until + // the provider's quota probe observes a reset; treating every 429 as a + // generic cooldown otherwise lets the member re-enter as soon as the short + // transient cooldown expires. + if status_code == 429 && error_body_indicates_quota_exhaustion(error_body) { + return Some(PoolMemberHardState::QuotaExhausted); + } + match status_code { 401 | 403 => Some(PoolMemberHardState::AuthInvalid), 402 => Some(PoolMemberHardState::QuotaExhausted), @@ -2082,6 +2091,27 @@ fn pool_score_hard_state_for_status( } } +fn error_body_indicates_quota_exhaustion(error_body: Option<&str>) -> bool { + let body = error_body.unwrap_or_default().to_ascii_lowercase(); + [ + "quota exhausted", + "quota_exhausted", + "quota exceeded", + "quota_exceeded", + "insufficient_quota", + "resource exhausted", + "resource has been exhausted", + "resource_exhausted", + "usage_limit_reached", + "limit_reached", + "quota limit reached", + "credits exhausted", + "insufficient credits", + ] + .iter() + .any(|marker| body.contains(marker)) +} + fn pool_score_hard_state_for_terminal_error_reason(reason: &str) -> PoolMemberHardState { if reason.starts_with("payment_required_") { PoolMemberHardState::QuotaExhausted @@ -3659,6 +3689,13 @@ mod tests { ), Some(PoolMemberHardState::QuotaExhausted) ); + assert_eq!( + pool_score_hard_state_for_status( + 429, + Some(r#"{"error":{"status":"RESOURCE_EXHAUSTED","message":"quota exhausted"}}"#), + ), + Some(PoolMemberHardState::QuotaExhausted) + ); } #[tokio::test] diff --git a/apps/aether-gateway/src/scheduler/candidate/runtime.rs b/apps/aether-gateway/src/scheduler/candidate/runtime.rs index 0227abd5d..25bb44260 100644 --- a/apps/aether-gateway/src/scheduler/candidate/runtime.rs +++ b/apps/aether-gateway/src/scheduler/candidate/runtime.rs @@ -341,18 +341,33 @@ fn read_key_account_quota_exhaustion_map( candidates .iter() .map(|candidate| { - let exhausted = provider_skip_exhausted_accounts - .get(candidate.provider_id.as_str()) - .copied() - .unwrap_or(false) - && provider_key_rpm_states - .get(candidate.key_id.as_str()) - .is_some_and(|key| { - admin_provider_pool_pure::admin_pool_key_account_quota_exhausted( + let exhausted = provider_key_rpm_states + .get(candidate.key_id.as_str()) + .is_some_and(|key| { + let account_exhausted = + admin_provider_pool_pure::admin_pool_key_model_quota_exhausted( key, candidate.provider_type.as_str(), + candidate.selected_provider_model_name.as_str(), ) - }); + .unwrap_or_else(|| { + admin_provider_pool_pure::admin_pool_key_account_quota_exhausted( + key, + candidate.provider_type.as_str(), + ) + }); + let hard_blocked = + admin_provider_pool_pure::admin_pool_key_model_quota_hard_blocked( + key, + candidate.provider_type.as_str(), + candidate.selected_provider_model_name.as_str(), + ); + let skip_configured = provider_skip_exhausted_accounts + .get(candidate.provider_id.as_str()) + .copied() + .unwrap_or(false); + hard_blocked || (skip_configured && account_exhausted) + }); (candidate.key_id.clone(), exhausted) }) .collect() diff --git a/apps/aether-gateway/src/tests/ai_execute/finalize_local_cli/compact.rs b/apps/aether-gateway/src/tests/ai_execute/finalize_local_cli/compact.rs index 2232af5ca..b6d665143 100644 --- a/apps/aether-gateway/src/tests/ai_execute/finalize_local_cli/compact.rs +++ b/apps/aether-gateway/src/tests/ai_execute/finalize_local_cli/compact.rs @@ -6,6 +6,7 @@ use super::{ EXECUTION_PATH_HEADER, TRACE_ID_HEADER, }; use crate::data::GatewayDataState; +use aether_ai_formats::openai_responses_message_item_id; use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY}; use aether_data::repository::auth::{ InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot, @@ -413,7 +414,7 @@ async fn gateway_executes_openai_responses_compact_openai_family_upstream_stream "output_text": "Hello Compact", "output": [{ "type": "message", - "id": "resp_compact_openai_family_123_msg", + "id": openai_responses_message_item_id("resp_compact_openai_family_123", 0), "role": "assistant", "status": "completed", "content": [{ diff --git a/apps/aether-gateway/src/tests/ai_execute/finalize_local_cli/direct.rs b/apps/aether-gateway/src/tests/ai_execute/finalize_local_cli/direct.rs index 21db78caf..2f217d05f 100644 --- a/apps/aether-gateway/src/tests/ai_execute/finalize_local_cli/direct.rs +++ b/apps/aether-gateway/src/tests/ai_execute/finalize_local_cli/direct.rs @@ -6,6 +6,7 @@ use super::{ TRACE_ID_HEADER, }; use crate::data::GatewayDataState; +use aether_ai_formats::openai_responses_message_item_id; use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY}; use aether_data::repository::auth::{ InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot, @@ -428,7 +429,7 @@ async fn gateway_executes_openai_responses_sync_upstream_stream_via_local_finali "output_text": "Hello", "output": [{ "type": "message", - "id": "resp_stream_001_msg", + "id": openai_responses_message_item_id("resp_stream_001", 0), "role": "assistant", "status": "completed", "content": [{ diff --git a/crates/aether-admin/src/provider/pool.rs b/crates/aether-admin/src/provider/pool.rs index 97021e99a..5da9fa198 100644 --- a/crates/aether-admin/src/provider/pool.rs +++ b/crates/aether-admin/src/provider/pool.rs @@ -118,6 +118,30 @@ pub fn admin_pool_key_quota_hard_blocked( aether_provider_pool::provider_pool_key_quota_hard_blocked(key, provider_type) } +pub fn admin_pool_key_model_quota_exhausted( + key: &StoredProviderCatalogKey, + provider_type: &str, + provider_model_name: &str, +) -> Option { + aether_provider_pool::provider_pool_key_model_quota_exhausted( + key, + provider_type, + provider_model_name, + ) +} + +pub fn admin_pool_key_model_quota_hard_blocked( + key: &StoredProviderCatalogKey, + provider_type: &str, + provider_model_name: &str, +) -> bool { + aether_provider_pool::provider_pool_key_model_quota_hard_blocked( + key, + provider_type, + provider_model_name, + ) +} + fn admin_pool_has_proxy(key: &StoredProviderCatalogKey) -> bool { match key.proxy.as_ref() { Some(Value::Object(values)) => !values.is_empty(), diff --git a/crates/aether-admin/src/provider/quota.rs b/crates/aether-admin/src/provider/quota.rs index 00de34d9d..81c833d23 100644 --- a/crates/aether-admin/src/provider/quota.rs +++ b/crates/aether-admin/src/provider/quota.rs @@ -10,6 +10,7 @@ const OAUTH_REFRESH_FAILED_PREFIX: &str = "[REFRESH_FAILED] "; const OAUTH_EXPIRED_PREFIX: &str = "[OAUTH_EXPIRED] "; const OAUTH_REQUEST_FAILED_PREFIX: &str = "[REQUEST_FAILED] "; const CODEX_SPARK_LIMIT_NAME: &str = "GPT-5.3-Codex-Spark"; +const CODEX_ADDITIONAL_QUOTA_WINDOWS_KEY: &str = "additional_quota_windows"; pub fn provider_auto_remove_banned_keys(config: Option<&serde_json::Value>) -> bool { config @@ -1893,6 +1894,112 @@ fn codex_find_spark_rate_limit( .and_then(serde_json::Value::as_object) } +/// Preserve every additional Codex rate-limit bucket in a model-addressable +/// representation. The upstream can add independent limits over time; do +/// not discard them merely because they are not one of the legacy flat +/// primary/secondary slots. +fn codex_additional_quota_windows( + root: &serde_json::Map, +) -> Vec { + let mut windows = Vec::new(); + let Some(items) = root + .get("additional_rate_limits") + .and_then(serde_json::Value::as_array) + else { + return windows; + }; + + for (index, item) in items.iter().enumerate() { + let Some(item_object) = item.as_object() else { + continue; + }; + let Some(limit_name) = ["limit_name", "metered_feature", "name", "id"] + .iter() + .find_map(|key| item_object.get(*key).and_then(serde_json::Value::as_str)) + .map(str::trim) + .filter(|value| !value.is_empty()) + else { + continue; + }; + let model_identity = item_object + .get("model") + .or_else(|| item_object.get("model_name")) + .or_else(|| item_object.get("model_id")) + .cloned() + .unwrap_or_else(|| json!(limit_name)); + let model_identities = item_object + .get("models") + .or_else(|| item_object.get("model_ids")) + .cloned(); + let Some(rate_limit) = item_object + .get("rate_limit") + .and_then(serde_json::Value::as_object) + else { + continue; + }; + + for (slot, slot_name) in [ + ("primary_window", "primary"), + ("secondary_window", "secondary"), + ] { + let Some(source) = rate_limit.get(slot).and_then(serde_json::Value::as_object) else { + continue; + }; + let used_percent = source.get("used_percent").and_then(coerce_json_f64); + let reset_after_seconds = source + .get("reset_after_seconds") + .or_else(|| source.get("reset_seconds")) + .and_then(coerce_json_u64); + let reset_at = source.get("reset_at").and_then(coerce_json_u64); + let window_minutes = source + .get("limit_window_seconds") + .or_else(|| source.get("window_seconds")) + .and_then(coerce_json_u64) + .map(|seconds| seconds.saturating_add(59) / 60) + .filter(|minutes| *minutes > 0); + if used_percent.is_none() + && reset_after_seconds.is_none() + && reset_at.is_none() + && window_minutes.is_none() + { + continue; + } + + let used_ratio = used_percent.map(|value| (value / 100.0).clamp(0.0, 1.0)); + let mut window = serde_json::Map::new(); + window.insert( + "code".to_string(), + json!(format!("additional_{index}_{slot_name}")), + ); + window.insert("label".to_string(), json!(limit_name)); + window.insert("scope".to_string(), json!("model")); + // The upstream limit name is the only stable identity available + // for an additional bucket. Keeping it in both fields allows + // exact model matching and generic family matching downstream. + window.insert("model".to_string(), model_identity.clone()); + if let Some(model_identities) = model_identities.clone() { + window.insert("models".to_string(), model_identities); + } + window.insert("quota_group".to_string(), json!(limit_name)); + window.insert("limit_name".to_string(), json!(limit_name)); + window.insert("used_ratio".to_string(), json!(used_ratio)); + window.insert( + "remaining_ratio".to_string(), + json!(used_ratio.map(|value| (1.0 - value).max(0.0))), + ); + window.insert( + "is_exhausted".to_string(), + json!(used_ratio.is_some_and(|value| value >= 1.0 - 1e-6)), + ); + window.insert("reset_at".to_string(), json!(reset_at)); + window.insert("reset_seconds".to_string(), json!(reset_after_seconds)); + window.insert("window_minutes".to_string(), json!(window_minutes)); + windows.push(serde_json::Value::Object(window)); + } + } + windows +} + fn codex_reset_credits_container( root: &serde_json::Map, ) -> Option<&serde_json::Map> { @@ -1993,6 +2100,14 @@ pub fn parse_codex_wham_usage_response( } } + // Keep all additional limits, including future ones with names unknown to + // this version of the gateway, so scheduling can select their bucket by + // metadata rather than by a hard-coded model name. + result.insert( + CODEX_ADDITIONAL_QUOTA_WINDOWS_KEY.to_string(), + serde_json::Value::Array(codex_additional_quota_windows(root)), + ); + if let Some(credits) = root.get("credits").and_then(serde_json::Value::as_object) { if let Some(value) = credits.get("has_credits").and_then(coerce_json_bool) { result.insert("has_credits".to_string(), json!(value)); @@ -5413,6 +5528,44 @@ mod tests { parsed.get("spark_secondary_window_minutes"), Some(&json!(10_080u64)) ); + let additional = parsed + .get("additional_quota_windows") + .and_then(serde_json::Value::as_array) + .expect("all additional quota windows should be preserved"); + assert_eq!(additional.len(), 2); + assert!(additional.iter().all(|window| { + window.get("model") == Some(&json!("GPT-5.3-Codex-Spark")) + && window.get("scope") == Some(&json!("model")) + })); + } + + #[test] + fn parses_unknown_additional_quota_buckets_without_model_specific_code() { + let parsed = parse_codex_wham_usage_response( + &json!({ + "rate_limit": { + "primary_window": {"used_percent": 100.0}, + "secondary_window": {"used_percent": 0.0} + }, + "additional_rate_limits": [{ + "limit_name": "Future-Model-Limit", + "rate_limit": { + "primary_window": { + "used_percent": 100.0, + "limit_window_seconds": 3600 + } + } + }] + }), + 1_700_000_000, + ) + .expect("quota response should parse"); + let additional = parsed["additional_quota_windows"] + .as_array() + .expect("additional windows should be an array"); + assert_eq!(additional.len(), 1); + assert_eq!(additional[0]["model"], json!("Future-Model-Limit")); + assert_eq!(additional[0]["is_exhausted"], json!(true)); } #[test] diff --git a/crates/aether-ai/formats/src/api.rs b/crates/aether-ai/formats/src/api.rs index e75172785..de48589f7 100644 --- a/crates/aether-ai/formats/src/api.rs +++ b/crates/aether-ai/formats/src/api.rs @@ -78,6 +78,7 @@ pub use crate::formats::openai::{ OpenAiProviderRequestFinalization, }, responses::{ + normalize_openai_responses_message_item_ids, openai_responses_message_item_id, openai_responses_synthetic_reasoning_item_id, strip_incompatible_openai_responses_reasoning_items, strip_incompatible_openai_responses_reasoning_items_with_policy, diff --git a/crates/aether-ai/formats/src/formats/conversion/response.rs b/crates/aether-ai/formats/src/formats/conversion/response.rs index 126fcf25f..646b99a9f 100644 --- a/crates/aether-ai/formats/src/formats/conversion/response.rs +++ b/crates/aether-ai/formats/src/formats/conversion/response.rs @@ -9,7 +9,7 @@ use serde_json::{json, Value}; use crate::formats::{ context::FormatContext, openai::responses::{ - openai_responses_synthetic_reasoning_item_id, + openai_responses_message_item_id, openai_responses_synthetic_reasoning_item_id, response::ensure_modern_openai_responses_response_fields, }, registry, @@ -218,7 +218,7 @@ pub fn build_openai_responses_response_with_content( if !content.is_empty() { output.push(json!({ "type": "message", - "id": format!("{response_id}_msg"), + "id": openai_responses_message_item_id(response_id, 0), "role": "assistant", "status": "completed", "content": content diff --git a/crates/aether-ai/formats/src/formats/openai/chat/stream.rs b/crates/aether-ai/formats/src/formats/openai/chat/stream.rs index 946efdf32..b90eb8b04 100644 --- a/crates/aether-ai/formats/src/formats/openai/chat/stream.rs +++ b/crates/aether-ai/formats/src/formats/openai/chat/stream.rs @@ -3,7 +3,7 @@ use std::collections::{BTreeMap, BTreeSet}; use serde_json::{json, Map, Value}; use crate::formats::openai::responses::{ - openai_responses_synthetic_reasoning_item_id, + openai_responses_message_item_id, openai_responses_synthetic_reasoning_item_id, response::{ ensure_modern_openai_responses_response_fields, openai_responses_current_timestamp, }, @@ -2240,7 +2240,7 @@ impl OpenAIResponsesClientEmitter { fn message_item_id(&self) -> String { self.message_item_id .clone() - .unwrap_or_else(|| format!("{}_msg", self.response_id())) + .unwrap_or_else(|| openai_responses_message_item_id(self.response_id(), 0)) } fn reasoning_item_id(&self) -> String { @@ -2259,7 +2259,7 @@ impl OpenAIResponsesClientEmitter { fn ensure_message_item_id(&mut self) -> String { if self.message_item_id.is_none() { - self.message_item_id = Some(format!("{}_msg", self.response_id())); + self.message_item_id = Some(openai_responses_message_item_id(self.response_id(), 0)); } self.message_item_id() } @@ -4167,7 +4167,7 @@ mod tests { assert!(sse.contains("event: response.output_item.done\n")); assert!(sse.contains("event: response.completed\n")); assert!(sse.contains("\"response_id\":\"resp_stream_123\"")); - assert!(sse.contains("\"item_id\":\"resp_stream_123_msg\"")); + assert!(sse.contains("\"item_id\":\"msg_aether_")); assert!(sse.contains("\"text\":\"Hello\"")); assert!(sse.contains("\"output_text\":\"Hello\"")); assert!(sse.contains("\"created_at\":")); @@ -4231,8 +4231,8 @@ mod tests { ); let sse = String::from_utf8(bytes).expect("sse should be utf8"); - assert!(sse.contains("\"item_id\":\"msg_first_msg\"")); - assert!(!sse.contains("\"item_id\":\"msg_second_msg\"")); + assert!(sse.contains("\"item_id\":\"msg_aether_")); + assert!(!sse.contains("msg_second_msg")); } #[test] diff --git a/crates/aether-ai/formats/src/formats/openai/request_contract.rs b/crates/aether-ai/formats/src/formats/openai/request_contract.rs index ebac54bd6..7eb7107bd 100644 --- a/crates/aether-ai/formats/src/formats/openai/request_contract.rs +++ b/crates/aether-ai/formats/src/formats/openai/request_contract.rs @@ -147,6 +147,13 @@ fn finalize_openai_provider_request_with_codex_model_capabilities_and_reasoning_ finalization.provider_api_format, reasoning_replay_policy, ); + if finalization + .provider_api_format + .trim() + .eq_ignore_ascii_case("openai:responses") + { + super::responses::normalize_openai_responses_message_item_ids(body); + } crate::enforce_request_body_stream_field( body, finalization.provider_api_format, @@ -408,6 +415,39 @@ mod tests { assert_eq!(input[1]["type"], "message"); } + #[test] + fn finalization_repairs_legacy_responses_message_ids_for_same_format_upstream() { + let mut body = json!({ + "model": "gpt-5.4", + "input": [{ + "type": "message", + "id": "1c938e58-32a8-4d28-9c34-538d78076895_msg", + "role": "assistant", + "content": [{"type": "input_text", "text": "previous answer"}] + }] + }); + + finalize_openai_provider_request( + &mut body, + OpenAiProviderRequestFinalization { + source_api_format: "openai:responses", + provider_api_format: "openai:responses", + provider_type: "codex", + provider_model: "gpt-5.4", + source_model: "gpt-5.4", + body_rules: None, + upstream_is_stream: false, + require_body_stream_field: false, + }, + ) + .expect("legacy message IDs should be repaired before provider validation"); + + let repaired_id = body["input"][0]["id"] + .as_str() + .expect("message ID should be a string"); + assert!(repaired_id.starts_with("msg_")); + } + #[test] fn non_responses_sources_receive_codex_responses_reasoning_defaults() { for source_api_format in ["openai:chat", "claude:messages", "gemini:generate_content"] { diff --git a/crates/aether-ai/formats/src/formats/openai/responses/mod.rs b/crates/aether-ai/formats/src/formats/openai/responses/mod.rs index 5be47f610..8db6327b8 100644 --- a/crates/aether-ai/formats/src/formats/openai/responses/mod.rs +++ b/crates/aether-ai/formats/src/formats/openai/responses/mod.rs @@ -9,6 +9,7 @@ pub mod stream; const TOOL_ERROR_PREFIX: &str = "[tool error]"; const AETHER_REASONING_ITEM_ID_PREFIX: &str = "rs_aether_"; +const AETHER_MESSAGE_ITEM_ID_PREFIX: &str = "msg_aether_"; /// Controls which provider-owned reasoning items may be replayed on a Responses request. /// @@ -38,6 +39,67 @@ pub fn openai_responses_synthetic_reasoning_item_id( ) } +/// Builds a stable, wire-compatible ID for a message item synthesized by Aether. +/// +/// Responses clients replay assistant message items verbatim on the next turn and +/// OpenAI requires those IDs to begin with `msg`. Upstream response IDs are not +/// guaranteed to have that prefix (Chat completion IDs and UUIDs are common), so +/// appending a suffix to the response ID is not sufficient. A deterministic UUID +/// keeps the ID stable across sync/stream projections while avoiding assumptions +/// about the upstream ID's shape or length. +pub fn openai_responses_message_item_id(response_id: &str, output_index: usize) -> String { + let seed = format!("{response_id}:{output_index}"); + format!( + "{AETHER_MESSAGE_ITEM_ID_PREFIX}{}", + uuid::Uuid::new_v5(&uuid::Uuid::NAMESPACE_OID, seed.as_bytes()).simple() + ) +} + +/// Repairs legacy/non-OpenAI message IDs in a Responses request in place. +/// +/// Aether versions before the `msg_` contract emitted IDs such as +/// `_msg`. Clients legitimately replay those assistant items on +/// the next turn, so merely fixing newly generated responses leaves existing +/// conversations broken. Preserve already-valid provider IDs and deterministically +/// remap only message items that do not begin with `msg`. +pub fn normalize_openai_responses_message_item_ids(body: &mut Value) -> usize { + let Some(items) = body.get_mut("input").and_then(Value::as_array_mut) else { + return 0; + }; + let mut repaired = 0usize; + for (index, item) in items.iter_mut().enumerate() { + let Some(object) = item.as_object_mut() else { + continue; + }; + if object.get("type").and_then(Value::as_str) != Some("message") { + continue; + } + let Some(raw_id) = object.get("id") else { + // IDs are optional for newly-authored input messages. Only repair + // an ID that a previous response actually supplied. + continue; + }; + let valid = object + .get("id") + .and_then(Value::as_str) + .is_some_and(|id| id.starts_with("msg")); + if valid { + continue; + } + let source_id = raw_id + .as_str() + .filter(|id| !id.trim().is_empty()) + .unwrap_or("missing") + .to_string(); + object.insert( + "id".to_string(), + Value::String(openai_responses_message_item_id(source_id.as_str(), index)), + ); + repaired += 1; + } + repaired +} + /// Removes reasoning history items that cannot be replayed against an OpenAI Responses backend. /// /// Reasoning IDs are opaque provider references and must never be repaired by changing their @@ -190,7 +252,9 @@ mod tests { use serde_json::json; use super::{ - openai_responses_request_operation, openai_responses_synthetic_reasoning_item_id, + normalize_openai_responses_message_item_ids, openai_responses_message_item_id, + openai_responses_request_operation, + openai_responses_synthetic_reasoning_item_id, strip_incompatible_openai_responses_reasoning_items, strip_incompatible_openai_responses_reasoning_items_with_policy, OpenAiResponsesReasoningReplayPolicy, OPENAI_RESPONSES_OPERATION_COMPACT, @@ -238,6 +302,42 @@ mod tests { assert_ne!(first, other); } + #[test] + fn synthetic_message_item_ids_are_stable_and_start_with_msg() { + let first = openai_responses_message_item_id( + "1c938e58-32a8-4d28-9c34-538d78076895", + 0, + ); + let second = openai_responses_message_item_id( + "1c938e58-32a8-4d28-9c34-538d78076895", + 0, + ); + let other = openai_responses_message_item_id("chatcmpl-123", 1); + + assert!(first.starts_with("msg_")); + assert_eq!(first, second); + assert_ne!(first, other); + } + + #[test] + fn normalizes_legacy_message_ids_but_preserves_valid_ids() { + let mut body = json!({ + "input": [ + {"type": "message", "id": "1c938e58-32a8-4d28-9c34-538d78076895_msg", "role": "assistant"}, + {"type": "message", "id": "msg_provider_123", "role": "assistant"}, + {"type": "function_call", "id": "legacy_call"}, + {"type": "message", "role": "user"} + ] + }); + + assert_eq!(normalize_openai_responses_message_item_ids(&mut body), 1); + let input = body["input"].as_array().expect("input array"); + assert!(input[0]["id"].as_str().is_some_and(|id| id.starts_with("msg_"))); + assert_eq!(input[1]["id"], "msg_provider_123"); + assert_eq!(input[2].get("id"), Some(&json!("legacy_call"))); + assert!(input[3].get("id").is_none()); + } + #[test] fn strips_foreign_and_non_replayable_synthetic_reasoning_items() { let portable_synthetic = openai_responses_synthetic_reasoning_item_id("resp_123", 1); diff --git a/crates/aether-ai/formats/src/formats/shared/sync_products.rs b/crates/aether-ai/formats/src/formats/shared/sync_products.rs index 57d512eb6..e4ed1bd60 100644 --- a/crates/aether-ai/formats/src/formats/shared/sync_products.rs +++ b/crates/aether-ai/formats/src/formats/shared/sync_products.rs @@ -7,7 +7,9 @@ use aether_ai_formats::formats::conversion::response::{ convert_openai_chat_response_to_openai_responses, convert_openai_responses_response_to_openai_chat, }; -use aether_ai_formats::formats::openai::responses::openai_responses_synthetic_reasoning_item_id; +use aether_ai_formats::formats::openai::responses::{ + openai_responses_message_item_id, openai_responses_synthetic_reasoning_item_id, +}; use aether_ai_formats::formats::openai::responses::response::ensure_modern_openai_responses_response_fields; use aether_ai_formats::formats::registry::{convert_response, FormatContext, FormatError}; use aether_ai_formats::{ @@ -2696,6 +2698,7 @@ fn aggregate_openai_responses_stream_sync_response_from_validated_terminal( if let Some(state) = message_states.remove(&output_index) { output.push(materialize_openai_responses_message_item( &response_id, + output_index, state, )); } @@ -3152,13 +3155,31 @@ fn resolve_openai_responses_tool_output_index( fn materialize_openai_responses_message_item( response_id: &str, + output_index: usize, state: OpenAIResponsesSyncMessageState, ) -> Value { let mut item = state.item; item.entry("type".to_string()) .or_insert_with(|| Value::String("message".to_string())); - item.entry("id".to_string()) - .or_insert_with(|| Value::String(format!("{response_id}_msg"))); + let message_id_is_valid = item + .get("id") + .and_then(Value::as_str) + .is_some_and(|id| id.starts_with("msg")); + if !message_id_is_valid { + let source_id = item + .get("id") + .and_then(Value::as_str) + .filter(|id| !id.trim().is_empty()) + .unwrap_or(response_id) + .to_string(); + item.insert( + "id".to_string(), + Value::String(openai_responses_message_item_id( + source_id.as_str(), + output_index, + )), + ); + } item.entry("role".to_string()) .or_insert_with(|| Value::String("assistant".to_string())); item.entry("status".to_string()) diff --git a/crates/aether-ai/formats/src/lib.rs b/crates/aether-ai/formats/src/lib.rs index fbb8a3f1e..bb3d736f2 100644 --- a/crates/aether-ai/formats/src/lib.rs +++ b/crates/aether-ai/formats/src/lib.rs @@ -56,7 +56,9 @@ pub use formats::openai::responses::request::{ validate_openai_responses_request_contract, OpenAiResponsesRequestContractViolation, }; pub use formats::openai::responses::{ - openai_responses_request_operation, openai_responses_synthetic_reasoning_item_id, + normalize_openai_responses_message_item_ids, openai_responses_message_item_id, + openai_responses_request_operation, + openai_responses_synthetic_reasoning_item_id, strip_incompatible_openai_responses_reasoning_items, strip_incompatible_openai_responses_reasoning_items_with_policy, OpenAiResponsesReasoningReplayPolicy, OPENAI_RESPONSES_OPERATION_COMPACT, diff --git a/crates/aether-ai/formats/src/protocol/canonical.rs b/crates/aether-ai/formats/src/protocol/canonical.rs index d6318775d..b042f7db4 100644 --- a/crates/aether-ai/formats/src/protocol/canonical.rs +++ b/crates/aether-ai/formats/src/protocol/canonical.rs @@ -4,6 +4,7 @@ use serde::{Deserialize, Serialize}; use serde_json::{json, Map, Value}; use crate::formats::openai::shared::map_thinking_budget_to_openai_reasoning_effort; +use crate::formats::openai::responses::openai_responses_message_item_id; use crate::formats::shared::model_directives::ReasoningEffort; use crate::formats::shared::response::remove_empty_pages_from_tool_input_value; @@ -3948,11 +3949,7 @@ pub(crate) fn flush_openai_responses_message_item( if message_content.is_empty() { return; } - let id = if *message_index == 0 { - format!("{response_id}_msg") - } else { - format!("{response_id}_msg_{message_index}") - }; + let id = openai_responses_message_item_id(response_id, *message_index); output.push(json!({ "type": "message", "id": id, diff --git a/crates/aether-data/runtime/src/lifecycle/backfill/postgres.rs b/crates/aether-data/runtime/src/lifecycle/backfill/postgres.rs index 80552bc67..e93521979 100644 --- a/crates/aether-data/runtime/src/lifecycle/backfill/postgres.rs +++ b/crates/aether-data/runtime/src/lifecycle/backfill/postgres.rs @@ -8,6 +8,8 @@ use tracing::{error, info, warn}; use super::types::PendingBackfillInfo; +// Historical backfill manifests must remain available after they are applied; +// deployed databases retain their versions for startup compatibility checks. static BACKFILL_MIGRATOR: Migrator = sqlx::migrate!("./backfills/postgres"); const SCHEMA_BACKFILLS_TABLE_EXISTS_SQL: &str = diff --git a/crates/aether-gateway/control/src/public/request_context.rs b/crates/aether-gateway/control/src/public/request_context.rs index 8c38439e0..d0126cf99 100644 --- a/crates/aether-gateway/control/src/public/request_context.rs +++ b/crates/aether-gateway/control/src/public/request_context.rs @@ -8,6 +8,7 @@ pub struct PublicRequestContext { pub request_query_string: Option, pub request_content_type: Option, pub host_header: Option, + pub client_ip: Option, pub control_decision: Option, } @@ -32,6 +33,7 @@ impl PublicRequestContext { request_query_string: uri.query().map(ToOwned::to_owned), request_content_type: header_value(headers, http::header::CONTENT_TYPE), host_header: header_value(headers, http::header::HOST), + client_ip: None, control_decision, } } diff --git a/crates/aether-pool-core/src/scheduler.rs b/crates/aether-pool-core/src/scheduler.rs index 64cf939b7..8cf3ebfc3 100644 --- a/crates/aether-pool-core/src/scheduler.rs +++ b/crates/aether-pool-core/src/scheduler.rs @@ -18,6 +18,8 @@ pub struct PoolSchedulingPreset { pub struct PoolSchedulingConfig { pub scheduling_presets: Vec, pub lru_enabled: bool, + /// Retained for configuration/API compatibility. Active quota exhaustion is + /// always an admission block; reset-aware adapters decide when it clears. pub skip_exhausted_accounts: bool, pub cost_limit_per_key_tokens: Option, } @@ -225,9 +227,14 @@ fn schedule_pool_group( continue; } - if item.key_context.quota_hard_blocked - || (pool_config.skip_exhausted_accounts && item.key_context.quota_exhausted) - { + // A quota snapshot is an account-level admission signal, not merely a + // ranking hint. Continuing to schedule a member whose quota is known to + // be exhausted causes a request-wide retry storm (the upstream returns + // 429 for every attempt). `quota_hard_blocked` remains available for + // providers that can distinguish an explicit permanent block, but every + // active exhaustion must be removed from the request's candidate set; + // reset-aware provider adapters clear the signal once capacity returns. + if item.key_context.quota_hard_blocked || item.key_context.quota_exhausted { skipped.push(PoolSkippedCandidate { candidate: item.candidate, skip_reason: POOL_ACCOUNT_EXHAUSTED_SKIP_REASON, @@ -928,6 +935,29 @@ mod tests { ); } + #[test] + fn pool_scheduler_skips_exhausted_accounts_even_when_legacy_flag_is_false() { + let ready = sample_candidate("provider-pool", "endpoint-1", "key-ready", 10, true); + let mut exhausted = + sample_candidate("provider-pool", "endpoint-1", "key-exhausted", 10, true); + exhausted.key_context.quota_exhausted = true; + + let outcome = run_pool_scheduler(vec![ready, exhausted], &BTreeMap::new(), "seed"); + + assert_eq!( + outcome + .candidates + .iter() + .map(|item| item.candidate.as_str()) + .collect::>(), + vec!["key-ready"] + ); + assert_eq!( + outcome.skipped_candidates[0].skip_reason, + POOL_ACCOUNT_EXHAUSTED_SKIP_REASON + ); + } + #[test] fn pool_scheduler_promotes_sticky_hit_before_other_sorted_keys() { let key_a = sample_candidate("provider-pool", "endpoint-1", "key-a", 10, true) diff --git a/crates/aether-provider/pool/src/lib.rs b/crates/aether-provider/pool/src/lib.rs index cb147c8f6..a8e83e85a 100644 --- a/crates/aether-provider/pool/src/lib.rs +++ b/crates/aether-provider/pool/src/lib.rs @@ -35,6 +35,7 @@ pub use providers::{ }; pub use quota::{ provider_pool_key_account_quota_exhausted, provider_pool_key_quota_hard_blocked, + provider_pool_key_model_quota_exhausted, provider_pool_key_model_quota_hard_blocked, provider_pool_key_scheduling_label, provider_pool_member_quota_snapshot, provider_pool_quota_metadata_provider_type, provider_pool_quota_metadata_updated_at, provider_pool_quota_snapshot_updated_at, @@ -561,7 +562,7 @@ mod tests { } }))); - let signals = service.member_signals("windsurf", &key, None); + let signals = service.member_signals("windsurf", &key, None, None); assert!(!signals.quota_exhausted); } @@ -588,7 +589,7 @@ mod tests { } })); - let signals = service.member_signals("windsurf", &key, None); + let signals = service.member_signals("windsurf", &key, None, None); assert!(signals.quota_exhausted); } @@ -876,6 +877,285 @@ mod tests { )); } + #[test] + fn model_quota_windows_are_isolated_without_provider_specific_names() { + let service = ProviderPoolService::with_builtin_adapters(); + let mut key = sample_key(None); + key.status_snapshot = Some(json!({ + "quota": { + "version": 2, + "provider_type": "codex", + "exhausted": true, + "windows": [ + { + "code": "alpha_short", + "quota_group": "alpha", + "model": "vendor-alpha-model", + "used_ratio": 1.0, + "is_exhausted": true + }, + { + "code": "alpha_long", + "quota_group": "alpha", + "model": "vendor-alpha-model", + "used_ratio": 0.2, + "is_exhausted": false + }, + { + "code": "beta_short", + "quota_group": "beta", + "model": "vendor-beta-model", + "used_ratio": 1.0, + "is_exhausted": true + } + ] + } + })); + + let alpha = service.member_signals( + "codex", + &key, + None, + Some("vendor-alpha-model"), + ); + let beta = service.member_signals("codex", &key, None, Some("vendor-beta-model")); + assert!(!alpha.quota_exhausted, "one available alpha window must keep it usable"); + assert!(beta.quota_exhausted); + assert!(!beta.quota_hard_blocked); + } + + #[test] + fn legacy_family_prefix_is_matched_by_model_tokens() { + let service = ProviderPoolService::with_builtin_adapters(); + let mut key = sample_key(None); + key.status_snapshot = Some(json!({ + "quota": { + "version": 2, + "provider_type": "codex", + "exhausted": true, + "windows": [ + { "code": "alpha_weekly", "used_ratio": 1.0, "is_exhausted": true }, + { "code": "default_weekly", "used_ratio": 0.1, "is_exhausted": false } + ] + } + })); + + let alpha = service.member_signals("codex", &key, None, Some("vendor-alpha-v2")); + assert!(alpha.quota_exhausted); + // An unrelated model has no identifiable bucket and therefore falls + // back to the account-level snapshot rather than guessing a family. + let unknown = service.member_signals("codex", &key, None, Some("vendor-gamma-v2")); + assert!(unknown.quota_exhausted); + } + + #[test] + fn compact_model_bucket_identity_matches_request_tokens() { + let service = ProviderPoolService::with_builtin_adapters(); + let mut key = sample_key(None); + key.status_snapshot = Some(json!({ + "quota": { + "version": 2, + "provider_type": "codex", + "exhausted": true, + "windows": [ + { + "code": "additional_0_primary_window", + "scope": "model", + "model": "spark", + "used_ratio": 0.0, + "is_exhausted": false + } + ] + } + })); + + let signals = service.member_signals( + "codex", + &key, + None, + Some("gpt-5.3-codex-spark"), + ); + assert!( + !signals.quota_exhausted, + "a compact bucket name should match a token in the selected model" + ); + } + + #[test] + fn model_family_token_matches_versioned_alias_without_hardcoding_name() { + let service = ProviderPoolService::with_builtin_adapters(); + let mut key = sample_key(None); + key.status_snapshot = Some(json!({ + "quota": { + "version": 2, + "provider_type": "codex", + "exhausted": true, + "windows": [{ + "code": "additional_0_primary_window", + "scope": "model", + "model": "gpt-5.3-codex-spark", + "used_ratio": 1.0, + "is_exhausted": true + }] + } + })); + + let signals = service.member_signals( + "codex", + &key, + None, + Some("gpt-5.4-codex-spark"), + ); + assert!(signals.quota_exhausted); + } + + #[test] + fn model_only_snapshot_does_not_poison_account_fallback() { + let service = ProviderPoolService::with_builtin_adapters(); + let mut key = sample_key(None); + key.status_snapshot = Some(json!({ + "quota": { + "version": 2, + "provider_type": "codex", + "exhausted": true, + "windows": [{ + "code": "model:vendor-alpha", + "scope": "model", + "model": "vendor-alpha", + "used_ratio": 1.0, + "is_exhausted": true + }] + } + })); + + let signals = service.member_signals("codex", &key, None, None); + assert!( + !signals.quota_exhausted, + "account-level inspection must ignore model-only buckets" + ); + } + + #[test] + fn model_map_quota_is_resolved_without_materialized_windows() { + let service = ProviderPoolService::with_builtin_adapters(); + let mut key = sample_key(None); + key.status_snapshot = Some(json!({ + "quota": { + "version": 2, + "provider_type": "grok", + "exhausted": false, + "quota_by_model": { + "vendor-alpha": { + "remaining": 0.0, + "total": 10.0 + }, + "vendor-beta": { + "remaining": 5.0, + "total": 10.0 + } + } + } + })); + + let alpha = service.member_signals("grok", &key, None, Some("vendor-alpha")); + let beta = service.member_signals("grok", &key, None, Some("vendor-beta")); + assert!(alpha.quota_exhausted); + assert!(!beta.quota_exhausted); + } + + #[test] + fn raw_provider_metadata_model_bucket_overrides_account_snapshot() { + let service = ProviderPoolService::with_builtin_adapters(); + let mut key = sample_key(Some(json!({ + "codex": { + "updated_at": 1_700_000_001u64, + "additional_quota_windows": [{ + "scope": "model", + "model": "future-spark", + "used_ratio": 0.0, + "is_exhausted": false + }] + } + }))); + key.status_snapshot = Some(json!({ + "quota": { + "version": 2, + "provider_type": "codex", + "exhausted": true, + "windows": [{ + "code": "primary", + "scope": "account", + "used_ratio": 1.0, + "is_exhausted": true + }] + } + })); + + let signals = service.member_signals("codex", &key, None, Some("future-spark")); + assert!( + !signals.quota_exhausted, + "a newer model bucket in provider metadata must not inherit an exhausted account bucket" + ); + } + + #[test] + fn model_quota_availability_suppresses_account_hard_block() { + let service = ProviderPoolService::with_builtin_adapters(); + let mut key = sample_key(None); + key.status_snapshot = Some(json!({ + "quota": { + "version": 2, + "provider_type": "codex", + "allowed": false, + "limit_reached": true, + "exhausted": true, + "windows": [{ + "scope": "model", + "model": "future-model", + "used_ratio": 0.0, + "is_exhausted": false + }] + } + })); + + let signals = service.member_signals("codex", &key, None, Some("future-model")); + assert!(!signals.quota_exhausted); + assert!(!signals.quota_hard_blocked); + } + + #[test] + fn newer_model_quota_observation_wins_over_stale_snapshot() { + let service = ProviderPoolService::with_builtin_adapters(); + let mut key = sample_key(Some(json!({ + "codex": { + "updated_at": 1_700_000_200u64, + "additional_quota_windows": [{ + "scope": "model", + "model": "future-model", + "used_ratio": 0.0, + "is_exhausted": false + }] + } + }))); + key.status_snapshot = Some(json!({ + "quota": { + "version": 2, + "provider_type": "codex", + "observed_at": 1_700_000_100u64, + "exhausted": true, + "windows": [{ + "scope": "model", + "model": "future-model", + "used_ratio": 1.0, + "is_exhausted": true + }] + } + })); + + let signals = service.member_signals("codex", &key, None, Some("future-model")); + assert!(!signals.quota_exhausted); + } + #[test] fn codex_explicit_quota_block_is_hard_until_reset() { let now = std::time::SystemTime::now() diff --git a/crates/aether-provider/pool/src/provider.rs b/crates/aether-provider/pool/src/provider.rs index c25ab672f..87f54e7f5 100644 --- a/crates/aether-provider/pool/src/provider.rs +++ b/crates/aether-provider/pool/src/provider.rs @@ -8,7 +8,8 @@ use crate::capability::{ProviderPoolCapabilities, ProviderPoolCapability}; use crate::plan::{derive_plan_tier, normalize_provider_plan_tier}; use crate::quota::{ provider_pool_account_blocked, provider_pool_quota_reset_seconds, - provider_pool_quota_snapshot_exhausted_decision, provider_pool_quota_usage_ratio, + provider_pool_model_quota_exhausted, provider_pool_quota_snapshot_exhausted_decision, + provider_pool_quota_usage_ratio, }; #[derive(Debug, Clone)] @@ -16,6 +17,11 @@ pub struct ProviderPoolMemberInput<'a> { pub provider_type: &'a str, pub key: &'a StoredProviderCatalogKey, pub auth_config: Option<&'a Map>, + /// The provider-side model selected for the current request. Quota + /// snapshots may contain several independent windows, so adapters use + /// this value to select the applicable bucket instead of treating the + /// whole account as one quota. + pub provider_model_name: Option<&'a str>, } pub trait ProviderPoolAdapter: Send + Sync { @@ -70,6 +76,11 @@ pub trait ProviderPoolAdapter: Send + Sync { } fn quota_exhausted(&self, input: &ProviderPoolMemberInput<'_>) -> bool { + if let Some(exhausted) = input.provider_model_name.and_then(|model| { + provider_pool_model_quota_exhausted(input.key, input.provider_type, model) + }) { + return exhausted; + } provider_pool_quota_snapshot_exhausted_decision(input.key, input.provider_type) .unwrap_or(false) } diff --git a/crates/aether-provider/pool/src/providers/antigravity.rs b/crates/aether-provider/pool/src/providers/antigravity.rs index 77e474504..24d9bb797 100644 --- a/crates/aether-provider/pool/src/providers/antigravity.rs +++ b/crates/aether-provider/pool/src/providers/antigravity.rs @@ -6,7 +6,9 @@ use serde_json::json; use crate::capability::ProviderPoolCapabilities; use crate::provider::{ provider_pool_endpoint_format_matches, provider_pool_matching_endpoint, ProviderPoolAdapter, + ProviderPoolMemberInput, }; +use crate::quota::provider_pool_model_quota_exhausted; use crate::quota_refresh::ProviderPoolQuotaRequestSpec; pub const ANTIGRAVITY_FETCH_AVAILABLE_MODELS_PATH: &str = "/v1internal:fetchAvailableModels"; @@ -26,6 +28,21 @@ impl ProviderPoolAdapter for AntigravityProviderPoolAdapter { } } + fn quota_exhausted(&self, input: &ProviderPoolMemberInput<'_>) -> bool { + input + .provider_model_name + .and_then(|model| { + provider_pool_model_quota_exhausted(input.key, input.provider_type, model) + }) + .unwrap_or_else(|| { + crate::quota::provider_pool_quota_snapshot_exhausted_decision( + input.key, + input.provider_type, + ) + .unwrap_or(false) + }) + } + fn quota_refresh_endpoint( &self, endpoints: &[StoredProviderCatalogEndpoint], diff --git a/crates/aether-provider/pool/src/providers/chatgpt_web.rs b/crates/aether-provider/pool/src/providers/chatgpt_web.rs index 536d5337c..3eb4b5754 100644 --- a/crates/aether-provider/pool/src/providers/chatgpt_web.rs +++ b/crates/aether-provider/pool/src/providers/chatgpt_web.rs @@ -12,7 +12,8 @@ use crate::provider::{ use crate::quota::{ provider_pool_current_unix_secs, provider_pool_json_bool, provider_pool_json_f64, provider_pool_metadata_bucket, provider_pool_quota_snapshot_exhausted_decision, - provider_pool_reset_deadline_elapsed, provider_pool_timestamp_unix_secs, + provider_pool_model_quota_exhausted, provider_pool_reset_deadline_elapsed, + provider_pool_timestamp_unix_secs, }; use crate::quota_refresh::ProviderPoolQuotaRequestSpec; @@ -40,6 +41,11 @@ impl ProviderPoolAdapter for ChatGptWebProviderPoolAdapter { } fn quota_exhausted(&self, input: &ProviderPoolMemberInput<'_>) -> bool { + if let Some(exhausted) = input.provider_model_name.and_then(|model| { + provider_pool_model_quota_exhausted(input.key, input.provider_type, model) + }) { + return exhausted; + } if let Some(exhausted) = provider_pool_quota_snapshot_exhausted_decision(input.key, input.provider_type) { diff --git a/crates/aether-provider/pool/src/providers/codex.rs b/crates/aether-provider/pool/src/providers/codex.rs index 66f7413a0..0086e813e 100644 --- a/crates/aether-provider/pool/src/providers/codex.rs +++ b/crates/aether-provider/pool/src/providers/codex.rs @@ -12,8 +12,8 @@ use crate::provider::{ use crate::quota::{ provider_pool_current_unix_secs, provider_pool_json_bool, provider_pool_json_f64, provider_pool_member_quota_snapshot, provider_pool_metadata_bucket, - provider_pool_quota_snapshot_exhausted_decision, provider_pool_reset_deadline_elapsed, - provider_pool_timestamp_unix_secs, + provider_pool_model_quota_exhausted, provider_pool_quota_snapshot_exhausted_decision, + provider_pool_reset_deadline_elapsed, provider_pool_timestamp_unix_secs, }; use crate::quota_refresh::ProviderPoolQuotaRequestSpec; @@ -49,6 +49,11 @@ impl ProviderPoolAdapter for CodexProviderPoolAdapter { } fn quota_exhausted(&self, input: &ProviderPoolMemberInput<'_>) -> bool { + if let Some(exhausted) = input.provider_model_name.and_then(|model| { + provider_pool_model_quota_exhausted(input.key, input.provider_type, model) + }) { + return exhausted; + } if let Some(quota_snapshot) = provider_pool_member_quota_snapshot(input.key, input.provider_type) { @@ -76,6 +81,11 @@ impl ProviderPoolAdapter for CodexProviderPoolAdapter { } fn quota_hard_blocked(&self, input: &ProviderPoolMemberInput<'_>) -> bool { + if input.provider_model_name.is_some_and(|model| { + provider_pool_model_quota_exhausted(input.key, input.provider_type, model).is_some() + }) { + return false; + } codex_explicit_quota_block_active(input.key, input.provider_type) } diff --git a/crates/aether-provider/pool/src/providers/grok.rs b/crates/aether-provider/pool/src/providers/grok.rs index c2bb4dc2c..ce9e2f91b 100644 --- a/crates/aether-provider/pool/src/providers/grok.rs +++ b/crates/aether-provider/pool/src/providers/grok.rs @@ -9,7 +9,8 @@ use crate::provider::{ use crate::quota::{ provider_pool_current_unix_secs, provider_pool_json_bool, provider_pool_json_f64, provider_pool_metadata_bucket, provider_pool_quota_snapshot_exhausted_decision, - provider_pool_reset_deadline_elapsed, provider_pool_timestamp_unix_secs, + provider_pool_model_quota_exhausted, provider_pool_reset_deadline_elapsed, + provider_pool_timestamp_unix_secs, }; pub const GROK_QUOTA_WINDOWS_BASIC: &[(&str, &str)] = &[("quota_fast", "fast")]; @@ -44,6 +45,11 @@ impl ProviderPoolAdapter for GrokProviderPoolAdapter { } fn quota_exhausted(&self, input: &ProviderPoolMemberInput<'_>) -> bool { + if let Some(exhausted) = input.provider_model_name.and_then(|model| { + provider_pool_model_quota_exhausted(input.key, input.provider_type, model) + }) { + return exhausted; + } if let Some(exhausted) = provider_pool_quota_snapshot_exhausted_decision(input.key, input.provider_type) { diff --git a/crates/aether-provider/pool/src/providers/kiro.rs b/crates/aether-provider/pool/src/providers/kiro.rs index 72c175315..4daab5111 100644 --- a/crates/aether-provider/pool/src/providers/kiro.rs +++ b/crates/aether-provider/pool/src/providers/kiro.rs @@ -12,8 +12,8 @@ use crate::provider::{ }; use crate::quota::{ provider_pool_current_unix_secs, provider_pool_json_f64, provider_pool_metadata_bucket, - provider_pool_quota_snapshot_exhausted_decision, provider_pool_reset_deadline_elapsed, - provider_pool_timestamp_unix_secs, + provider_pool_model_quota_exhausted, provider_pool_quota_snapshot_exhausted_decision, + provider_pool_reset_deadline_elapsed, provider_pool_timestamp_unix_secs, }; use crate::quota_refresh::ProviderPoolQuotaRequestSpec; @@ -46,6 +46,11 @@ impl ProviderPoolAdapter for KiroProviderPoolAdapter { } fn quota_exhausted(&self, input: &ProviderPoolMemberInput<'_>) -> bool { + if let Some(exhausted) = input.provider_model_name.and_then(|model| { + provider_pool_model_quota_exhausted(input.key, input.provider_type, model) + }) { + return exhausted; + } if let Some(exhausted) = provider_pool_quota_snapshot_exhausted_decision(input.key, input.provider_type) { diff --git a/crates/aether-provider/pool/src/providers/windsurf.rs b/crates/aether-provider/pool/src/providers/windsurf.rs index 112cbbb68..806f9136e 100644 --- a/crates/aether-provider/pool/src/providers/windsurf.rs +++ b/crates/aether-provider/pool/src/providers/windsurf.rs @@ -13,7 +13,8 @@ use crate::provider::{ }; use crate::quota::{ provider_pool_json_bool, provider_pool_json_f64, provider_pool_member_quota_snapshot, - provider_pool_metadata_bucket, provider_pool_quota_snapshot_exhausted_decision, + provider_pool_metadata_bucket, provider_pool_model_quota_exhausted, + provider_pool_quota_snapshot_exhausted_decision, }; use crate::quota_refresh::ProviderPoolQuotaRequestSpec; @@ -50,6 +51,11 @@ impl ProviderPoolAdapter for WindsurfProviderPoolAdapter { } fn quota_exhausted(&self, input: &ProviderPoolMemberInput<'_>) -> bool { + if let Some(exhausted) = input.provider_model_name.and_then(|model| { + provider_pool_model_quota_exhausted(input.key, input.provider_type, model) + }) { + return exhausted; + } if windsurf_quota_snapshot_hard_exhausted(input.key, input.provider_type) { return true; } diff --git a/crates/aether-provider/pool/src/quota.rs b/crates/aether-provider/pool/src/quota.rs index 86ff4e3d6..a3a1e78a0 100644 --- a/crates/aether-provider/pool/src/quota.rs +++ b/crates/aether-provider/pool/src/quota.rs @@ -15,6 +15,7 @@ pub fn provider_pool_key_account_quota_exhausted( provider_type, key, auth_config: None, + provider_model_name: None, }) } @@ -27,6 +28,25 @@ pub fn provider_pool_key_quota_hard_blocked( provider_type, key, auth_config: None, + provider_model_name: None, + }) +} + +/// Model-aware hard-block lookup for pre-scheduler candidate filtering. Some +/// providers expose permanent account flags alongside independent model +/// buckets; adapters can suppress the account flag when the selected model +/// has its own usable quota. +pub fn provider_pool_key_model_quota_hard_blocked( + key: &StoredProviderCatalogKey, + provider_type: &str, + provider_model_name: &str, +) -> bool { + let adapter = ProviderPoolService::with_builtin_adapters().adapter(provider_type); + adapter.quota_hard_blocked(&ProviderPoolMemberInput { + provider_type, + key, + auth_config: None, + provider_model_name: Some(provider_model_name), }) } @@ -44,6 +64,331 @@ pub fn provider_pool_member_quota_snapshot<'a>( .then_some(quota_snapshot) } +/// Resolve exhaustion for the quota bucket applicable to one provider model. +/// +/// Providers are free to expose quota windows in different shapes. Newer +/// snapshots should put an explicit `model`/`models` (or `quota_group`) on a +/// window; legacy Codex snapshots use a family prefix in `code` (for example +/// `spark_5h`). We deliberately do not name any product or model here: the +/// resolver compares the metadata supplied by the provider with the selected +/// model and only falls back to account-level exhaustion when no model bucket +/// can be identified. +pub(crate) fn provider_pool_model_quota_exhausted( + key: &StoredProviderCatalogKey, + provider_type: &str, + provider_model_name: &str, +) -> Option { + let requested = provider_pool_identifier_tokens(provider_model_name); + if requested.is_empty() { + return None; + } + + // Prefer the materialized status snapshot, but also inspect the raw + // provider metadata. A quota refresh and a request can race, leaving the + // latter newer than the snapshot; resolving both avoids falling back to an + // account-wide signal and incorrectly blocking an unrelated model bucket. + let sources = [ + provider_pool_member_quota_snapshot(key, provider_type), + provider_pool_metadata_bucket(key.upstream_metadata.as_ref(), provider_type), + ]; + let mut resolved = None::<(Option, bool)>; + for source in sources.into_iter().flatten() { + let windows = provider_pool_collect_quota_windows(source); + if windows.is_empty() { + continue; + } + let observed_at = provider_pool_timestamp_unix_secs(source.get("observed_at")) + .or_else(|| provider_pool_timestamp_unix_secs(source.get("updated_at"))); + let model_matches = windows + .iter() + .filter(|window| { + provider_pool_window_explicitly_matches_model(window, provider_model_name) + }) + .collect::>(); + if !model_matches.is_empty() { + let exhausted = provider_pool_matching_windows_exhausted( + model_matches, + observed_at, + ); + if resolved.is_none() + || provider_pool_should_replace_model_quota_resolution( + resolved.as_ref().and_then(|(observed_at, _)| *observed_at), + observed_at, + ) + { + resolved = Some((observed_at, exhausted)); + } + continue; + } + + // Legacy snapshots may not carry a model field. Match an opaque + // family token (the prefix before `_`/`:` in `code`, or an explicit + // family key) against the model's tokens. This keeps independent + // windows isolated without baking in names such as "spark". + let family_matches = windows + .iter() + .filter(|window| provider_pool_window_family_matches_model(window, &requested)) + .collect::>(); + if !family_matches.is_empty() { + let exhausted = provider_pool_matching_windows_exhausted( + family_matches, + observed_at, + ); + if resolved.is_none() + || provider_pool_should_replace_model_quota_resolution( + resolved.as_ref().and_then(|(observed_at, _)| *observed_at), + observed_at, + ) + { + resolved = Some((observed_at, exhausted)); + } + } + } + + resolved.map(|(_, exhausted)| exhausted) +} + +fn provider_pool_should_replace_model_quota_resolution( + previous_observed_at: Option, + next_observed_at: Option, +) -> bool { + match (previous_observed_at, next_observed_at) { + (Some(previous), Some(next)) => next >= previous, + (None, Some(_)) => true, + (Some(_), None) => false, + // Preserve source order when neither side carries freshness metadata; + // the materialized status snapshot is preferred over raw metadata. + (None, None) => false, + } +} + +/// Public adapter-independent model quota lookup used by schedulers that need +/// to prefilter candidates before constructing provider-pool signals. +pub fn provider_pool_key_model_quota_exhausted( + key: &StoredProviderCatalogKey, + provider_type: &str, + provider_model_name: &str, +) -> Option { + provider_pool_model_quota_exhausted(key, provider_type, provider_model_name) +} + +fn provider_pool_matching_windows_exhausted( + windows: Vec<&Map>, + snapshot_observed_at: Option, +) -> bool { + let now_unix_secs = provider_pool_current_unix_secs(); + + // A request can use a bucket as long as at least one of its independent + // windows still has capacity (e.g. a short and a long rolling window). + windows.iter().all(|window| { + provider_pool_quota_window_is_exhausted(window) + && !now_unix_secs.is_some_and(|now| { + provider_pool_reset_deadline_elapsed(window, snapshot_observed_at, now) + }) + }) +} + +fn provider_pool_window_explicitly_matches_model( + window: &Map, + requested_model: &str, +) -> bool { + let requested = provider_pool_normalize_identifier(requested_model); + let requested_tokens = provider_pool_identifier_tokens(requested_model); + if requested.is_empty() { + return false; + } + [ + "model", + "model_name", + "model_id", + "quota_model", + "quota_model_name", + "target_model", + "limit_name", + ] + .iter() + .filter_map(|key| window.get(*key)) + .any(|value| match value { + Value::String(value) => { + provider_pool_identifiers_match(&requested, value, &requested_tokens) + } + Value::Array(values) => values + .iter() + .filter_map(Value::as_str) + .any(|value| provider_pool_identifiers_match(&requested, value, &requested_tokens)), + _ => false, + }) || ["models", "model_ids"] + .iter() + .filter_map(|key| window.get(*key).and_then(Value::as_array)) + .flatten() + .filter_map(Value::as_str) + .any(|value| provider_pool_identifiers_match(&requested, value, &requested_tokens)) +} + +fn provider_pool_window_family_matches_model( + window: &Map, + requested_tokens: &std::collections::BTreeSet, +) -> bool { + let explicit_scope = window + .get("scope") + .and_then(Value::as_str) + .map(|scope| scope.trim().to_ascii_lowercase()); + // A model-scoped window without an explicit model must not accidentally + // match a token from its opaque code. + if explicit_scope.as_deref() == Some("model") + || provider_pool_window_has_explicit_model(window) + { + return false; + } + + let mut families = Vec::new(); + for key in ["quota_group", "quota_family", "family", "bucket"] { + if let Some(value) = window.get(key).and_then(Value::as_str) { + families.push(value.to_string()); + } + } + if let Some(code) = window.get("code").and_then(Value::as_str) { + let code = code.trim(); + if let Some((prefix, _)) = code.split_once(['_', ':', '/']) { + families.push(prefix.to_string()); + } + } + families.into_iter().any(|family| { + let normalized = provider_pool_normalize_identifier(&family); + if normalized.is_empty() + || [ + "account", + "quota", + "window", + "primary", + "secondary", + "rate", + "reset", + ] + .iter() + .any(|generic| normalized == *generic) + { + return false; + } + requested_tokens.iter().any(|token| { + token.len() >= 3 + && (token == &normalized + || token.contains(&normalized) + || normalized.contains(token)) + }) + }) +} + +fn provider_pool_identifiers_match( + requested: &str, + candidate: &str, + requested_tokens: &std::collections::BTreeSet, +) -> bool { + let candidate_tokens = provider_pool_identifier_tokens(candidate); + let candidate = provider_pool_normalize_identifier(candidate); + if candidate.is_empty() { + return false; + } + requested == candidate + || candidate_tokens + .iter() + .any(|token| { + requested_tokens.contains(token) && provider_pool_is_specific_model_token(token) + }) + // Handle compact upstream identifiers such as `spark` embedded in a + // provider model name (`vendor-codex-spark`) while avoiding accidental + // one/two-character matches. + || (candidate.len() >= 4 + && requested.len() >= 6 + && (requested.contains(&candidate) || candidate.contains(requested))) +} + +fn provider_pool_is_specific_model_token(token: &str) -> bool { + token.len() >= 4 + && !token.chars().all(|character| character.is_ascii_digit()) + && ![ + "auto", "base", "claude", "codex", "default", "fast", "flash", "free", + "gemini", "gpt", "latest", "mini", "model", "plus", "pro", "reasoning", + "team", "think", "thinking", "vendor", + ] + .contains(&token) +} + +fn provider_pool_window_has_explicit_model(window: &Map) -> bool { + [ + "model", + "model_name", + "model_id", + "quota_model", + "quota_model_name", + "target_model", + "models", + "model_ids", + ] + .iter() + .any(|key| match window.get(*key) { + Some(Value::String(value)) => !value.trim().is_empty(), + Some(Value::Array(values)) => values.iter().any(|value| { + value.as_str().is_some_and(|value| !value.trim().is_empty()) + }), + _ => false, + }) +} + +fn provider_pool_normalize_identifier(value: &str) -> String { + value + .trim() + .to_ascii_lowercase() + .chars() + .filter(|character| character.is_ascii_alphanumeric()) + .collect() +} + +fn provider_pool_identifier_tokens(value: &str) -> std::collections::BTreeSet { + value + .split(|character: char| !character.is_ascii_alphanumeric()) + .map(|token| token.trim().to_ascii_lowercase()) + .filter(|token| token.len() >= 3) + .collect() +} + +/// Materialize quota windows from the small set of shapes emitted by current +/// and legacy adapters. Model maps (`quota_by_model`/`models`) are converted +/// to the same window representation used by status snapshots, with the map +/// key retained as the model identity. Keeping this normalization here means +/// provider adapters do not need to grow model-specific quota code whenever an +/// upstream introduces another independent bucket. +fn provider_pool_collect_quota_windows(source: &Map) -> Vec> { + let mut windows = Vec::new(); + for key in ["windows", "additional_quota_windows"] { + if let Some(values) = source.get(key).and_then(Value::as_array) { + windows.extend(values.iter().filter_map(Value::as_object).cloned()); + } + } + for key in ["quota_by_model", "models", "model_quotas"] { + let Some(models) = source.get(key).and_then(Value::as_object) else { + continue; + }; + for (model_name, item) in models { + let Some(item) = item.as_object() else { + continue; + }; + let mut window = item.clone(); + window + .entry("model".to_string()) + .or_insert_with(|| json!(model_name)); + window + .entry("scope".to_string()) + .or_insert_with(|| json!("model")); + window + .entry("code".to_string()) + .or_insert_with(|| json!(format!("model:{model_name}"))); + windows.push(window); + } + } + windows +} + pub fn provider_pool_quota_snapshot_updated_at( key: &StoredProviderCatalogKey, provider_type: &str, @@ -197,12 +542,51 @@ pub(crate) fn provider_pool_reset_deadline_elapsed( fn provider_pool_quota_window_is_exhausted(window: &Map) -> bool { provider_pool_json_bool(window.get("is_exhausted")) + .or_else(|| provider_pool_json_bool(window.get("exhausted"))) .or_else(|| { - provider_pool_json_f64(window.get("used_ratio")).map(|value| value >= 1.0 - 1e-6) + provider_pool_json_f64(window.get("used_ratio").or_else(|| window.get("usage_ratio"))) + .map(|value| value >= 1.0 - 1e-6) + }) + .or_else(|| { + provider_pool_json_f64(window.get("used_percent")).map(|value| value >= 100.0 - 1e-6) + }) + .or_else(|| { + provider_pool_json_f64( + window + .get("remaining_ratio") + .or_else(|| window.get("remaining_fraction")), + ) + .map(|value| value <= 1e-6) + }) + .or_else(|| { + provider_pool_json_f64(window.get("remaining_percent")) + .map(|value| value <= 1e-6) + }) + .or_else(|| { + let remaining = provider_pool_json_f64( + window + .get("remaining") + .or_else(|| window.get("remaining_value")), + )?; + let limit = provider_pool_json_f64( + window + .get("limit") + .or_else(|| window.get("limit_value")) + .or_else(|| window.get("total")), + )?; + (limit > 0.0).then_some(remaining <= 0.0) }) .unwrap_or(false) } +fn provider_pool_window_is_model_scoped(window: &Map) -> bool { + window + .get("scope") + .and_then(Value::as_str) + .is_some_and(|scope| scope.trim().eq_ignore_ascii_case("model")) + || provider_pool_window_has_explicit_model(window) +} + fn provider_pool_quota_snapshot_matches_provider( quota_snapshot: &Map, provider_type: &str, @@ -237,6 +621,18 @@ fn provider_pool_quota_snapshot_matches_provider( .get("windows") .and_then(Value::as_array) .is_some_and(|windows| !windows.is_empty()) + || quota_snapshot + .get("additional_quota_windows") + .and_then(Value::as_array) + .is_some_and(|windows| !windows.is_empty()) + || ["quota_by_model", "models", "model_quotas"] + .iter() + .any(|key| { + quota_snapshot + .get(*key) + .and_then(Value::as_object) + .is_some_and(|models| !models.is_empty()) + }) || quota_snapshot .get("credits") .and_then(Value::as_object) @@ -265,16 +661,28 @@ pub(crate) fn provider_pool_quota_snapshot_exhausted_decision( provider_pool_timestamp_unix_secs(quota_snapshot.get("observed_at")) .or_else(|| provider_pool_timestamp_unix_secs(quota_snapshot.get("updated_at"))); - if let Some(windows) = quota_snapshot - .get("windows") - .and_then(Value::as_array) - .filter(|windows| !windows.is_empty()) - { + let materialized_windows = provider_pool_collect_quota_windows(quota_snapshot); + if !materialized_windows.is_empty() { + // Model-scoped windows are evaluated only when the request model + // is known. They must not turn the account-level fallback into + // an exhausted state for unrelated models. + let account_scoped_windows = materialized_windows + .iter() + .filter(|window| !provider_pool_window_is_model_scoped(window)) + .collect::>(); + // A snapshot containing only model-scoped buckets has no + // account-wide signal to apply when the caller did not provide a + // model name (for example, an admin status listing). Do not let a + // single exhausted model poison every sibling bucket. + if account_scoped_windows.is_empty() { + return Some(false); + } + let windows = account_scoped_windows; let mut saw_exhausted_window = false; let mut saw_active_exhausted_window = false; let mut windows_max_ratio = None::; - for window in windows.iter().filter_map(Value::as_object) { + for window in windows.iter() { if let Some(ratio) = provider_pool_json_f64(window.get("used_ratio")) { windows_max_ratio = Some(windows_max_ratio.map_or(ratio, |current| current.max(ratio))); diff --git a/crates/aether-provider/pool/src/service.rs b/crates/aether-provider/pool/src/service.rs index 403783e6d..f61b07a7f 100644 --- a/crates/aether-provider/pool/src/service.rs +++ b/crates/aether-provider/pool/src/service.rs @@ -123,12 +123,14 @@ impl ProviderPoolService { provider_type: &str, key: &StoredProviderCatalogKey, auth_config: Option<&Map>, + provider_model_name: Option<&str>, ) -> aether_pool_core::PoolMemberSignals { let adapter = self.adapter(provider_type); let input = ProviderPoolMemberInput { provider_type, key, auth_config, + provider_model_name, }; adapter.member_signals(&input) } From 77229943d13ad7882e021cc07e8d082e28c15737 Mon Sep 17 00:00:00 2001 From: zhefox Date: Wed, 2 Sep 2026 15:29:56 +0800 Subject: [PATCH 02/13] fix(pool): preserve generic quota windows for model scheduling --- crates/aether-provider/pool/src/quota.rs | 76 ++++++++++++++++++++++-- 1 file changed, 71 insertions(+), 5 deletions(-) diff --git a/crates/aether-provider/pool/src/quota.rs b/crates/aether-provider/pool/src/quota.rs index a3a1e78a0..c31d99105 100644 --- a/crates/aether-provider/pool/src/quota.rs +++ b/crates/aether-provider/pool/src/quota.rs @@ -106,7 +106,7 @@ pub(crate) fn provider_pool_model_quota_exhausted( }) .collect::>(); if !model_matches.is_empty() { - let exhausted = provider_pool_matching_windows_exhausted( + let exhausted = provider_pool_explicit_model_windows_exhausted( model_matches, observed_at, ); @@ -130,7 +130,7 @@ pub(crate) fn provider_pool_model_quota_exhausted( .filter(|window| provider_pool_window_family_matches_model(window, &requested)) .collect::>(); if !family_matches.is_empty() { - let exhausted = provider_pool_matching_windows_exhausted( + let exhausted = provider_pool_any_window_exhausted( family_matches, observed_at, ); @@ -142,6 +142,28 @@ pub(crate) fn provider_pool_model_quota_exhausted( { resolved = Some((observed_at, exhausted)); } + continue; + } + + // Account-scoped windows (for example the ordinary weekly and short + // windows emitted by Codex) apply to every model that has no more + // specific family. Restrict this fallback to well-known structural + // window names; opaque family names such as `alpha_weekly` must not + // accidentally make an unrelated model look schedulable. + let generic_matches = windows + .iter() + .filter(|window| provider_pool_window_is_generic(window)) + .collect::>(); + if !generic_matches.is_empty() { + let exhausted = provider_pool_any_window_exhausted(generic_matches, observed_at); + if resolved.is_none() + || provider_pool_should_replace_model_quota_resolution( + resolved.as_ref().and_then(|(observed_at, _)| *observed_at), + observed_at, + ) + { + resolved = Some((observed_at, exhausted)); + } } } @@ -172,14 +194,15 @@ pub fn provider_pool_key_model_quota_exhausted( provider_pool_model_quota_exhausted(key, provider_type, provider_model_name) } -fn provider_pool_matching_windows_exhausted( +fn provider_pool_explicit_model_windows_exhausted( windows: Vec<&Map>, snapshot_observed_at: Option, ) -> bool { let now_unix_secs = provider_pool_current_unix_secs(); - // A request can use a bucket as long as at least one of its independent - // windows still has capacity (e.g. a short and a long rolling window). + // Explicit model buckets represent independent windows for one model. The + // model remains usable while at least one of those windows still has + // capacity, so exhaustion is reported only when all are exhausted. windows.iter().all(|window| { provider_pool_quota_window_is_exhausted(window) && !now_unix_secs.is_some_and(|now| { @@ -188,6 +211,49 @@ fn provider_pool_matching_windows_exhausted( }) } +fn provider_pool_any_window_exhausted( + windows: Vec<&Map>, + snapshot_observed_at: Option, +) -> bool { + let now_unix_secs = provider_pool_current_unix_secs(); + windows.iter().any(|window| { + provider_pool_quota_window_is_exhausted(window) + && !now_unix_secs.is_some_and(|now| { + provider_pool_reset_deadline_elapsed(window, snapshot_observed_at, now) + }) + }) +} + +fn provider_pool_window_is_generic(window: &Map) -> bool { + if provider_pool_window_has_explicit_model(window) + || window + .get("scope") + .and_then(Value::as_str) + .is_some_and(|scope| scope.trim().eq_ignore_ascii_case("model")) + { + return false; + } + + let code = window.get("code").and_then(Value::as_str).unwrap_or_default(); + let family = code + .split_once(['_', ':', '/']) + .map(|(prefix, _)| prefix) + .unwrap_or(code) + .trim() + .to_ascii_lowercase(); + if family.is_empty() { + return window + .get("scope") + .and_then(Value::as_str) + .is_some_and(|scope| scope.trim().eq_ignore_ascii_case("account")); + } + [ + "weekly", "5h", "daily", "monthly", "primary", "secondary", "account", "quota", + "window", "rate", "reset", + ] + .contains(&family.as_str()) +} + fn provider_pool_window_explicitly_matches_model( window: &Map, requested_model: &str, From b1bf7837cf5244664ea20a0c650d00416766e54f Mon Sep 17 00:00:00 2001 From: zhefox Date: Wed, 2 Sep 2026 15:40:40 +0800 Subject: [PATCH 03/13] style: format quota and response changes --- .../formats/src/formats/openai/chat/stream.rs | 3 +- .../src/formats/openai/responses/mod.rs | 17 ++--- .../src/formats/shared/sync_products.rs | 1 - crates/aether-ai/formats/src/lib.rs | 3 +- .../formats/src/protocol/canonical.rs | 2 +- crates/aether-provider/pool/src/lib.rs | 30 +++----- crates/aether-provider/pool/src/provider.rs | 4 +- .../pool/src/providers/chatgpt_web.rs | 4 +- .../pool/src/providers/grok.rs | 4 +- crates/aether-provider/pool/src/quota.rs | 72 +++++++++++++------ 10 files changed, 73 insertions(+), 67 deletions(-) diff --git a/crates/aether-ai/formats/src/formats/openai/chat/stream.rs b/crates/aether-ai/formats/src/formats/openai/chat/stream.rs index fe5f8f4ec..22dc3e502 100644 --- a/crates/aether-ai/formats/src/formats/openai/chat/stream.rs +++ b/crates/aether-ai/formats/src/formats/openai/chat/stream.rs @@ -4,8 +4,7 @@ use serde_json::{json, Map, Value}; use crate::formats::openai::namespace::NamespaceToolAliases; use crate::formats::openai::responses::{ - encode_gemini_tool_signature_carrier_with_direction, - openai_responses_message_item_id, + encode_gemini_tool_signature_carrier_with_direction, openai_responses_message_item_id, openai_responses_synthetic_reasoning_item_id, response::{ ensure_modern_openai_responses_response_fields, openai_responses_current_timestamp, diff --git a/crates/aether-ai/formats/src/formats/openai/responses/mod.rs b/crates/aether-ai/formats/src/formats/openai/responses/mod.rs index f775841a5..9a9328ff8 100644 --- a/crates/aether-ai/formats/src/formats/openai/responses/mod.rs +++ b/crates/aether-ai/formats/src/formats/openai/responses/mod.rs @@ -318,8 +318,7 @@ mod tests { use super::{ decode_gemini_tool_signature_carrier, encode_gemini_tool_signature_carrier_with_direction, normalize_openai_responses_message_item_ids, openai_responses_message_item_id, - openai_responses_request_operation, - openai_responses_synthetic_reasoning_item_id, + openai_responses_request_operation, openai_responses_synthetic_reasoning_item_id, strip_incompatible_openai_responses_reasoning_items, strip_incompatible_openai_responses_reasoning_items_with_policy, GeminiToolSignatureCarrierDirection, OpenAiResponsesReasoningReplayPolicy, @@ -414,14 +413,8 @@ mod tests { #[test] fn synthetic_message_item_ids_are_stable_and_start_with_msg() { - let first = openai_responses_message_item_id( - "1c938e58-32a8-4d28-9c34-538d78076895", - 0, - ); - let second = openai_responses_message_item_id( - "1c938e58-32a8-4d28-9c34-538d78076895", - 0, - ); + let first = openai_responses_message_item_id("1c938e58-32a8-4d28-9c34-538d78076895", 0); + let second = openai_responses_message_item_id("1c938e58-32a8-4d28-9c34-538d78076895", 0); let other = openai_responses_message_item_id("chatcmpl-123", 1); assert!(first.starts_with("msg_")); @@ -442,7 +435,9 @@ mod tests { assert_eq!(normalize_openai_responses_message_item_ids(&mut body), 1); let input = body["input"].as_array().expect("input array"); - assert!(input[0]["id"].as_str().is_some_and(|id| id.starts_with("msg_"))); + assert!(input[0]["id"] + .as_str() + .is_some_and(|id| id.starts_with("msg_"))); assert_eq!(input[1]["id"], "msg_provider_123"); assert_eq!(input[2].get("id"), Some(&json!("legacy_call"))); assert!(input[3].get("id").is_none()); diff --git a/crates/aether-ai/formats/src/formats/shared/sync_products.rs b/crates/aether-ai/formats/src/formats/shared/sync_products.rs index d429b8057..96dbc4ea1 100644 --- a/crates/aether-ai/formats/src/formats/shared/sync_products.rs +++ b/crates/aether-ai/formats/src/formats/shared/sync_products.rs @@ -10,7 +10,6 @@ use aether_ai_formats::formats::conversion::response::{ use aether_ai_formats::formats::openai::responses::{ openai_responses_message_item_id, openai_responses_synthetic_reasoning_item_id, }; -use aether_ai_formats::formats::openai::responses::response::ensure_modern_openai_responses_response_fields; use aether_ai_formats::formats::registry::{convert_response, FormatContext, FormatError}; use aether_ai_formats::{ canonical_response_unknown_block_count, canonical_to_claude_response, diff --git a/crates/aether-ai/formats/src/lib.rs b/crates/aether-ai/formats/src/lib.rs index 0f141368c..be40c6d6f 100644 --- a/crates/aether-ai/formats/src/lib.rs +++ b/crates/aether-ai/formats/src/lib.rs @@ -58,8 +58,7 @@ pub use formats::openai::responses::request::{ }; pub use formats::openai::responses::{ normalize_openai_responses_message_item_ids, openai_responses_message_item_id, - openai_responses_request_operation, - openai_responses_synthetic_reasoning_item_id, + openai_responses_request_operation, openai_responses_synthetic_reasoning_item_id, strip_incompatible_openai_responses_reasoning_items, strip_incompatible_openai_responses_reasoning_items_with_policy, OpenAiResponsesReasoningReplayPolicy, OPENAI_RESPONSES_OPERATION_COMPACT, diff --git a/crates/aether-ai/formats/src/protocol/canonical.rs b/crates/aether-ai/formats/src/protocol/canonical.rs index 51863f7ed..bf4db8aab 100644 --- a/crates/aether-ai/formats/src/protocol/canonical.rs +++ b/crates/aether-ai/formats/src/protocol/canonical.rs @@ -3,11 +3,11 @@ use std::collections::{BTreeMap, BTreeSet, VecDeque}; use serde::{Deserialize, Serialize}; use serde_json::{json, Map, Value}; +use crate::formats::openai::responses::openai_responses_message_item_id; use crate::formats::openai::responses::{ decode_gemini_tool_signature_carrier, GeminiToolSignatureCarrierDirection, }; use crate::formats::openai::shared::map_thinking_budget_to_openai_reasoning_effort; -use crate::formats::openai::responses::openai_responses_message_item_id; use crate::formats::shared::model_directives::ReasoningEffort; use crate::formats::shared::response::remove_empty_pages_from_tool_input_value; diff --git a/crates/aether-provider/pool/src/lib.rs b/crates/aether-provider/pool/src/lib.rs index ceb921a8f..f8c813614 100644 --- a/crates/aether-provider/pool/src/lib.rs +++ b/crates/aether-provider/pool/src/lib.rs @@ -34,8 +34,8 @@ pub use providers::{ WINDSURF_MODEL_CONFIGS_PATH, WINDSURF_RATE_LIMIT_PATH, WINDSURF_USER_STATUS_PATH, }; pub use quota::{ - provider_pool_key_account_quota_exhausted, provider_pool_key_quota_hard_blocked, - provider_pool_key_model_quota_exhausted, provider_pool_key_model_quota_hard_blocked, + provider_pool_key_account_quota_exhausted, provider_pool_key_model_quota_exhausted, + provider_pool_key_model_quota_hard_blocked, provider_pool_key_quota_hard_blocked, provider_pool_key_scheduling_label, provider_pool_member_quota_snapshot, provider_pool_quota_metadata_provider_type, provider_pool_quota_metadata_updated_at, provider_pool_quota_snapshot_updated_at, @@ -1005,14 +1005,12 @@ mod tests { } })); - let alpha = service.member_signals( - "codex", - &key, - None, - Some("vendor-alpha-model"), - ); + let alpha = service.member_signals("codex", &key, None, Some("vendor-alpha-model")); let beta = service.member_signals("codex", &key, None, Some("vendor-beta-model")); - assert!(!alpha.quota_exhausted, "one available alpha window must keep it usable"); + assert!( + !alpha.quota_exhausted, + "one available alpha window must keep it usable" + ); assert!(beta.quota_exhausted); assert!(!beta.quota_hard_blocked); } @@ -1062,12 +1060,7 @@ mod tests { } })); - let signals = service.member_signals( - "codex", - &key, - None, - Some("gpt-5.3-codex-spark"), - ); + let signals = service.member_signals("codex", &key, None, Some("gpt-5.3-codex-spark")); assert!( !signals.quota_exhausted, "a compact bucket name should match a token in the selected model" @@ -1093,12 +1086,7 @@ mod tests { } })); - let signals = service.member_signals( - "codex", - &key, - None, - Some("gpt-5.4-codex-spark"), - ); + let signals = service.member_signals("codex", &key, None, Some("gpt-5.4-codex-spark")); assert!(signals.quota_exhausted); } diff --git a/crates/aether-provider/pool/src/provider.rs b/crates/aether-provider/pool/src/provider.rs index 87f54e7f5..47cace82e 100644 --- a/crates/aether-provider/pool/src/provider.rs +++ b/crates/aether-provider/pool/src/provider.rs @@ -7,8 +7,8 @@ use serde_json::{Map, Value}; use crate::capability::{ProviderPoolCapabilities, ProviderPoolCapability}; use crate::plan::{derive_plan_tier, normalize_provider_plan_tier}; use crate::quota::{ - provider_pool_account_blocked, provider_pool_quota_reset_seconds, - provider_pool_model_quota_exhausted, provider_pool_quota_snapshot_exhausted_decision, + provider_pool_account_blocked, provider_pool_model_quota_exhausted, + provider_pool_quota_reset_seconds, provider_pool_quota_snapshot_exhausted_decision, provider_pool_quota_usage_ratio, }; diff --git a/crates/aether-provider/pool/src/providers/chatgpt_web.rs b/crates/aether-provider/pool/src/providers/chatgpt_web.rs index 3eb4b5754..574e2f107 100644 --- a/crates/aether-provider/pool/src/providers/chatgpt_web.rs +++ b/crates/aether-provider/pool/src/providers/chatgpt_web.rs @@ -11,8 +11,8 @@ use crate::provider::{ }; use crate::quota::{ provider_pool_current_unix_secs, provider_pool_json_bool, provider_pool_json_f64, - provider_pool_metadata_bucket, provider_pool_quota_snapshot_exhausted_decision, - provider_pool_model_quota_exhausted, provider_pool_reset_deadline_elapsed, + provider_pool_metadata_bucket, provider_pool_model_quota_exhausted, + provider_pool_quota_snapshot_exhausted_decision, provider_pool_reset_deadline_elapsed, provider_pool_timestamp_unix_secs, }; use crate::quota_refresh::ProviderPoolQuotaRequestSpec; diff --git a/crates/aether-provider/pool/src/providers/grok.rs b/crates/aether-provider/pool/src/providers/grok.rs index ce9e2f91b..460e927cd 100644 --- a/crates/aether-provider/pool/src/providers/grok.rs +++ b/crates/aether-provider/pool/src/providers/grok.rs @@ -8,8 +8,8 @@ use crate::provider::{ }; use crate::quota::{ provider_pool_current_unix_secs, provider_pool_json_bool, provider_pool_json_f64, - provider_pool_metadata_bucket, provider_pool_quota_snapshot_exhausted_decision, - provider_pool_model_quota_exhausted, provider_pool_reset_deadline_elapsed, + provider_pool_metadata_bucket, provider_pool_model_quota_exhausted, + provider_pool_quota_snapshot_exhausted_decision, provider_pool_reset_deadline_elapsed, provider_pool_timestamp_unix_secs, }; diff --git a/crates/aether-provider/pool/src/quota.rs b/crates/aether-provider/pool/src/quota.rs index c31d99105..f19200fc8 100644 --- a/crates/aether-provider/pool/src/quota.rs +++ b/crates/aether-provider/pool/src/quota.rs @@ -106,10 +106,8 @@ pub(crate) fn provider_pool_model_quota_exhausted( }) .collect::>(); if !model_matches.is_empty() { - let exhausted = provider_pool_explicit_model_windows_exhausted( - model_matches, - observed_at, - ); + let exhausted = + provider_pool_explicit_model_windows_exhausted(model_matches, observed_at); if resolved.is_none() || provider_pool_should_replace_model_quota_resolution( resolved.as_ref().and_then(|(observed_at, _)| *observed_at), @@ -130,10 +128,7 @@ pub(crate) fn provider_pool_model_quota_exhausted( .filter(|window| provider_pool_window_family_matches_model(window, &requested)) .collect::>(); if !family_matches.is_empty() { - let exhausted = provider_pool_any_window_exhausted( - family_matches, - observed_at, - ); + let exhausted = provider_pool_any_window_exhausted(family_matches, observed_at); if resolved.is_none() || provider_pool_should_replace_model_quota_resolution( resolved.as_ref().and_then(|(observed_at, _)| *observed_at), @@ -234,7 +229,10 @@ fn provider_pool_window_is_generic(window: &Map) -> bool { return false; } - let code = window.get("code").and_then(Value::as_str).unwrap_or_default(); + let code = window + .get("code") + .and_then(Value::as_str) + .unwrap_or_default(); let family = code .split_once(['_', ':', '/']) .map(|(prefix, _)| prefix) @@ -248,8 +246,17 @@ fn provider_pool_window_is_generic(window: &Map) -> bool { .is_some_and(|scope| scope.trim().eq_ignore_ascii_case("account")); } [ - "weekly", "5h", "daily", "monthly", "primary", "secondary", "account", "quota", - "window", "rate", "reset", + "weekly", + "5h", + "daily", + "monthly", + "primary", + "secondary", + "account", + "quota", + "window", + "rate", + "reset", ] .contains(&family.as_str()) } @@ -301,8 +308,7 @@ fn provider_pool_window_family_matches_model( .map(|scope| scope.trim().to_ascii_lowercase()); // A model-scoped window without an explicit model must not accidentally // match a token from its opaque code. - if explicit_scope.as_deref() == Some("model") - || provider_pool_window_has_explicit_model(window) + if explicit_scope.as_deref() == Some("model") || provider_pool_window_has_explicit_model(window) { return false; } @@ -373,9 +379,26 @@ fn provider_pool_is_specific_model_token(token: &str) -> bool { token.len() >= 4 && !token.chars().all(|character| character.is_ascii_digit()) && ![ - "auto", "base", "claude", "codex", "default", "fast", "flash", "free", - "gemini", "gpt", "latest", "mini", "model", "plus", "pro", "reasoning", - "team", "think", "thinking", "vendor", + "auto", + "base", + "claude", + "codex", + "default", + "fast", + "flash", + "free", + "gemini", + "gpt", + "latest", + "mini", + "model", + "plus", + "pro", + "reasoning", + "team", + "think", + "thinking", + "vendor", ] .contains(&token) } @@ -394,9 +417,9 @@ fn provider_pool_window_has_explicit_model(window: &Map) -> bool .iter() .any(|key| match window.get(*key) { Some(Value::String(value)) => !value.trim().is_empty(), - Some(Value::Array(values)) => values.iter().any(|value| { - value.as_str().is_some_and(|value| !value.trim().is_empty()) - }), + Some(Value::Array(values)) => values + .iter() + .any(|value| value.as_str().is_some_and(|value| !value.trim().is_empty())), _ => false, }) } @@ -610,8 +633,12 @@ fn provider_pool_quota_window_is_exhausted(window: &Map) -> bool provider_pool_json_bool(window.get("is_exhausted")) .or_else(|| provider_pool_json_bool(window.get("exhausted"))) .or_else(|| { - provider_pool_json_f64(window.get("used_ratio").or_else(|| window.get("usage_ratio"))) - .map(|value| value >= 1.0 - 1e-6) + provider_pool_json_f64( + window + .get("used_ratio") + .or_else(|| window.get("usage_ratio")), + ) + .map(|value| value >= 1.0 - 1e-6) }) .or_else(|| { provider_pool_json_f64(window.get("used_percent")).map(|value| value >= 100.0 - 1e-6) @@ -625,8 +652,7 @@ fn provider_pool_quota_window_is_exhausted(window: &Map) -> bool .map(|value| value <= 1e-6) }) .or_else(|| { - provider_pool_json_f64(window.get("remaining_percent")) - .map(|value| value <= 1e-6) + provider_pool_json_f64(window.get("remaining_percent")).map(|value| value <= 1e-6) }) .or_else(|| { let remaining = provider_pool_json_f64( From 2ed2cc66efc25c86bb07c97eaa0e543b6e9cb72c Mon Sep 17 00:00:00 2001 From: zhefox Date: Wed, 2 Sep 2026 15:44:15 +0800 Subject: [PATCH 04/13] fix(formats): retain response field normalization import --- crates/aether-ai/formats/src/formats/shared/sync_products.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/crates/aether-ai/formats/src/formats/shared/sync_products.rs b/crates/aether-ai/formats/src/formats/shared/sync_products.rs index 96dbc4ea1..8424884fa 100644 --- a/crates/aether-ai/formats/src/formats/shared/sync_products.rs +++ b/crates/aether-ai/formats/src/formats/shared/sync_products.rs @@ -7,6 +7,7 @@ use aether_ai_formats::formats::conversion::response::{ convert_openai_chat_response_to_openai_responses, convert_openai_responses_response_to_openai_chat, }; +use aether_ai_formats::formats::openai::responses::response::ensure_modern_openai_responses_response_fields; use aether_ai_formats::formats::openai::responses::{ openai_responses_message_item_id, openai_responses_synthetic_reasoning_item_id, }; From 45c840b8d30612b51f6b375a9a10e23ff3097abd Mon Sep 17 00:00:00 2001 From: ZheFox <77232781+zhefox@users.noreply.github.com> Date: Thu, 3 Sep 2026 10:46:44 +0800 Subject: [PATCH 05/13] fix(providers): refresh Antigravity grouped quotas --- .../standard/openai/chat/decision/request.rs | 4 +- .../admin/provider/oauth/quota/antigravity.rs | 114 ++++++++- .../src/handlers/shared/catalog.rs | 223 +++++++++++++++++- .../src/tests/ai_execute/finalize_local.rs | 2 +- .../finalize_local_cli/cross_format.rs | 2 +- .../finalize_local_provider/gemini.rs | 2 +- .../stream_provider_gemini/local_cli.rs | 2 +- .../src/tests/ai_execute/sync/gemini/cli.rs | 2 +- .../tests/control/admin/endpoints/quota.rs | 164 ++++++++++--- crates/aether-admin/src/provider/quota.rs | 175 +++++++++++++- crates/aether-model-fetch/src/logic.rs | 22 ++ crates/aether-model-fetch/src/strategy.rs | 92 ++++++-- crates/aether-model-fetch/src/transport.rs | 2 +- crates/aether-provider/pool/src/lib.rs | 80 ++++++- .../pool/src/providers/antigravity.rs | 84 +++++++ .../aether-provider/pool/src/providers/mod.rs | 3 +- crates/aether-provider/pool/src/quota.rs | 1 + .../transport/src/antigravity/auth.rs | 23 +- .../components/AntigravityQuotaDialog.vue | 48 +++- .../__tests__/AntigravityQuotaDialog.spec.ts | 34 ++- 20 files changed, 976 insertions(+), 103 deletions(-) diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs index 4300a1733..36e3d8b4c 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs @@ -2741,7 +2741,7 @@ mod tests { .provider_request_headers .get("x-client-version") .map(String::as_str), - Some("1.2.3") + Some("4.3.0") ); assert_eq!( payload @@ -2761,7 +2761,7 @@ mod tests { assert_eq!(payload.provider_request_body["model"], "gemini-2.5-pro"); assert_eq!( payload.provider_request_body["userAgent"], - "antigravity/cli/1.0.16 (aidev_client; os_type=linux; arch=arm64; auth_method=consumer)" + "vscode/1.X.X (Antigravity/4.3.0)" ); assert_eq!(payload.provider_request_body["requestType"], "agent"); assert!(payload.provider_request_body.get("contents").is_none()); diff --git a/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/antigravity.rs b/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/antigravity.rs index 3e60995f1..d3e26fc79 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/antigravity.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/antigravity.rs @@ -7,15 +7,20 @@ use super::shared::{ }; use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot}; use crate::GatewayError; -use aether_admin::provider::quota::parse_antigravity_usage_response; +use aether_admin::provider::quota::{ + parse_antigravity_quota_summary_response, parse_antigravity_usage_response, +}; use aether_contracts::ProxySnapshot; use aether_data_contracts::repository::provider_catalog::{ StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider, }; -use aether_provider_pool::build_antigravity_pool_quota_request; +use aether_provider_pool::{ + build_antigravity_pool_quota_request, build_antigravity_pool_quota_summary_request, +}; use serde_json::json; use std::collections::BTreeMap; use std::time::{SystemTime, UNIX_EPOCH}; +use tracing::warn; async fn execute_antigravity_quota_plan( state: &AdminAppState<'_>, @@ -55,6 +60,79 @@ async fn execute_antigravity_quota_plan( execute_provider_quota_plan(state, transport, plan, "antigravity").await } +async fn fetch_antigravity_quota_summary_best_effort( + state: &AdminAppState<'_>, + transport: &AdminGatewayProviderTransportSnapshot, + authorization: (String, String), + project_id: &str, + identity_headers: BTreeMap, + proxy_override: Option<&ProxySnapshot>, +) -> Option { + let mut request_project_id = Some(project_id); + + loop { + let proxy = match proxy_override { + Some(proxy) => Some(proxy.clone()), + None => { + state + .resolve_transport_proxy_snapshot_with_tunnel_affinity(transport) + .await + } + }; + let timeouts = Some(resolve_provider_quota_execution_timeouts( + state.resolve_transport_execution_timeouts(transport), + proxy.as_ref(), + )); + let spec = build_antigravity_pool_quota_summary_request( + &transport.key.id, + &transport.endpoint.base_url, + authorization.clone(), + request_project_id, + identity_headers.clone(), + ); + let plan = build_provider_quota_execution_plan( + transport, + spec, + proxy, + state.resolve_transport_profile(transport), + timeouts, + ); + let outcome = match execute_provider_quota_plan(state, transport, plan, "antigravity").await + { + Ok(outcome) => outcome, + Err(error) => { + warn!(error = ?error, "Antigravity grouped quota request failed"); + return None; + } + }; + let result = match outcome { + ProviderQuotaExecutionOutcome::Response(result) => result, + ProviderQuotaExecutionOutcome::Failure(detail) => { + warn!(detail = %detail, "Antigravity grouped quota execution failed"); + return None; + } + }; + + if result.status_code == 200 { + return result + .body + .as_ref() + .and_then(|body| body.json_body.as_ref()) + .and_then(parse_antigravity_quota_summary_response); + } + if result.status_code == 403 && request_project_id.is_some() { + request_project_id = None; + continue; + } + + warn!( + status_code = result.status_code, + "Antigravity grouped quota request returned a non-success status" + ); + return None; + } +} + pub(crate) async fn refresh_antigravity_provider_quota_locally( state: &AdminAppState<'_>, provider: &StoredProviderCatalogProvider, @@ -130,9 +208,9 @@ pub(crate) async fn refresh_antigravity_provider_quota_locally( let result = match execute_antigravity_quota_plan( state, &transport, - authorization, + authorization.clone(), &project_id, - identity_headers, + identity_headers.clone(), proxy_override.as_ref(), ) .await? @@ -168,9 +246,31 @@ pub(crate) async fn refresh_antigravity_provider_quota_locally( .as_ref() .and_then(|body| body.json_body.as_ref()) { - metadata_update = parse_antigravity_usage_response(body_json, now_unix_secs) - .map(|metadata| json!({ "antigravity": metadata })); - if metadata_update.is_some() { + if let Some(mut metadata) = + parse_antigravity_usage_response(body_json, now_unix_secs) + { + if let Some(metadata) = metadata.as_object_mut() { + metadata.insert("project_id".to_string(), json!(project_id)); + } + if let Some(quota_groups) = fetch_antigravity_quota_summary_best_effort( + state, + &transport, + authorization, + &project_id, + identity_headers, + proxy_override.as_ref(), + ) + .await + { + if let Some(metadata) = metadata.as_object_mut() { + metadata.insert("quota_groups".to_string(), quota_groups); + metadata.insert( + "quota_groups_updated_at".to_string(), + json!(now_unix_secs), + ); + } + } + metadata_update = Some(json!({ "antigravity": metadata })); status = "success".to_string(); } else { status = "no_metadata".to_string(); diff --git a/apps/aether-gateway/src/handlers/shared/catalog.rs b/apps/aether-gateway/src/handlers/shared/catalog.rs index fe8791567..bd842358c 100644 --- a/apps/aether-gateway/src/handlers/shared/catalog.rs +++ b/apps/aether-gateway/src/handlers/shared/catalog.rs @@ -654,6 +654,112 @@ fn antigravity_model_quota_window_snapshot( Some(window) } +fn antigravity_grouped_quota_window_snapshots( + metadata: &Map, + observed_at_unix_secs: Option, +) -> Vec { + let mut windows = Vec::new(); + let Some(groups) = metadata.get("quota_groups").and_then(Value::as_array) else { + return windows; + }; + + for (group_index, group) in groups.iter().filter_map(Value::as_object).enumerate() { + let group_code = group + .get("group_id") + .or_else(|| group.get("groupId")) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) + .unwrap_or_else(|| format!("group:{group_index}")); + let group_label = group + .get("display_name") + .or_else(|| group.get("displayName")) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) + .unwrap_or_else(|| format!("Quota group {}", group_index + 1)); + let group_description = group + .get("description") + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned); + + for (bucket_index, bucket) in group + .get("buckets") + .and_then(Value::as_array) + .into_iter() + .flatten() + .filter_map(Value::as_object) + .enumerate() + { + let bucket_id = bucket + .get("bucket_id") + .or_else(|| bucket.get("bucketId")) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) + .unwrap_or_else(|| format!("bucket-{}", bucket_index + 1)); + let Some(mut window) = + model_quota_window_snapshot(&bucket_id, bucket, observed_at_unix_secs) + else { + continue; + }; + let Some(window) = window.as_object_mut() else { + continue; + }; + let period = bucket + .get("window") + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned); + let bucket_detail = bucket + .get("display_name") + .or_else(|| bucket.get("displayName")) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) + .or_else(|| period.clone()); + let bucket_label = bucket_detail + .filter(|detail| !detail.eq_ignore_ascii_case(&group_label)) + .map(|detail| format!("{group_label} · {detail}")) + .unwrap_or_else(|| group_label.clone()); + let description = bucket + .get("description") + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) + .or_else(|| group_description.clone()); + + window.insert( + "code".to_string(), + json!(format!("group:{group_index}:{bucket_id}")), + ); + window.insert("label".to_string(), json!(bucket_label)); + window.insert("scope".to_string(), json!("quota_group")); + window.remove("model"); + window.insert("quota_group".to_string(), json!(group_code)); + window.insert("quota_group_label".to_string(), json!(group_label)); + window.insert("bucket_id".to_string(), json!(bucket_id)); + if let Some(period) = period { + window.insert("window".to_string(), json!(period)); + } + if let Some(description) = description { + window.insert("description".to_string(), json!(description)); + } + windows.push(Value::Object(window.clone())); + } + } + + windows +} + fn provider_quota_metadata_string( metadata: &Map, fields: &[&str], @@ -1622,7 +1728,7 @@ fn build_antigravity_quota_status_snapshot( .map(str::trim) .filter(|value| !value.is_empty()) .map(ToOwned::to_owned); - let windows = provider_quota_model_bucket(metadata) + let mut windows = provider_quota_model_bucket(metadata) .map(|models| { models .iter() @@ -1636,6 +1742,13 @@ fn build_antigravity_quota_status_snapshot( .collect::>() }) .unwrap_or_default(); + let grouped_observed_at_unix_secs = + provider_quota_timestamp_unix_secs(metadata.get("quota_groups_updated_at")) + .or(observed_at_unix_secs); + windows.extend(antigravity_grouped_quota_window_snapshots( + metadata, + grouped_observed_at_unix_secs, + )); if windows.is_empty() && observed_at_unix_secs.is_none() && !is_forbidden { return None; @@ -3132,6 +3245,52 @@ mod tests { assert_eq!(spark_weekly.get("remaining_ratio"), Some(&json!(0.95))); } + #[test] + fn codex_wham_snapshot_does_not_duplicate_normalized_spark_windows() { + let codex = admin_provider_quota_pure::parse_codex_wham_usage_response( + &json!({ + "plan_type": "plus", + "rate_limit": { + "primary_window": {"used_percent": 25.0}, + "secondary_window": {"used_percent": 10.0} + }, + "additional_rate_limits": [{ + "limit_name": "GPT-5.3-Codex-Spark", + "rate_limit": { + "primary_window": { + "used_percent": 40.0, + "limit_window_seconds": 18_000 + }, + "secondary_window": { + "used_percent": 5.0, + "limit_window_seconds": 604_800 + } + } + }] + }), + 1_777_000_000, + ) + .expect("Codex WHAM quota should parse"); + let upstream_metadata = json!({"codex": codex}); + let payload = sync_provider_key_quota_status_snapshot( + None, + "codex", + Some(&upstream_metadata), + "refresh_api", + ) + .expect("Codex quota snapshot should sync"); + let windows = payload["quota"]["windows"] + .as_array() + .expect("quota windows should exist"); + let spark_codes = windows + .iter() + .filter_map(|window| window.get("code").and_then(Value::as_str)) + .filter(|code| code.starts_with("spark_")) + .collect::>(); + + assert_eq!(spark_codes, vec!["spark_5h", "spark_weekly"]); + } + #[test] fn provider_key_status_snapshot_payload_keeps_codex_free_window_quota_available() { let mut key = sample_catalog_key(); @@ -4166,6 +4325,68 @@ mod tests { assert_eq!(claude_window.get("reset_seconds"), Some(&json!(12_011u64))); } + #[test] + fn sync_provider_key_quota_status_snapshot_materializes_antigravity_group_windows() { + let upstream_metadata = json!({ + "antigravity": { + "updated_at": 1_777_000_000u64, + "models": { + "gemini-3.7-flash-tiered": { + "remaining_fraction": 0.9 + } + }, + "quota_groups": [{ + "display_name": "Claude and GPT models", + "description": "Shared quota", + "buckets": [{ + "bucket_id": "3p-5h", + "window": "5h", + "remaining_fraction": 0.25, + "reset_time": "2026-05-05T05:00:00Z", + "display_name": "5 hour" + }, { + "bucket_id": "3p-weekly", + "window": "weekly", + "remaining_fraction": 0.8, + "reset_time": "2026-05-11T00:00:00Z" + }] + }] + } + }); + + let payload = sync_provider_key_quota_status_snapshot( + None, + "antigravity", + Some(&upstream_metadata), + "refresh_api", + ) + .expect("Antigravity quota snapshot should sync"); + let windows = payload["quota"]["windows"] + .as_array() + .expect("quota windows should exist"); + let five_hour = windows + .iter() + .find(|window| window["code"] == "group:0:3p-5h") + .expect("5h grouped quota window should exist"); + let weekly = windows + .iter() + .find(|window| window["code"] == "group:0:3p-weekly") + .expect("weekly grouped quota window should exist"); + + assert_eq!(windows.len(), 3); + assert_eq!(five_hour["scope"], json!("quota_group")); + assert_eq!(five_hour["bucket_id"], json!("3p-5h")); + assert_eq!(five_hour["label"], json!("Claude and GPT models · 5 hour")); + assert_eq!( + five_hour["quota_group_label"], + json!("Claude and GPT models") + ); + assert_eq!(five_hour["remaining_ratio"], json!(0.25)); + assert_eq!(five_hour["used_ratio"], json!(0.75)); + assert!(five_hour.get("model").is_none()); + assert_eq!(weekly["window"], json!("weekly")); + } + #[test] fn provider_key_status_snapshot_payload_backfills_account_block_from_oauth_invalid_reason() { let mut key = sample_catalog_key(); diff --git a/apps/aether-gateway/src/tests/ai_execute/finalize_local.rs b/apps/aether-gateway/src/tests/ai_execute/finalize_local.rs index c2103c842..2a699a237 100644 --- a/apps/aether-gateway/src/tests/ai_execute/finalize_local.rs +++ b/apps/aether-gateway/src/tests/ai_execute/finalize_local.rs @@ -1949,7 +1949,7 @@ async fn gateway_executes_openai_chat_antigravity_cross_format_sync_via_local_fi "Bearer imported-antigravity-chat-token" ); assert_eq!(seen_execution_runtime_request.x_client_name, "antigravity"); - assert_eq!(seen_execution_runtime_request.x_client_version, "1.2.3"); + assert_eq!(seen_execution_runtime_request.x_client_version, "4.3.0"); assert_eq!( seen_execution_runtime_request.x_vscode_sessionid, "sess-antigravity-chat-local-123" diff --git a/apps/aether-gateway/src/tests/ai_execute/finalize_local_cli/cross_format.rs b/apps/aether-gateway/src/tests/ai_execute/finalize_local_cli/cross_format.rs index 03471fe8b..0f1a1338c 100644 --- a/apps/aether-gateway/src/tests/ai_execute/finalize_local_cli/cross_format.rs +++ b/apps/aether-gateway/src/tests/ai_execute/finalize_local_cli/cross_format.rs @@ -1537,7 +1537,7 @@ async fn gateway_executes_openai_responses_antigravity_cross_format_upstream_str ); assert_eq!( seen_remote_execution_runtime_request.x_client_version, - "1.2.3" + "4.3.0" ); assert_eq!( seen_remote_execution_runtime_request.x_vscode_sessionid, diff --git a/apps/aether-gateway/src/tests/ai_execute/finalize_local_provider/gemini.rs b/apps/aether-gateway/src/tests/ai_execute/finalize_local_provider/gemini.rs index 8fa7f5745..509186eab 100644 --- a/apps/aether-gateway/src/tests/ai_execute/finalize_local_provider/gemini.rs +++ b/apps/aether-gateway/src/tests/ai_execute/finalize_local_provider/gemini.rs @@ -2123,7 +2123,7 @@ async fn gateway_executes_antigravity_gemini_cli_sync_upstream_stream_via_local_ ); assert_eq!( seen_remote_execution_runtime_request.x_client_version, - "1.2.3" + "4.3.0" ); assert_eq!( seen_remote_execution_runtime_request.x_vscode_sessionid, diff --git a/apps/aether-gateway/src/tests/ai_execute/stream_provider_gemini/local_cli.rs b/apps/aether-gateway/src/tests/ai_execute/stream_provider_gemini/local_cli.rs index 7b7d98ad1..5de65b509 100644 --- a/apps/aether-gateway/src/tests/ai_execute/stream_provider_gemini/local_cli.rs +++ b/apps/aether-gateway/src/tests/ai_execute/stream_provider_gemini/local_cli.rs @@ -2017,7 +2017,7 @@ async fn gateway_executes_antigravity_gemini_cli_stream_via_local_decision_gate_ "Bearer refreshed-antigravity-cli-stream-access-token" ); assert_eq!(seen_execution_runtime_request.x_client_name, "antigravity"); - assert_eq!(seen_execution_runtime_request.x_client_version, "1.2.3"); + assert_eq!(seen_execution_runtime_request.x_client_version, "4.3.0"); assert_eq!( seen_execution_runtime_request.x_vscode_sessionid, "sess-antigravity-stream-local-123" diff --git a/apps/aether-gateway/src/tests/ai_execute/sync/gemini/cli.rs b/apps/aether-gateway/src/tests/ai_execute/sync/gemini/cli.rs index 770bd5846..6beb48a01 100644 --- a/apps/aether-gateway/src/tests/ai_execute/sync/gemini/cli.rs +++ b/apps/aether-gateway/src/tests/ai_execute/sync/gemini/cli.rs @@ -2321,7 +2321,7 @@ async fn gateway_executes_antigravity_gemini_cli_sync_via_local_decision_gate_af "Bearer refreshed-antigravity-cli-access-token" ); assert_eq!(seen_execution_runtime_request.x_client_name, "antigravity"); - assert_eq!(seen_execution_runtime_request.x_client_version, "1.2.3"); + assert_eq!(seen_execution_runtime_request.x_client_version, "4.3.0"); assert_eq!( seen_execution_runtime_request.x_vscode_sessionid, "sess-antigravity-local-123" diff --git a/apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs b/apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs index 0c6bc508b..c4d52ce31 100644 --- a/apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs +++ b/apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs @@ -2264,6 +2264,8 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru struct SeenExecutionRuntimeRequest { url: String, authorization: String, + user_agent: String, + x_client_version: String, provider_api_format: String, request_body: Option, } @@ -2281,7 +2283,7 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru }), ); - let seen_execution_runtime = Arc::new(Mutex::new(None::)); + let seen_execution_runtime = Arc::new(Mutex::new(Vec::::new())); let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime); let execution_runtime = Router::new().route( "/v1/execute/sync", @@ -2294,26 +2296,30 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru .expect("body should read"), ) .expect("plan should parse"); - *seen_execution_runtime_inner + let request_body = plan.body.json_body.clone(); + seen_execution_runtime_inner .lock() - .expect("mutex should lock") = Some(SeenExecutionRuntimeRequest { - url: plan.url.clone(), - authorization: plan - .headers - .get("authorization") - .cloned() - .unwrap_or_default(), - provider_api_format: plan.provider_api_format.clone(), - request_body: plan.body.json_body.clone(), - }); - let result = aether_contracts::ExecutionResult { - request_id: plan.request_id, - candidate_id: None, - status_code: 200, - headers: BTreeMap::new(), - response_observation: None, - body: Some(aether_contracts::ResponseBody { - json_body: Some(json!({ + .expect("mutex should lock") + .push(SeenExecutionRuntimeRequest { + url: plan.url.clone(), + authorization: plan + .headers + .get("authorization") + .cloned() + .unwrap_or_default(), + user_agent: plan.headers.get("user-agent").cloned().unwrap_or_default(), + x_client_version: plan + .headers + .get("x-client-version") + .cloned() + .unwrap_or_default(), + provider_api_format: plan.provider_api_format.clone(), + request_body: request_body.clone(), + }); + let (status_code, json_body) = match plan.provider_api_format.as_str() { + "antigravity:fetch_available_models" => ( + 200, + json!({ "models": { "claude-sonnet-4": { "displayName": "Claude Sonnet 4", @@ -2326,7 +2332,47 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru "displayName": "Gemini 2.5 Pro" } } - })), + }), + ), + "antigravity:retrieve_user_quota_summary" + if request_body + .as_ref() + .and_then(|body| body.get("project")) + .is_some() => + { + (403, json!({"error": {"message": "project not accepted"}})) + } + "antigravity:retrieve_user_quota_summary" => ( + 200, + json!({ + "groups": [{ + "displayName": "Claude and GPT models", + "description": "Shared quota", + "buckets": [{ + "bucketId": "3p-5h", + "window": "5h", + "remainingFraction": 0.25, + "resetTime": "2026-05-05T05:00:00Z", + "displayName": "5 hour" + }, { + "bucketId": "3p-weekly", + "window": "weekly", + "remainingFraction": 0.8, + "resetTime": "2026-05-11T00:00:00Z" + }] + }] + }), + ), + unexpected => panic!("unexpected quota request format: {unexpected}"), + }; + let result = aether_contracts::ExecutionResult { + request_id: plan.request_id, + candidate_id: None, + status_code, + headers: BTreeMap::new(), + response_observation: None, + body: Some(aether_contracts::ResponseBody { + json_body: Some(json_body), body_bytes_b64: None, }), telemetry: None, @@ -2429,31 +2475,55 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru payload["results"][0]["quota_snapshot"]["windows"] .as_array() .map(Vec::len), - Some(1usize) + Some(3usize) ); assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0); - let seen_execution_runtime_request = seen_execution_runtime + let seen_execution_runtime_requests = seen_execution_runtime .lock() .expect("mutex should lock") - .clone() - .expect("execution runtime request should be captured"); + .clone(); + assert_eq!(seen_execution_runtime_requests.len(), 3); + let fetch_models_request = &seen_execution_runtime_requests[0]; assert_eq!( - seen_execution_runtime_request.url, + fetch_models_request.url, "https://daily-cloudcode-pa.googleapis.com/v1internal:fetchAvailableModels" ); + assert_eq!(fetch_models_request.authorization, "Bearer ya29.ant-token"); assert_eq!( - seen_execution_runtime_request.authorization, - "Bearer ya29.ant-token" + fetch_models_request.user_agent, + "vscode/1.X.X (Antigravity/4.3.0)" ); + assert_eq!(fetch_models_request.x_client_version, "4.3.0"); assert_eq!( - seen_execution_runtime_request.provider_api_format, + fetch_models_request.provider_api_format, "antigravity:fetch_available_models" ); assert_eq!( - seen_execution_runtime_request.request_body, + fetch_models_request.request_body, Some(json!({ "project": "project-ant-123" })) ); + let grouped_with_project = &seen_execution_runtime_requests[1]; + let grouped_without_project = &seen_execution_runtime_requests[2]; + assert_eq!( + grouped_with_project.url, + "https://daily-cloudcode-pa.googleapis.com/v1internal:retrieveUserQuotaSummary" + ); + assert_eq!( + grouped_with_project.provider_api_format, + "antigravity:retrieve_user_quota_summary" + ); + assert_eq!( + grouped_with_project.request_body, + Some(json!({"project": "project-ant-123"})) + ); + assert_eq!(grouped_without_project.url, grouped_with_project.url); + assert_eq!(grouped_without_project.request_body, Some(json!({}))); + assert!(seen_execution_runtime_requests.iter().all(|request| { + request.authorization == "Bearer ya29.ant-token" + && request.user_agent == "vscode/1.X.X (Antigravity/4.3.0)" + && request.x_client_version == "4.3.0" + })); let reloaded = provider_catalog_repository .list_keys_by_ids(&["key-antigravity-a".to_string()]) @@ -2481,6 +2551,26 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru .and_then(|value| value.get("used_percent")), Some(&json!(75.0)) ); + assert_eq!( + reloaded[0] + .upstream_metadata + .as_ref() + .and_then(|value| value.pointer("/antigravity/quota_groups/0/buckets/0/bucket_id")), + Some(&json!("3p-5h")) + ); + assert_eq!( + reloaded[0] + .upstream_metadata + .as_ref() + .and_then(|value| value.pointer("/antigravity/project_id")), + Some(&json!("project-ant-123")) + ); + assert!(reloaded[0] + .upstream_metadata + .as_ref() + .and_then(|value| value.pointer("/antigravity/quota_groups_updated_at")) + .and_then(serde_json::Value::as_u64) + .is_some()); assert_eq!( reloaded[0] .status_snapshot @@ -2505,7 +2595,19 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru .and_then(|value| value.get("windows")) .and_then(|value| value.as_array()) .map(Vec::len), - Some(1usize) + Some(3usize) + ); + assert_eq!( + reloaded[0] + .status_snapshot + .as_ref() + .and_then(|value| value.pointer("/quota/windows")) + .and_then(serde_json::Value::as_array) + .and_then(|windows| windows + .iter() + .find(|window| window["code"] == "group:0:3p-5h")) + .and_then(|window| window.get("remaining_ratio")), + Some(&json!(0.25)) ); gateway_handle.abort(); diff --git a/crates/aether-admin/src/provider/quota.rs b/crates/aether-admin/src/provider/quota.rs index e70af0420..7fde73523 100644 --- a/crates/aether-admin/src/provider/quota.rs +++ b/crates/aether-admin/src/provider/quota.rs @@ -286,6 +286,124 @@ pub fn parse_antigravity_usage_response( })) } +pub fn parse_antigravity_quota_summary_response( + value: &serde_json::Value, +) -> Option { + let groups = value.get("groups")?.as_array()?; + let mut parsed_groups = Vec::new(); + + for (group_index, group) in groups.iter().enumerate() { + let Some(group) = group.as_object() else { + continue; + }; + let group_id = coerce_json_string( + group + .get("groupId") + .or_else(|| group.get("group_id")) + .or_else(|| group.get("id")), + ); + let display_name = coerce_json_string( + group + .get("displayName") + .or_else(|| group.get("display_name")), + ) + .unwrap_or_else(|| format!("Quota group {}", group_index + 1)); + let description = coerce_json_string(group.get("description")); + let mut parsed_buckets = Vec::new(); + + for (bucket_index, bucket) in group + .get("buckets") + .and_then(serde_json::Value::as_array) + .into_iter() + .flatten() + .enumerate() + { + let Some(bucket) = bucket.as_object() else { + continue; + }; + let bucket_id = coerce_json_string( + bucket + .get("bucketId") + .or_else(|| bucket.get("bucket_id")) + .or_else(|| bucket.get("id")), + ) + .unwrap_or_else(|| format!("bucket-{}", bucket_index + 1)); + let window = coerce_json_string(bucket.get("window")); + let remaining_fraction = bucket + .get("remainingFraction") + .or_else(|| bucket.get("remaining_fraction")) + .and_then(coerce_json_f64) + .map(|value| value.clamp(0.0, 1.0)); + let reset_time = bucket + .get("resetTime") + .or_else(|| bucket.get("reset_time")) + .cloned() + .filter(|value| !value.is_null()); + let bucket_display_name = coerce_json_string( + bucket + .get("displayName") + .or_else(|| bucket.get("display_name")), + ); + let bucket_description = coerce_json_string(bucket.get("description")); + + if window.is_none() + && remaining_fraction.is_none() + && reset_time.is_none() + && bucket_display_name.is_none() + && bucket_description.is_none() + { + continue; + } + + let mut parsed_bucket = serde_json::Map::new(); + parsed_bucket.insert("bucket_id".to_string(), json!(bucket_id)); + if let Some(window) = window { + parsed_bucket.insert("window".to_string(), json!(window)); + } + if let Some(remaining_fraction) = remaining_fraction { + parsed_bucket.insert("remaining_fraction".to_string(), json!(remaining_fraction)); + parsed_bucket.insert( + "used_percent".to_string(), + json!((1.0 - remaining_fraction) * 100.0), + ); + parsed_bucket.insert( + "is_exhausted".to_string(), + json!(remaining_fraction <= 1e-6), + ); + } + if let Some(reset_time) = reset_time { + parsed_bucket.insert("reset_time".to_string(), reset_time); + } + if let Some(display_name) = bucket_display_name { + parsed_bucket.insert("display_name".to_string(), json!(display_name)); + } + if let Some(description) = bucket_description { + parsed_bucket.insert("description".to_string(), json!(description)); + } + parsed_buckets.push(serde_json::Value::Object(parsed_bucket)); + } + + if parsed_buckets.is_empty() { + continue; + } + let mut parsed_group = serde_json::Map::new(); + if let Some(group_id) = group_id { + parsed_group.insert("group_id".to_string(), json!(group_id)); + } + parsed_group.insert("display_name".to_string(), json!(display_name)); + if let Some(description) = description { + parsed_group.insert("description".to_string(), json!(description)); + } + parsed_group.insert( + "buckets".to_string(), + serde_json::Value::Array(parsed_buckets), + ); + parsed_groups.push(serde_json::Value::Object(parsed_group)); + } + + (!parsed_groups.is_empty()).then_some(serde_json::Value::Array(parsed_groups)) +} + pub fn parse_gemini_cli_retrieve_user_quota_response( value: &serde_json::Value, updated_at_unix_secs: u64, @@ -2042,6 +2160,13 @@ fn codex_additional_quota_windows( let Some(item_object) = item.as_object() else { continue; }; + if item_object + .get("limit_name") + .and_then(serde_json::Value::as_str) + .is_some_and(|name| name.trim() == CODEX_SPARK_LIMIT_NAME) + { + continue; + } let Some(limit_name) = ["limit_name", "metered_feature", "name", "id"] .iter() .find_map(|key| item_object.get(*key).and_then(serde_json::Value::as_str)) @@ -3944,11 +4069,11 @@ mod tests { codex_rate_limit_metadata_exhausted, codex_runtime_invalid_reason, codex_websocket_response_has_usage_limit_error, codex_websocket_usage_limit_reset_at, extract_execution_error_detail, merge_codex_quota_metadata_snapshot, - normalize_codex_reset_credit_consume_outcome, parse_antigravity_usage_response, - parse_chatgpt_web_conversation_init_response, parse_codex_backend_me_response, - parse_codex_usage_headers, parse_codex_websocket_rate_limits_response, - parse_codex_wham_reset_credits_detail_response, parse_codex_wham_usage_response, - parse_gemini_cli_retrieve_user_quota_response, + normalize_codex_reset_credit_consume_outcome, parse_antigravity_quota_summary_response, + parse_antigravity_usage_response, parse_chatgpt_web_conversation_init_response, + parse_codex_backend_me_response, parse_codex_usage_headers, + parse_codex_websocket_rate_limits_response, parse_codex_wham_reset_credits_detail_response, + parse_codex_wham_usage_response, parse_gemini_cli_retrieve_user_quota_response, parse_gemini_cli_v1internal_credits_response, parse_windsurf_model_configs_response, parse_windsurf_rate_limit_response, parse_windsurf_user_status_response, provider_auto_remove_quota_exhausted_keys, quota_refresh_success_invalid_state, @@ -6368,11 +6493,10 @@ mod tests { .get("additional_quota_windows") .and_then(serde_json::Value::as_array) .expect("all additional quota windows should be preserved"); - assert_eq!(additional.len(), 2); - assert!(additional.iter().all(|window| { - window.get("model") == Some(&json!("GPT-5.3-Codex-Spark")) - && window.get("scope") == Some(&json!("model")) - })); + assert!( + additional.is_empty(), + "the Spark windows are already normalized into the dedicated spark fields" + ); } #[test] @@ -6916,6 +7040,37 @@ mod tests { ); } + #[test] + fn parses_antigravity_grouped_weekly_and_five_hour_quota() { + let groups = parse_antigravity_quota_summary_response(&json!({ + "groups": [{ + "displayName": "Claude and GPT models", + "description": "Shared quota", + "buckets": [{ + "bucketId": "3p-5h", + "window": "5h", + "remainingFraction": 0.25, + "resetTime": "2026-05-05T05:00:00Z", + "displayName": "5 hour" + }, { + "bucketId": "3p-weekly", + "window": "weekly", + "remainingFraction": 0.8, + "resetTime": "2026-05-11T00:00:00Z" + }] + }] + })) + .expect("grouped Antigravity quota should parse"); + + assert_eq!(groups[0]["display_name"], json!("Claude and GPT models")); + assert_eq!(groups[0]["description"], json!("Shared quota")); + assert_eq!(groups[0]["buckets"][0]["bucket_id"], json!("3p-5h")); + assert_eq!(groups[0]["buckets"][0]["window"], json!("5h")); + assert_eq!(groups[0]["buckets"][0]["remaining_fraction"], json!(0.25)); + assert_eq!(groups[0]["buckets"][0]["used_percent"], json!(75.0)); + assert_eq!(groups[0]["buckets"][1]["bucket_id"], json!("3p-weekly")); + } + #[test] fn parses_gemini_cli_retrieve_user_quota_buckets() { let parsed = parse_gemini_cli_retrieve_user_quota_response( diff --git a/crates/aether-model-fetch/src/logic.rs b/crates/aether-model-fetch/src/logic.rs index b4c330257..e09f69057 100644 --- a/crates/aether-model-fetch/src/logic.rs +++ b/crates/aether-model-fetch/src/logic.rs @@ -624,6 +624,15 @@ pub fn merge_upstream_metadata(current: Option<&Value>, incoming: &Value) -> Val next_value.as_object_mut(), merged.get(namespace).and_then(Value::as_object), ) { + if namespace.eq_ignore_ascii_case("antigravity") { + for field in ["quota_groups", "quota_groups_updated_at"] { + if !next_namespace.contains_key(field) { + if let Some(value) = old_namespace.get(field) { + next_namespace.insert(field.to_string(), value.clone()); + } + } + } + } if let (Some(new_quota), Some(old_quota)) = ( next_namespace .get_mut("quota_by_model") @@ -1738,6 +1747,11 @@ mod tests { let merged = merge_upstream_metadata( Some(&json!({ "antigravity": { + "quota_groups": [{ + "display_name": "Claude and GPT models", + "buckets": [{"bucket_id": "3p-5h", "window": "5h"}] + }], + "quota_groups_updated_at": 1_777_000_000u64, "quota_by_model": { "gemini-2.5-pro": { "remaining_fraction": 0.3, @@ -1767,6 +1781,14 @@ mod tests { assert!(merged["antigravity"]["quota_by_model"] .get("stale-model") .is_none()); + assert_eq!( + merged["antigravity"]["quota_groups"][0]["buckets"][0]["bucket_id"], + "3p-5h" + ); + assert_eq!( + merged["antigravity"]["quota_groups_updated_at"], + json!(1_777_000_000u64) + ); } #[test] diff --git a/crates/aether-model-fetch/src/strategy.rs b/crates/aether-model-fetch/src/strategy.rs index 1412f7331..64951adb8 100644 --- a/crates/aether-model-fetch/src/strategy.rs +++ b/crates/aether-model-fetch/src/strategy.rs @@ -1037,7 +1037,9 @@ fn parse_antigravity_models_response(body: &Value) -> Result<(Vec, Option })); let quota_payload = build_antigravity_quota_payload(model_object.get("quotaInfo")); - quota_by_model.insert(model_id.to_string(), Value::Object(quota_payload)); + if !quota_payload.is_empty() { + quota_by_model.insert(model_id.to_string(), Value::Object(quota_payload)); + } } let upstream_metadata = (!quota_by_model.is_empty()).then(|| { @@ -1141,31 +1143,37 @@ fn infer_kiro_model_owner(model_id: &str) -> &'static str { } fn build_antigravity_quota_payload(quota_info: Option<&Value>) -> serde_json::Map { - let quota_info = quota_info.and_then(Value::as_object); + let Some(quota_info) = quota_info.and_then(Value::as_object) else { + return serde_json::Map::new(); + }; let reset_time = quota_info - .and_then(|value| value.get("resetTime")) + .get("resetTime") .and_then(Value::as_str) .map(str::trim) .filter(|value| !value.is_empty()) .map(ToOwned::to_owned); let remaining_fraction = quota_info - .and_then(|value| value.get("remainingFraction")) - .and_then(Value::as_f64); + .get("remainingFraction") + .and_then(|value| { + value.as_f64().or_else(|| { + value + .as_str() + .map(str::trim) + .filter(|value| !value.is_empty()) + .and_then(|value| value.parse::().ok()) + }) + }) + .filter(|value| value.is_finite()) + .map(|value| value.clamp(0.0, 1.0)); let mut payload = serde_json::Map::new(); - match remaining_fraction { - Some(remaining_fraction) => { - let used_percent = ((1.0 - remaining_fraction) * 100.0).clamp(0.0, 100.0); - payload.insert( - "remaining_fraction".to_string(), - Value::from(remaining_fraction), - ); - payload.insert("used_percent".to_string(), Value::from(used_percent)); - } - None => { - payload.insert("remaining_fraction".to_string(), Value::from(0.0)); - payload.insert("used_percent".to_string(), Value::from(100.0)); - } + if let Some(remaining_fraction) = remaining_fraction { + let used_percent = (1.0 - remaining_fraction) * 100.0; + payload.insert( + "remaining_fraction".to_string(), + Value::from(remaining_fraction), + ); + payload.insert("used_percent".to_string(), Value::from(used_percent)); } if let Some(reset_time) = reset_time { payload.insert("reset_time".to_string(), Value::String(reset_time)); @@ -1618,8 +1626,8 @@ mod tests { use super::{ build_vertex_google_list_url, build_vertex_service_account_list_url, - parse_codex_models_response_for_request, select_model_fetch_strategy, ModelFetchStrategy, - ModelFetchStrategyKind, + parse_antigravity_models_response, parse_codex_models_response_for_request, + select_model_fetch_strategy, ModelFetchStrategy, ModelFetchStrategyKind, }; use crate::transport::ModelFetchTransportRuntime; use crate::{fetch_models_from_transports, fetch_models_from_transports_for_client_version}; @@ -2700,6 +2708,50 @@ mod tests { ); } + #[test] + fn antigravity_models_without_explicit_quota_are_not_marked_exhausted() { + let (models, metadata) = parse_antigravity_models_response(&json!({ + "models": { + "gemini-3.7-flash-tiered": { + "displayName": "Gemini 3.7 Flash" + }, + "gemini-3.7-flash-high": { + "displayName": "Gemini 3.7 Flash High", + "quotaInfo": { + "remainingFraction": "0.75", + "resetTime": "2030-01-01T00:00:00Z" + } + }, + "gemini-3.7-flash-low": { + "displayName": "Gemini 3.7 Flash Low", + "quotaInfo": { + "remainingFraction": 0.0 + } + } + } + })) + .expect("Antigravity models should parse"); + + assert_eq!(models.len(), 3); + let metadata = metadata.expect("explicit quota should produce metadata"); + let antigravity = &metadata["antigravity"]; + assert!(antigravity["quota_by_model"] + .get("gemini-3.7-flash-tiered") + .is_none()); + assert_eq!( + antigravity["quota_by_model"]["gemini-3.7-flash-high"]["remaining_fraction"], + json!(0.75) + ); + assert_eq!( + antigravity["quota_by_model"]["gemini-3.7-flash-high"]["used_percent"], + json!(25.0) + ); + assert_eq!( + antigravity["quota_by_model"]["gemini-3.7-flash-low"]["used_percent"], + json!(100.0) + ); + } + #[tokio::test] async fn kiro_transport_fetches_list_available_models() { let executed_urls = Arc::new(Mutex::new(Vec::new())); diff --git a/crates/aether-model-fetch/src/transport.rs b/crates/aether-model-fetch/src/transport.rs index 15fa42366..1fffba4af 100644 --- a/crates/aether-model-fetch/src/transport.rs +++ b/crates/aether-model-fetch/src/transport.rs @@ -1185,7 +1185,7 @@ mod tests { ); assert_eq!( plan.headers.get("x-client-version").map(String::as_str), - Some("1.2.3") + Some("4.3.0") ); assert_eq!( plan.headers.get("x-vscode-sessionid").map(String::as_str), diff --git a/crates/aether-provider/pool/src/lib.rs b/crates/aether-provider/pool/src/lib.rs index f8c813614..7b29abc51 100644 --- a/crates/aether-provider/pool/src/lib.rs +++ b/crates/aether-provider/pool/src/lib.rs @@ -15,10 +15,11 @@ pub use presets::{ }; pub use provider::{ProviderPoolAdapter, ProviderPoolMemberInput}; pub use providers::{ - build_antigravity_pool_quota_request, build_chatgpt_web_pool_quota_request, - build_codex_pool_quota_request, build_codex_pool_reset_credit_consume_request, - build_codex_pool_reset_credits_request, build_gemini_cli_pool_quota_request, - build_kiro_pool_quota_request, build_windsurf_pool_model_configs_request, + build_antigravity_pool_quota_request, build_antigravity_pool_quota_summary_request, + build_chatgpt_web_pool_quota_request, build_codex_pool_quota_request, + build_codex_pool_reset_credit_consume_request, build_codex_pool_reset_credits_request, + build_gemini_cli_pool_quota_request, build_kiro_pool_quota_request, + build_windsurf_pool_model_configs_request, build_windsurf_pool_model_configs_request_with_base_url, build_windsurf_pool_quota_request, build_windsurf_pool_quota_request_with_base_url, build_windsurf_pool_rate_limit_request, build_windsurf_pool_rate_limit_request_with_base_url, enrich_chatgpt_web_quota_metadata, @@ -27,11 +28,12 @@ pub use providers::{ AntigravityProviderPoolAdapter, ChatGptWebProviderPoolAdapter, CodexProviderPoolAdapter, DefaultProviderPoolAdapter, GeminiCliProviderPoolAdapter, GrokProviderPoolAdapter, KiroPoolQuotaAuthInput, KiroProviderPoolAdapter, UnsupportedQuotaProviderPoolAdapter, - ANTIGRAVITY_FETCH_AVAILABLE_MODELS_PATH, CHATGPT_WEB_CONVERSATION_INIT_PATH, - CHATGPT_WEB_DEFAULT_BASE_URL, CODEX_WHAM_RESET_CREDITS_CONSUME_URL, - CODEX_WHAM_RESET_CREDITS_URL, CODEX_WHAM_USAGE_URL, GEMINI_CLI_RETRIEVE_USER_QUOTA_PATH, - GEMINI_CLI_USER_AGENT, KIRO_USAGE_LIMITS_PATH, KIRO_USAGE_SDK_VERSION, - WINDSURF_MODEL_CONFIGS_PATH, WINDSURF_RATE_LIMIT_PATH, WINDSURF_USER_STATUS_PATH, + ANTIGRAVITY_FETCH_AVAILABLE_MODELS_PATH, ANTIGRAVITY_RETRIEVE_USER_QUOTA_SUMMARY_PATH, + CHATGPT_WEB_CONVERSATION_INIT_PATH, CHATGPT_WEB_DEFAULT_BASE_URL, + CODEX_WHAM_RESET_CREDITS_CONSUME_URL, CODEX_WHAM_RESET_CREDITS_URL, CODEX_WHAM_USAGE_URL, + GEMINI_CLI_RETRIEVE_USER_QUOTA_PATH, GEMINI_CLI_USER_AGENT, KIRO_USAGE_LIMITS_PATH, + KIRO_USAGE_SDK_VERSION, WINDSURF_MODEL_CONFIGS_PATH, WINDSURF_RATE_LIMIT_PATH, + WINDSURF_USER_STATUS_PATH, }; pub use quota::{ provider_pool_key_account_quota_exhausted, provider_pool_key_model_quota_exhausted, @@ -883,6 +885,66 @@ mod tests { assert!(!available.quota_exhausted); } + #[test] + fn antigravity_tiered_model_quota_wins_over_display_only_group_windows() { + let service = ProviderPoolService::with_builtin_adapters(); + let mut key = sample_key(None); + key.status_snapshot = Some(json!({ + "quota": { + "version": 2, + "provider_type": "antigravity", + "exhausted": true, + "windows": [{ + "code": "model:gemini-3.7-flash-tiered", + "scope": "model", + "model": "gemini-3.7-flash-tiered", + "remaining_ratio": 0.906, + "used_ratio": 0.094, + "is_exhausted": false + }, { + "code": "group:0:3p-5h", + "scope": "quota_group", + "quota_group": "group:0", + "bucket_id": "3p-5h", + "used_ratio": 1.0, + "is_exhausted": true + }] + } + })); + + let signals = + service.member_signals("antigravity", &key, None, Some("gemini-3.7-flash-tiered")); + + assert!(!signals.quota_exhausted); + assert!(!signals.quota_hard_blocked); + } + + #[test] + fn antigravity_tiered_variant_does_not_match_another_model_family() { + let service = ProviderPoolService::with_builtin_adapters(); + let mut key = sample_key(None); + key.status_snapshot = Some(json!({ + "quota": { + "version": 2, + "provider_type": "antigravity", + "exhausted": false, + "windows": [{ + "code": "model:gemini-3.7-pro-tiered", + "scope": "model", + "model": "gemini-3.7-pro-tiered", + "used_ratio": 1.0, + "is_exhausted": true + }] + } + })); + + let signals = + service.member_signals("antigravity", &key, None, Some("gemini-3.7-flash-tiered")); + + assert!(!signals.quota_exhausted); + assert!(!signals.quota_hard_blocked); + } + #[test] fn codex_standard_and_spark_quota_families_are_independent() { let service = ProviderPoolService::with_builtin_adapters(); diff --git a/crates/aether-provider/pool/src/providers/antigravity.rs b/crates/aether-provider/pool/src/providers/antigravity.rs index 24d9bb797..9b87ae7c6 100644 --- a/crates/aether-provider/pool/src/providers/antigravity.rs +++ b/crates/aether-provider/pool/src/providers/antigravity.rs @@ -12,6 +12,8 @@ use crate::quota::provider_pool_model_quota_exhausted; use crate::quota_refresh::ProviderPoolQuotaRequestSpec; pub const ANTIGRAVITY_FETCH_AVAILABLE_MODELS_PATH: &str = "/v1internal:fetchAvailableModels"; +pub const ANTIGRAVITY_RETRIEVE_USER_QUOTA_SUMMARY_PATH: &str = + "/v1internal:retrieveUserQuotaSummary"; #[derive(Debug, Clone, Default)] pub struct AntigravityProviderPoolAdapter; @@ -92,3 +94,85 @@ pub fn build_antigravity_pool_quota_request( accept_invalid_certs: false, } } + +pub fn build_antigravity_pool_quota_summary_request( + key_id: &str, + endpoint_base_url: &str, + authorization: (String, String), + project_id: Option<&str>, + mut identity_headers: BTreeMap, +) -> ProviderPoolQuotaRequestSpec { + let mut headers = std::mem::take(&mut identity_headers); + headers.insert("authorization".to_string(), authorization.1); + headers.insert("content-type".to_string(), "application/json".to_string()); + headers.insert("accept".to_string(), "application/json".to_string()); + headers + .entry("user-agent".to_string()) + .or_insert_with(|| "antigravity".to_string()); + + let json_body = project_id + .map(str::trim) + .filter(|project_id| !project_id.is_empty()) + .map_or_else(|| json!({}), |project_id| json!({ "project": project_id })); + + ProviderPoolQuotaRequestSpec { + request_id: format!("antigravity-quota-summary:{key_id}"), + provider_name: "antigravity".to_string(), + quota_kind: "antigravity".to_string(), + method: "POST".to_string(), + url: format!( + "{}{}", + endpoint_base_url.trim_end_matches('/'), + ANTIGRAVITY_RETRIEVE_USER_QUOTA_SUMMARY_PATH + ), + headers, + content_type: Some("application/json".to_string()), + json_body: Some(json_body), + client_api_format: "gemini:generate_content".to_string(), + provider_api_format: "antigravity:retrieve_user_quota_summary".to_string(), + model_name: Some("retrieveUserQuotaSummary".to_string()), + accept_invalid_certs: false, + } +} + +#[cfg(test)] +mod tests { + use std::collections::BTreeMap; + + use serde_json::json; + + use super::{ + build_antigravity_pool_quota_summary_request, ANTIGRAVITY_RETRIEVE_USER_QUOTA_SUMMARY_PATH, + }; + + #[test] + fn grouped_quota_request_can_retry_without_project_on_the_same_endpoint() { + let with_project = build_antigravity_pool_quota_summary_request( + "key-1", + "https://daily-cloudcode-pa.googleapis.com/", + ("authorization".to_string(), "Bearer token".to_string()), + Some("project-1"), + BTreeMap::new(), + ); + let without_project = build_antigravity_pool_quota_summary_request( + "key-1", + "https://daily-cloudcode-pa.googleapis.com/", + ("authorization".to_string(), "Bearer token".to_string()), + None, + BTreeMap::new(), + ); + + assert_eq!( + with_project.url, + format!( + "https://daily-cloudcode-pa.googleapis.com{ANTIGRAVITY_RETRIEVE_USER_QUOTA_SUMMARY_PATH}" + ) + ); + assert_eq!( + with_project.json_body, + Some(json!({"project": "project-1"})) + ); + assert_eq!(without_project.url, with_project.url); + assert_eq!(without_project.json_body, Some(json!({}))); + } +} diff --git a/crates/aether-provider/pool/src/providers/mod.rs b/crates/aether-provider/pool/src/providers/mod.rs index 0239dde79..72734c81e 100644 --- a/crates/aether-provider/pool/src/providers/mod.rs +++ b/crates/aether-provider/pool/src/providers/mod.rs @@ -10,7 +10,8 @@ pub mod windsurf; pub use antigravity::AntigravityProviderPoolAdapter; pub use antigravity::{ - build_antigravity_pool_quota_request, ANTIGRAVITY_FETCH_AVAILABLE_MODELS_PATH, + build_antigravity_pool_quota_request, build_antigravity_pool_quota_summary_request, + ANTIGRAVITY_FETCH_AVAILABLE_MODELS_PATH, ANTIGRAVITY_RETRIEVE_USER_QUOTA_SUMMARY_PATH, }; pub use chatgpt_web::ChatGptWebProviderPoolAdapter; pub use chatgpt_web::{ diff --git a/crates/aether-provider/pool/src/quota.rs b/crates/aether-provider/pool/src/quota.rs index f19200fc8..9d18d506c 100644 --- a/crates/aether-provider/pool/src/quota.rs +++ b/crates/aether-provider/pool/src/quota.rs @@ -398,6 +398,7 @@ fn provider_pool_is_specific_model_token(token: &str) -> bool { "team", "think", "thinking", + "tiered", "vendor", ] .contains(&token) diff --git a/crates/aether-provider/transport/src/antigravity/auth.rs b/crates/aether-provider/transport/src/antigravity/auth.rs index 52925427d..c8043a7b5 100644 --- a/crates/aether-provider/transport/src/antigravity/auth.rs +++ b/crates/aether-provider/transport/src/antigravity/auth.rs @@ -5,8 +5,8 @@ use serde_json::Value; use super::super::snapshot::GatewayProviderTransportSnapshot; pub const ANTIGRAVITY_PROVIDER_TYPE: &str = "antigravity"; -pub const ANTIGRAVITY_REQUEST_USER_AGENT: &str = - "antigravity/cli/1.0.16 (aidev_client; os_type=linux; arch=arm64; auth_method=consumer)"; +pub const ANTIGRAVITY_CLIENT_VERSION: &str = "4.3.0"; +pub const ANTIGRAVITY_REQUEST_USER_AGENT: &str = "vscode/1.X.X (Antigravity/4.3.0)"; const ANTIGRAVITY_CLIENT_NAME: &str = "antigravity"; const ANTIGRAVITY_GOOG_API_CLIENT: &str = "gl-node/18.18.2 fire/0.8.6 grpc/1.10.x"; @@ -109,7 +109,7 @@ pub fn build_antigravity_static_identity_headers( } pub fn build_antigravity_static_client_headers( - client_version: Option<&str>, + _client_version: Option<&str>, session_id: Option<&str>, ) -> BTreeMap { let mut headers = BTreeMap::from([ @@ -125,14 +125,12 @@ pub fn build_antigravity_static_client_headers( String::from("user-agent"), String::from(ANTIGRAVITY_REQUEST_USER_AGENT), ), + ( + String::from("x-client-version"), + String::from(ANTIGRAVITY_CLIENT_VERSION), + ), ]); - if let Some(client_version) = client_version - .map(str::trim) - .filter(|value| !value.is_empty()) - { - headers.insert(String::from("x-client-version"), client_version.to_string()); - } if let Some(session_id) = session_id.map(str::trim).filter(|value| !value.is_empty()) { headers.insert(String::from("x-vscode-sessionid"), session_id.to_string()); } @@ -273,7 +271,8 @@ mod tests { use super::{ build_antigravity_static_client_headers, resolve_local_antigravity_request_auth, - AntigravityRequestAuth, AntigravityRequestAuthSupport, ANTIGRAVITY_REQUEST_USER_AGENT, + AntigravityRequestAuth, AntigravityRequestAuthSupport, ANTIGRAVITY_CLIENT_VERSION, + ANTIGRAVITY_REQUEST_USER_AGENT, }; use crate::snapshot::{ GatewayProviderTransportEndpoint, GatewayProviderTransportKey, @@ -383,7 +382,7 @@ mod tests { } #[test] - fn static_client_headers_use_native_antigravity_cli_user_agent() { + fn static_client_headers_pin_the_known_good_antigravity_identity() { let headers = build_antigravity_static_client_headers(Some("1.0.16"), Some("session-abc")); assert_eq!( @@ -396,7 +395,7 @@ mod tests { ); assert_eq!( headers.get("x-client-version").map(String::as_str), - Some("1.0.16") + Some(ANTIGRAVITY_CLIENT_VERSION) ); assert_eq!( headers.get("x-vscode-sessionid").map(String::as_str), diff --git a/frontend/src/features/providers/components/AntigravityQuotaDialog.vue b/frontend/src/features/providers/components/AntigravityQuotaDialog.vue index 2fec683f7..0a60ba705 100644 --- a/frontend/src/features/providers/components/AntigravityQuotaDialog.vue +++ b/frontend/src/features/providers/components/AntigravityQuotaDialog.vue @@ -45,11 +45,11 @@
@@ -238,6 +238,48 @@ function buildItemsFromQuotaSnapshot(quota: QuotaStatusSnapshot | null | undefin return dedupeAntigravityQuotaItemsByLabel(items) } +function buildGroupedItemsFromQuotaSnapshot( + quota: QuotaStatusSnapshot | null | undefined, +): QuotaItem[] { + if (!quota) return [] + + const providerType = String(quota.provider_type || '').trim().toLowerCase() + if (providerType && providerType !== 'antigravity') return [] + + const windows = Array.isArray(quota.windows) + ? quota.windows.filter(window => String(window?.scope || '').trim().toLowerCase() === 'quota_group') + : [] + + return windows + .map((window) => { + const code = String(window.code || '').trim() + const label = String(window.label || code).trim() + if (!code || !label) return null + + const usedPercent = + typeof window.used_ratio === 'number' + ? Math.max(Math.min(window.used_ratio * 100, 100), 0) + : typeof window.remaining_ratio === 'number' + ? Math.max(Math.min((1 - window.remaining_ratio) * 100, 100), 0) + : null + if (usedPercent == null) return null + + const remainingPercent = + typeof window.remaining_ratio === 'number' + ? Math.max(Math.min(window.remaining_ratio * 100, 100), 0) + : Math.max(100 - usedPercent, 0) + + return { + model: code, + label, + usedPercent, + remainingPercent, + resetSeconds: getQuotaWindowLiveResetSeconds(quota, window), + } satisfies QuotaItem + }) + .filter((item): item is QuotaItem => item !== null) +} + const rawItems = computed(() => { const snapshotItems = buildItemsFromQuotaSnapshot(props.quotaSnapshot) if (snapshotItems.length > 0) return snapshotItems @@ -291,6 +333,8 @@ const rawItems = computed(() => { }) const items = computed(() => summarizeAntigravityQuotaItems(rawItems.value)) +const groupedItems = computed(() => buildGroupedItemsFromQuotaSnapshot(props.quotaSnapshot)) +const displayItems = computed(() => [...groupedItems.value, ...items.value]) async function handleTestModel(modelName: string) { if (!props.providerId || testingModel.value) return diff --git a/frontend/src/features/providers/components/__tests__/AntigravityQuotaDialog.spec.ts b/frontend/src/features/providers/components/__tests__/AntigravityQuotaDialog.spec.ts index 8a7bda80b..ded9c2119 100644 --- a/frontend/src/features/providers/components/__tests__/AntigravityQuotaDialog.spec.ts +++ b/frontend/src/features/providers/components/__tests__/AntigravityQuotaDialog.spec.ts @@ -2,7 +2,7 @@ import { describe, expect, it, vi } from 'vitest' import { createApp, defineComponent, h } from 'vue' import AntigravityQuotaDialog from '@/features/providers/components/AntigravityQuotaDialog.vue' -import type { UpstreamMetadata } from '@/api/endpoints/types' +import type { QuotaStatusSnapshot, UpstreamMetadata } from '@/api/endpoints/types' vi.mock('@/components/ui', async () => { const { defineComponent, h } = await import('vue') @@ -77,7 +77,7 @@ vi.mock('@/utils/errorParser', () => ({ parseApiError: (value: unknown) => String(value), })) -function mount(metadata: UpstreamMetadata) { +function mount(metadata: UpstreamMetadata, quotaSnapshot?: QuotaStatusSnapshot) { const root = document.createElement('div') document.body.appendChild(root) @@ -86,6 +86,7 @@ function mount(metadata: UpstreamMetadata) { return () => h(AntigravityQuotaDialog, { open: true, metadata, + quotaSnapshot, keyName: 'Key-1', }) }, @@ -102,6 +103,35 @@ function mount(metadata: UpstreamMetadata) { } describe('AntigravityQuotaDialog', () => { + it('renders grouped five-hour and weekly quota windows', () => { + const { root, unmount } = mount({}, { + code: 'ok', + provider_type: 'antigravity', + exhausted: false, + windows: [{ + code: 'group:0:3p-5h', + label: 'Claude and GPT models · 5 hour', + scope: 'quota_group', + used_ratio: 0.75, + remaining_ratio: 0.25, + }, { + code: 'group:0:3p-weekly', + label: 'Claude and GPT models · weekly', + scope: 'quota_group', + used_ratio: 0.2, + remaining_ratio: 0.8, + }], + }) + const text = root.textContent || '' + + expect(text).toContain('Claude and GPT models · 5 hour') + expect(text).toContain('25.0%') + expect(text).toContain('Claude and GPT models · weekly') + expect(text).toContain('80.0%') + + unmount() + }) + it('hides quota buckets outside the shared pool summary families', () => { const rawIdentifier = 'RateLimitResetCredit_05cbb6eeeb9c81918e011d8300f9ebfb' const { root, unmount } = mount({ From 76fb8905c9df5f969c6c99237e6769fd5123c920 Mon Sep 17 00:00:00 2001 From: ZheFox <77232781+zhefox@users.noreply.github.com> Date: Thu, 3 Sep 2026 11:04:18 +0800 Subject: [PATCH 06/13] test(gateway): align antigravity response message id --- .../src/tests/ai_execute/finalize_local_cli/cross_format.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/apps/aether-gateway/src/tests/ai_execute/finalize_local_cli/cross_format.rs b/apps/aether-gateway/src/tests/ai_execute/finalize_local_cli/cross_format.rs index 0f1a1338c..d5d79c9a2 100644 --- a/apps/aether-gateway/src/tests/ai_execute/finalize_local_cli/cross_format.rs +++ b/apps/aether-gateway/src/tests/ai_execute/finalize_local_cli/cross_format.rs @@ -6,6 +6,7 @@ use super::{ EXECUTION_PATH_HEADER, TRACE_ID_HEADER, }; use crate::data::GatewayDataState; +use aether_ai_formats::openai_responses_message_item_id; use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY}; use aether_data::repository::auth::{ InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot, @@ -1465,7 +1466,7 @@ async fn gateway_executes_openai_responses_antigravity_cross_format_upstream_str "output_text": "Hello Antigravity", "output": [{ "type": "message", - "id": "resp-local-stream_msg", + "id": openai_responses_message_item_id("resp-local-stream", 0), "role": "assistant", "status": "completed", "content": [{ From 40a5e1470d1f095b3e90c76d879b3d3191766691 Mon Sep 17 00:00:00 2001 From: ZheFox <77232781+zhefox@users.noreply.github.com> Date: Thu, 3 Sep 2026 13:16:40 +0800 Subject: [PATCH 07/13] fix(gateway): persist quota refresh from strong catalog reads --- .../src/data/state/testing/mod.rs | 15 ++ .../admin/provider/oauth/quota/shared.rs | 146 +++++++++++++++++- .../tests/control/admin/endpoints/quota.rs | 4 +- crates/aether-admin/src/provider/quota.rs | 9 +- 4 files changed, 163 insertions(+), 11 deletions(-) diff --git a/apps/aether-gateway/src/data/state/testing/mod.rs b/apps/aether-gateway/src/data/state/testing/mod.rs index 486d6f435..230c77a3d 100644 --- a/apps/aether-gateway/src/data/state/testing/mod.rs +++ b/apps/aether-gateway/src/data/state/testing/mod.rs @@ -213,6 +213,21 @@ impl GatewayDataState { self } + #[cfg(test)] + pub(crate) fn with_cached_provider_catalog_reader_for_tests( + mut self, + repository: Arc, + ) -> Self + where + T: ProviderCatalogReadRepository + 'static, + { + let inner: Arc = repository; + self.provider_catalog_reader = Some(Arc::new( + super::provider_catalog_cache::CachedProviderCatalogReadRepository::new(inner), + )); + self + } + #[cfg(test)] pub(crate) fn with_request_candidate_reader( mut self, diff --git a/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/shared.rs b/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/shared.rs index 6c5e55937..55f3d0837 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/shared.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/shared.rs @@ -240,7 +240,16 @@ fn merge_upstream_metadata( .unwrap_or_default(); if let Some(update_object) = updates.as_object() { for (key, value) in update_object { - merged.insert(key.clone(), value.clone()); + let mut next = value.clone(); + if let (Some(current_namespace), Some(next_namespace)) = ( + merged.get(key).and_then(serde_json::Value::as_object), + next.as_object_mut(), + ) { + let mut combined = current_namespace.clone(); + combined.extend(next_namespace.clone()); + next = serde_json::Value::Object(combined); + } + merged.insert(key.clone(), next); } } serde_json::Value::Object(merged) @@ -1306,7 +1315,8 @@ where F: std::future::Future, { let Some(mut latest_key) = state - .read_provider_catalog_keys_by_ids(&[key_id.to_string()]) + .app() + .list_provider_catalog_keys_by_ids_strong(&[key_id.to_string()]) .await? .into_iter() .next() @@ -1350,9 +1360,18 @@ where let metadata_updates = metadata_update .and_then(serde_json::Value::as_object) .map(|updates| { + let merged = latest_key + .upstream_metadata + .as_ref() + .and_then(serde_json::Value::as_object); updates - .iter() - .map(|(namespace, value)| (namespace.clone(), value.clone())) + .keys() + .filter_map(|namespace| { + merged + .and_then(|metadata| metadata.get(namespace)) + .cloned() + .map(|value| (namespace.clone(), value)) + }) .collect::>() }) .unwrap_or_default(); @@ -1384,7 +1403,7 @@ where } else { serde_json::json!({}) }; - let mut expected = observed_upstream_metadata + let expected = observed_upstream_metadata .as_ref() .and_then(serde_json::Value::as_object) .and_then(|metadata| metadata.get(namespace)) @@ -2807,4 +2826,121 @@ mod tests { json!({"remaining":4}) ); } + + #[tokio::test] + async fn quota_refresh_strong_read_bypasses_stale_provider_catalog_cache() { + let mut key = StoredProviderCatalogKey::new( + "key-antigravity-stale-cache".to_string(), + "provider-antigravity-stale-cache".to_string(), + "Antigravity stale cache".to_string(), + "oauth".to_string(), + None, + true, + ) + .expect("key should build"); + key.upstream_metadata = Some(json!({ + "antigravity": { + "project_id": "project-1", + "quota_by_model": { + "gemini-3.7-flash-tiered": {"remaining_fraction": 0.9} + } + } + })); + + let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( + vec![], + vec![], + vec![key], + )); + let data = + GatewayDataState::with_provider_catalog_repository_for_tests(Arc::clone(&repository)) + .with_cached_provider_catalog_reader_for_tests(Arc::clone(&repository)); + let app = AppState::new() + .expect("app should build") + .with_data_state_for_tests(data); + let admin_state = AdminAppState::new(&app); + let key_ids = ["key-antigravity-stale-cache".to_string()]; + + let cached = app + .read_provider_catalog_keys_by_ids(&key_ids) + .await + .expect("initial cached read should succeed"); + assert_eq!( + cached[0].upstream_metadata.as_ref().unwrap()["antigravity"]["quota_by_model"] + ["gemini-3.7-flash-tiered"]["remaining_fraction"], + json!(0.9) + ); + + let current_namespace = json!({ + "project_id": "project-1", + "model_fetch_revision": 2, + "quota_by_model": { + "gemini-3.7-flash-tiered": {"remaining_fraction": 0.7} + } + }); + assert!(repository + .upsert_key_upstream_metadata_namespace( + "key-antigravity-stale-cache", + "antigravity", + ¤t_namespace, + None, + ) + .await + .expect("out-of-band metadata update should succeed")); + let still_cached = app + .read_provider_catalog_keys_by_ids(&key_ids) + .await + .expect("stale cached read should succeed"); + assert_eq!( + still_cached[0].upstream_metadata.as_ref().unwrap()["antigravity"]["quota_by_model"] + ["gemini-3.7-flash-tiered"]["remaining_fraction"], + json!(0.9), + "regression setup must keep the ordinary read stale" + ); + + let metadata_update = json!({ + "antigravity": { + "project_id": "project-1", + "quota_by_model": { + "gemini-3.7-flash-tiered": {"remaining_fraction": 0.6} + }, + "quota_groups": [{ + "display_name": "Gemini models", + "buckets": [{"bucket_id": "gemini-weekly", "window": "weekly"}] + }] + } + }); + assert!(persist_provider_quota_refresh_state( + &admin_state, + "key-antigravity-stale-cache", + Some(&metadata_update), + None, + None, + None, + ) + .await + .expect("quota refresh persistence should not error")); + + let stored = repository + .list_keys_by_ids(&key_ids) + .await + .expect("key should reload") + .pop() + .expect("key should exist"); + assert_eq!( + stored.upstream_metadata.as_ref().unwrap()["antigravity"]["quota_groups"][0]["buckets"] + [0]["bucket_id"], + json!("gemini-weekly") + ); + assert_eq!( + stored.upstream_metadata.as_ref().unwrap()["antigravity"]["model_fetch_revision"], + json!(2), + "quota refresh must preserve fields written by another Antigravity metadata producer" + ); + assert_eq!( + stored.upstream_metadata.as_ref().unwrap()["antigravity"]["quota_by_model"] + ["gemini-3.7-flash-tiered"]["remaining_fraction"], + json!(0.6) + ); + } } diff --git a/apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs b/apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs index c4d52ce31..444d550f2 100644 --- a/apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs +++ b/apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs @@ -2536,7 +2536,7 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru .upstream_metadata .as_ref() .and_then(|value| value.get("antigravity")) - .and_then(|value| value.get("models")) + .and_then(|value| value.get("quota_by_model")) .and_then(|value| value.get("claude-sonnet-4")) .and_then(|value| value.get("remaining_fraction")), Some(&json!(0.25)) @@ -2546,7 +2546,7 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru .upstream_metadata .as_ref() .and_then(|value| value.get("antigravity")) - .and_then(|value| value.get("models")) + .and_then(|value| value.get("quota_by_model")) .and_then(|value| value.get("claude-sonnet-4")) .and_then(|value| value.get("used_percent")), Some(&json!(75.0)) diff --git a/crates/aether-admin/src/provider/quota.rs b/crates/aether-admin/src/provider/quota.rs index 7fde73523..349f26855 100644 --- a/crates/aether-admin/src/provider/quota.rs +++ b/crates/aether-admin/src/provider/quota.rs @@ -282,7 +282,7 @@ pub fn parse_antigravity_usage_response( "is_forbidden": false, "forbidden_reason": serde_json::Value::Null, "forbidden_at": serde_json::Value::Null, - "models": quota_by_model, + "quota_by_model": quota_by_model, })) } @@ -7026,16 +7026,17 @@ mod tests { .expect("antigravity quota should parse"); assert_eq!( - parsed["models"]["RateLimitResetCredit_05cbb6eeeb9c81918e011d8300f9ebfb"] + parsed["quota_by_model"]["RateLimitResetCredit_05cbb6eeeb9c81918e011d8300f9ebfb"] ["display_name"], json!("Key-1") ); assert_eq!( - parsed["models"]["RateLimitResetCredit_05cbb6eeeb9c81918e011d8300f9ebfb"]["reset_time"], + parsed["quota_by_model"]["RateLimitResetCredit_05cbb6eeeb9c81918e011d8300f9ebfb"] + ["reset_time"], json!("2030-01-01T00:00:00Z") ); assert_eq!( - parsed["models"]["gemini-3-pro-preview"]["display_name"], + parsed["quota_by_model"]["gemini-3-pro-preview"]["display_name"], json!("Gemini 3 Pro Preview") ); } From d672ba20687b641b80c994c512155a7397eaa16c Mon Sep 17 00:00:00 2001 From: ZheFox <77232781+zhefox@users.noreply.github.com> Date: Thu, 3 Sep 2026 13:45:57 +0800 Subject: [PATCH 08/13] test(gateway): seed routing strategy for antigravity flows --- .../src/data/state/testing/mod.rs | 29 +++++++++++++++++++ .../src/tests/ai_execute/finalize_local.rs | 3 +- .../finalize_local_cli/cross_format.rs | 3 +- .../finalize_local_provider/gemini.rs | 3 +- .../stream_provider_gemini/local_cli.rs | 1 + .../src/tests/ai_execute/sync/gemini/cli.rs | 3 +- 6 files changed, 38 insertions(+), 4 deletions(-) diff --git a/apps/aether-gateway/src/data/state/testing/mod.rs b/apps/aether-gateway/src/data/state/testing/mod.rs index 230c77a3d..0babed923 100644 --- a/apps/aether-gateway/src/data/state/testing/mod.rs +++ b/apps/aether-gateway/src/data/state/testing/mod.rs @@ -1,10 +1,15 @@ use std::collections::BTreeMap; use std::sync::{Arc, RwLock}; +use aether_data::repository::routing_profiles::InMemoryRoutingGroupRepository; use aether_data_contracts::repository::candidates::RequestCandidateRepository; use aether_data_contracts::repository::pool_scores::PoolMemberScoreRepository; use aether_data_contracts::repository::quota::ProviderQuotaRepository; +use aether_data_contracts::repository::routing_profiles::{ + StoredRoutingGroup, StoredRoutingGroupBinding, StoredRoutingGroupVersion, +}; use aether_data_contracts::repository::usage::UsageRepository; +use aether_routing_core::RoutingGroupConfig; use super::{ AnnouncementReadRepository, AnnouncementWriteRepository, AuthApiKeyReadRepository, @@ -892,6 +897,30 @@ impl GatewayDataState { self } + #[cfg(test)] + pub(crate) fn with_system_default_routing_group_for_tests(self) -> Self { + let now = 1; + let repository = Arc::new(InMemoryRoutingGroupRepository::seed( + [StoredRoutingGroup { + id: "system-default".to_string(), + name: "system-default".to_string(), + description: Some("test system default routing strategy".to_string()), + enabled: true, + is_system_default: true, + sort_order: 0, + config_json: serde_json::to_value(RoutingGroupConfig::default()) + .expect("default routing config should serialize"), + version: 1, + created_at: now, + updated_at: now, + published_at: Some(now), + }], + std::iter::empty::(), + std::iter::empty::(), + )); + self.with_routing_group_repository_for_tests(repository) + } + #[cfg(test)] pub(crate) fn with_auth_api_key_reader( mut self, diff --git a/apps/aether-gateway/src/tests/ai_execute/finalize_local.rs b/apps/aether-gateway/src/tests/ai_execute/finalize_local.rs index 2a699a237..3bac90c5d 100644 --- a/apps/aether-gateway/src/tests/ai_execute/finalize_local.rs +++ b/apps/aether-gateway/src/tests/ai_execute/finalize_local.rs @@ -1877,7 +1877,8 @@ async fn gateway_executes_openai_chat_antigravity_cross_format_sync_via_local_fi Arc::clone(&request_candidate_repository), Arc::clone(&usage_repository), DEVELOPMENT_ENCRYPTION_KEY, - ), + ) + .with_system_default_routing_group_for_tests(), ); let gateway = build_router_with_state(gateway_state); let (gateway_url, gateway_handle) = start_server(gateway).await; diff --git a/apps/aether-gateway/src/tests/ai_execute/finalize_local_cli/cross_format.rs b/apps/aether-gateway/src/tests/ai_execute/finalize_local_cli/cross_format.rs index d5d79c9a2..613dc13d0 100644 --- a/apps/aether-gateway/src/tests/ai_execute/finalize_local_cli/cross_format.rs +++ b/apps/aether-gateway/src/tests/ai_execute/finalize_local_cli/cross_format.rs @@ -1425,7 +1425,8 @@ async fn gateway_executes_openai_responses_antigravity_cross_format_upstream_str Arc::clone(&request_candidate_repository), Arc::clone(&usage_repository), DEVELOPMENT_ENCRYPTION_KEY, - ), + ) + .with_system_default_routing_group_for_tests(), ) .with_oauth_refresh_coordinator_for_tests(oauth_refresh); let gateway = build_router_with_state(gateway_state); diff --git a/apps/aether-gateway/src/tests/ai_execute/finalize_local_provider/gemini.rs b/apps/aether-gateway/src/tests/ai_execute/finalize_local_provider/gemini.rs index 509186eab..25da1ed30 100644 --- a/apps/aether-gateway/src/tests/ai_execute/finalize_local_provider/gemini.rs +++ b/apps/aether-gateway/src/tests/ai_execute/finalize_local_provider/gemini.rs @@ -2004,7 +2004,8 @@ async fn gateway_executes_antigravity_gemini_cli_sync_upstream_stream_via_local_ Arc::clone(&request_candidate_repository), Arc::clone(&usage_repository), DEVELOPMENT_ENCRYPTION_KEY, - ), + ) + .with_system_default_routing_group_for_tests(), ) .with_oauth_refresh_coordinator_for_tests(oauth_refresh); let gateway = build_router_with_state(gateway_state); diff --git a/apps/aether-gateway/src/tests/ai_execute/stream_provider_gemini/local_cli.rs b/apps/aether-gateway/src/tests/ai_execute/stream_provider_gemini/local_cli.rs index 5de65b509..5e2e11b21 100644 --- a/apps/aether-gateway/src/tests/ai_execute/stream_provider_gemini/local_cli.rs +++ b/apps/aether-gateway/src/tests/ai_execute/stream_provider_gemini/local_cli.rs @@ -1915,6 +1915,7 @@ async fn gateway_executes_antigravity_gemini_cli_stream_via_local_decision_gate_ Arc::clone(&request_candidate_repository), DEVELOPMENT_ENCRYPTION_KEY, ) + .with_system_default_routing_group_for_tests() .with_system_config_values_for_tests([( crate::constants::ANTIGRAVITY_BEARER_BRIDGE_CONFIG_KEY.to_string(), json!({ diff --git a/apps/aether-gateway/src/tests/ai_execute/sync/gemini/cli.rs b/apps/aether-gateway/src/tests/ai_execute/sync/gemini/cli.rs index 6beb48a01..707a995d1 100644 --- a/apps/aether-gateway/src/tests/ai_execute/sync/gemini/cli.rs +++ b/apps/aether-gateway/src/tests/ai_execute/sync/gemini/cli.rs @@ -2234,7 +2234,8 @@ async fn gateway_executes_antigravity_gemini_cli_sync_via_local_decision_gate_af provider_catalog_repository, Arc::clone(&request_candidate_repository), DEVELOPMENT_ENCRYPTION_KEY, - ), + ) + .with_system_default_routing_group_for_tests(), ) .with_oauth_refresh_coordinator_for_tests(oauth_refresh); let gateway = build_router_with_state(gateway_state); From 979dbc4b330701de9a7279a7369a0b28fd3a2ebc Mon Sep 17 00:00:00 2001 From: ZheFox <77232781+zhefox@users.noreply.github.com> Date: Thu, 3 Sep 2026 14:37:10 +0800 Subject: [PATCH 09/13] fix(data): preserve jsonb numbers for metadata CAS --- Cargo.toml | 2 +- .../adapters/postgres/src/provider_catalog.rs | 11 +++++++++++ 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/Cargo.toml b/Cargo.toml index ae184e1cb..34bd0926f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -125,7 +125,7 @@ regex = "1" rustls = { version = "0.23", features = ["ring"] } semver = "1" serde = { version = "1", features = ["derive"] } -serde_json = { version = "1", features = ["preserve_order"] } +serde_json = { version = "1", features = ["arbitrary_precision", "preserve_order"] } serde_path_to_error = "0.1" sha2 = "0.10" socket2 = "0.6" diff --git a/crates/aether-data/adapters/postgres/src/provider_catalog.rs b/crates/aether-data/adapters/postgres/src/provider_catalog.rs index 85704aef4..3c68881b5 100644 --- a/crates/aether-data/adapters/postgres/src/provider_catalog.rs +++ b/crates/aether-data/adapters/postgres/src/provider_catalog.rs @@ -3671,6 +3671,17 @@ mod tests { assert!(!sql.contains("is_active")); } + #[test] + fn runtime_metadata_cas_preserves_postgres_jsonb_numeric_lexemes() { + // PostgreSQL jsonb keeps the exact decimal value. Losing that lexeme while + // decoding into serde_json::Value makes a subsequent namespace CAS compare + // a nearby-but-different number and report a false conflict. + let postgres_json = r#"{"used_percent":9.373679999999995}"#; + let decoded: serde_json::Value = serde_json::from_str(postgres_json).unwrap(); + + assert_eq!(serde_json::to_string(&decoded).unwrap(), postgres_json); + } + #[test] fn ordinary_key_update_does_not_own_runtime_observation_fields() { let sql = super::KEY_UPDATE_SQL.to_ascii_lowercase(); From 89fe9e9f0a75b6a868ec56e4a03f642e17d59bac Mon Sep 17 00:00:00 2001 From: ZheFox <77232781+zhefox@users.noreply.github.com> Date: Thu, 3 Sep 2026 16:10:30 +0800 Subject: [PATCH 10/13] fix(runtime): avoid JSON precision stream regression --- Cargo.toml | 2 +- .../adapters/postgres/src/provider_catalog.rs | 281 +++++++++++++++--- .../execution/src/stream/ndjson.rs | 22 +- 3 files changed, 259 insertions(+), 46 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 34bd0926f..ae184e1cb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -125,7 +125,7 @@ regex = "1" rustls = { version = "0.23", features = ["ring"] } semver = "1" serde = { version = "1", features = ["derive"] } -serde_json = { version = "1", features = ["arbitrary_precision", "preserve_order"] } +serde_json = { version = "1", features = ["preserve_order"] } serde_path_to_error = "0.1" sha2 = "0.10" socket2 = "0.6" diff --git a/crates/aether-data/adapters/postgres/src/provider_catalog.rs b/crates/aether-data/adapters/postgres/src/provider_catalog.rs index 3c68881b5..a8e79e44b 100644 --- a/crates/aether-data/adapters/postgres/src/provider_catalog.rs +++ b/crates/aether-data/adapters/postgres/src/provider_catalog.rs @@ -398,7 +398,18 @@ WHERE id = $1 AND ($6::text IS NULL OR auth_config IS NOT DISTINCT FROM $6) "#; -const KEY_RUNTIME_METADATA_CAS_SQL: &str = r#" +const KEY_RUNTIME_METADATA_NAMESPACE_LOCK_SQL: &str = r#" +SELECT + jsonb_typeof(COALESCE(upstream_metadata, '{}'::jsonb)) = 'object' + AS metadata_is_object, + COALESCE(upstream_metadata, '{}'::jsonb) ? $2 AS namespace_exists, + COALESCE(upstream_metadata, '{}'::jsonb) -> $2 AS namespace_value +FROM provider_api_keys +WHERE id = $1 +FOR UPDATE +"#; + +const KEY_RUNTIME_METADATA_UPDATE_SQL: &str = r#" UPDATE provider_api_keys SET upstream_metadata = COALESCE(upstream_metadata, '{}'::jsonb) @@ -410,10 +421,58 @@ SET END WHERE id = $1 AND jsonb_typeof(COALESCE(upstream_metadata, '{}'::jsonb)) = 'object' - AND (COALESCE(upstream_metadata, '{}'::jsonb) -> $2) - IS NOT DISTINCT FROM $6::jsonb "#; +fn runtime_metadata_namespace_matches( + metadata_is_object: bool, + namespace_exists: bool, + current: Option<&serde_json::Value>, + expected: Option<&serde_json::Value>, +) -> bool { + metadata_is_object + && match expected { + Some(expected) => namespace_exists && current == Some(expected), + None => !namespace_exists, + } +} + +async fn lock_runtime_metadata_namespace_matches( + tx: &mut sqlx::Transaction<'_, Postgres>, + key_id: &str, + namespace: &str, + expected: Option<&serde_json::Value>, +) -> Result { + let Some(row) = sqlx::query(KEY_RUNTIME_METADATA_NAMESPACE_LOCK_SQL) + .bind(key_id) + .bind(namespace) + .fetch_optional(&mut **tx) + .await + .map_postgres_err()? + else { + return Ok(false); + }; + let metadata_is_object = row + .try_get::("metadata_is_object") + .map_postgres_err()?; + let namespace_exists = row + .try_get::("namespace_exists") + .map_postgres_err()?; + let current = row + .try_get::, _>("namespace_value") + .map_postgres_err()?; + + // PostgreSQL jsonb retains decimal lexemes that serde_json's default + // Number representation rounds to f64. Re-read and compare while holding + // the row lock instead of binding that rounded value back into a jsonb + // equality predicate, which would report a false CAS conflict. + Ok(runtime_metadata_namespace_matches( + metadata_is_object, + namespace_exists, + current.as_ref(), + expected, + )) +} + fn validate_key_for_update(key: &StoredProviderCatalogKey) -> Result<(), DataLayerError> { if key.id.trim().is_empty() { return Err(DataLayerError::InvalidInput( @@ -1041,6 +1100,20 @@ WHERE id = $1 .to_string(), )); } + let mut tx = self.pool.begin().await.map_postgres_err()?; + if let Some(expected) = update.expected_upstream_metadata_namespace.as_ref() { + let matches = lock_runtime_metadata_namespace_matches( + &mut tx, + &update.key_id, + &expected.namespace, + expected.expected_value.as_ref(), + ) + .await?; + if !matches { + tx.rollback().await.map_postgres_err()?; + return Ok(false); + } + } let rows_affected = sqlx::query( r#" UPDATE provider_api_keys @@ -1089,14 +1162,6 @@ WHERE id = $1 AND providers.provider_type = $18 ) ) - AND ( - $19::boolean IS FALSE - OR ( - jsonb_typeof(COALESCE(upstream_metadata, '{}'::jsonb)) = 'object' - AND (COALESCE(upstream_metadata, '{}'::jsonb) -> $20) - IS NOT DISTINCT FROM $21::jsonb - ) - ) "#, ) .bind(&update.key_id) @@ -1142,24 +1207,16 @@ WHERE id = $1 .as_ref() .map(|expected| expected.provider_type.as_str()), ) - .bind(update.expected_upstream_metadata_namespace.is_some()) - .bind( - update - .expected_upstream_metadata_namespace - .as_ref() - .map(|expected| expected.namespace.as_str()), - ) - .bind( - update - .expected_upstream_metadata_namespace - .as_ref() - .and_then(|expected| expected.expected_value.as_ref()), - ) - .execute(&self.pool) + .execute(&mut *tx) .await .map_postgres_err()? .rows_affected(); - Ok(rows_affected > 0) + if rows_affected == 0 { + tx.rollback().await.map_postgres_err()?; + return Ok(false); + } + tx.commit().await.map_postgres_err()?; + Ok(true) } pub async fn create_provider( @@ -2702,18 +2759,34 @@ WHERE id = $1 update: &ProviderCatalogKeyRuntimeMetadataUpdate, ) -> Result { validate_runtime_metadata_update(update)?; - let rows_affected = sqlx::query(KEY_RUNTIME_METADATA_CAS_SQL) + let mut tx = self.pool.begin().await.map_postgres_err()?; + if !lock_runtime_metadata_namespace_matches( + &mut tx, + &update.key_id, + &update.namespace, + update.expected_upstream_metadata_value.as_ref(), + ) + .await? + { + tx.rollback().await.map_postgres_err()?; + return Ok(false); + } + let rows_affected = sqlx::query(KEY_RUNTIME_METADATA_UPDATE_SQL) .bind(&update.key_id) .bind(&update.namespace) .bind(&update.upstream_metadata_value) .bind(&update.status_snapshot_patch) .bind(update.updated_at_unix_secs.map(|value| value as f64)) - .bind(update.expected_upstream_metadata_value.as_ref()) - .execute(&self.pool) + .execute(&mut *tx) .await .map_postgres_err()? .rows_affected(); - Ok(rows_affected > 0) + if rows_affected == 0 { + tx.rollback().await.map_postgres_err()?; + return Ok(false); + } + tx.commit().await.map_postgres_err()?; + Ok(true) } pub async fn update_key_status_snapshot( @@ -3544,6 +3617,9 @@ fn map_key_row(row: &PgRow) -> Result #[cfg(test)] mod tests { + use aether_data_contracts::repository::provider_catalog::ProviderCatalogKeyRuntimeMetadataUpdate; + use serde_json::json; + use super::SqlxProviderCatalogReadRepository; use crate::{PostgresPoolConfig, PostgresPoolFactory}; @@ -3662,24 +3738,141 @@ mod tests { } #[test] - fn runtime_metadata_cas_compares_only_the_requested_namespace() { - let sql = super::KEY_RUNTIME_METADATA_CAS_SQL.to_ascii_lowercase(); - assert!(sql.contains("upstream_metadata, '{}'::jsonb) -> $2")); - assert!(sql.contains("jsonb_typeof(coalesce(upstream_metadata, '{}'::jsonb)) = 'object'")); - assert!(sql.contains("is not distinct from $6::jsonb")); - assert!(sql.contains("status_snapshot::jsonb")); - assert!(!sql.contains("is_active")); + fn runtime_metadata_cas_locks_only_the_requested_namespace() { + let lock_sql = super::KEY_RUNTIME_METADATA_NAMESPACE_LOCK_SQL.to_ascii_lowercase(); + let update_sql = super::KEY_RUNTIME_METADATA_UPDATE_SQL.to_ascii_lowercase(); + + assert!(lock_sql.contains("upstream_metadata, '{}'::jsonb) -> $2")); + assert!(lock_sql.contains("upstream_metadata, '{}'::jsonb) ? $2")); + assert!(lock_sql.contains("for update")); + assert!(update_sql + .contains("jsonb_typeof(coalesce(upstream_metadata, '{}'::jsonb)) = 'object'")); + assert!(update_sql.contains("status_snapshot::jsonb")); + assert!(!update_sql.contains("is_active")); } #[test] - fn runtime_metadata_cas_preserves_postgres_jsonb_numeric_lexemes() { - // PostgreSQL jsonb keeps the exact decimal value. Losing that lexeme while - // decoding into serde_json::Value makes a subsequent namespace CAS compare - // a nearby-but-different number and report a false conflict. - let postgres_json = r#"{"used_percent":9.373679999999995}"#; - let decoded: serde_json::Value = serde_json::from_str(postgres_json).unwrap(); + fn runtime_metadata_namespace_cas_distinguishes_missing_from_json_null() { + assert!(super::runtime_metadata_namespace_matches( + true, false, None, None, + )); + assert!(!super::runtime_metadata_namespace_matches( + true, + true, + Some(&serde_json::Value::Null), + None, + )); + assert!(super::runtime_metadata_namespace_matches( + true, + true, + Some(&serde_json::Value::Null), + Some(&serde_json::Value::Null), + )); + assert!(!super::runtime_metadata_namespace_matches( + false, false, None, None, + )); + } - assert_eq!(serde_json::to_string(&decoded).unwrap(), postgres_json); + #[tokio::test] + #[ignore = "requires AETHER_TEST_DATABASE_URL and PostgreSQL migrations"] + async fn live_runtime_metadata_cas_handles_high_precision_jsonb_numbers() { + let database_url = std::env::var("AETHER_TEST_DATABASE_URL") + .expect("AETHER_TEST_DATABASE_URL must point at the test database"); + let factory = PostgresPoolFactory::new(PostgresPoolConfig { + database_url, + min_connections: 1, + max_connections: 2, + acquire_timeout_ms: 10_000, + idle_timeout_ms: 30_000, + max_lifetime_ms: 60_000, + statement_cache_capacity: 64, + require_ssl: false, + }) + .expect("factory should build"); + let repository = SqlxProviderCatalogReadRepository::new( + factory.connect_lazy().expect("lazy pool should build"), + ); + crate::run_migrations(repository.pool()) + .await + .expect("test database migrations should succeed"); + + let suffix = uuid::Uuid::new_v4().simple().to_string(); + let provider_id = uuid::Uuid::new_v4().to_string(); + let key_id = uuid::Uuid::new_v4().to_string(); + let provider_name = format!("provider-metadata-cas-{suffix}"); + let key_name = format!("key-metadata-cas-{suffix}"); + sqlx::query( + "INSERT INTO providers (id, name, provider_type) VALUES ($1, $2, 'antigravity')", + ) + .bind(&provider_id) + .bind(&provider_name) + .execute(repository.pool()) + .await + .expect("provider fixture should insert"); + sqlx::query( + r#" +INSERT INTO provider_api_keys ( + id, name, provider_id, total_tokens, total_cost_usd, upstream_metadata +) +VALUES ($1, $2, $3, 0, 0, $4::jsonb) +"#, + ) + .bind(&key_id) + .bind(&key_name) + .bind(&provider_id) + .bind(r#"{"antigravity":{"used_percent":0.123456789012345678901234567890}}"#) + .execute(repository.pool()) + .await + .expect("provider key fixture should insert"); + + let observed = sqlx::query_scalar::<_, serde_json::Value>( + "SELECT upstream_metadata -> 'antigravity' FROM provider_api_keys WHERE id = $1", + ) + .bind(&key_id) + .fetch_one(repository.pool()) + .await + .expect("metadata namespace should load"); + assert_ne!( + serde_json::to_string(&observed).expect("metadata should serialize"), + r#"{"used_percent":0.123456789012345678901234567890}"#, + "the fixture must exercise precision loss in serde_json's default number representation", + ); + + let updated = repository + .update_key_runtime_metadata(&ProviderCatalogKeyRuntimeMetadataUpdate { + key_id: key_id.clone(), + namespace: "antigravity".to_string(), + expected_upstream_metadata_value: Some(observed), + upstream_metadata_value: json!({"used_percent": 12.5}), + status_snapshot_patch: json!({"quota": {"used_percent": 12.5}}), + updated_at_unix_secs: Some(1_700_000_000), + }) + .await + .expect("runtime metadata CAS should execute"); + assert!( + updated, + "matching metadata must not report a false CAS conflict" + ); + + let stored = sqlx::query_scalar::<_, serde_json::Value>( + "SELECT upstream_metadata -> 'antigravity' FROM provider_api_keys WHERE id = $1", + ) + .bind(&key_id) + .fetch_one(repository.pool()) + .await + .expect("updated metadata namespace should load"); + assert_eq!(stored, json!({"used_percent": 12.5})); + + sqlx::query("DELETE FROM provider_api_keys WHERE id = $1") + .bind(&key_id) + .execute(repository.pool()) + .await + .expect("provider key fixture should delete"); + sqlx::query("DELETE FROM providers WHERE id = $1") + .bind(&provider_id) + .execute(repository.pool()) + .await + .expect("provider fixture should delete"); } #[test] diff --git a/crates/aether-gateway/execution/src/stream/ndjson.rs b/crates/aether-gateway/execution/src/stream/ndjson.rs index 990eb2411..d7ba00e66 100644 --- a/crates/aether-gateway/execution/src/stream/ndjson.rs +++ b/crates/aether-gateway/execution/src/stream/ndjson.rs @@ -17,7 +17,9 @@ pub fn decode_stream_frame_ndjson(line: &[u8]) -> Result { mod tests { use std::collections::BTreeMap; - use aether_contracts::{StreamFramePayload, StreamFrameType}; + use aether_contracts::{ + ExecutionStreamTerminalSummary, StandardizedUsage, StreamFramePayload, StreamFrameType, + }; use super::{decode_stream_frame_ndjson, encode_stream_frame_ndjson}; @@ -41,4 +43,22 @@ mod tests { decode_stream_frame_ndjson(raw.trim_ascii_end()).expect("frame should decode"); assert_eq!(decoded, frame); } + + #[test] + fn ndjson_round_trip_preserves_terminal_usage_with_fractional_fields() { + let mut usage = StandardizedUsage::new(); + usage.cache_storage_token_hours = 0.125; + let frame = + aether_contracts::StreamFrame::eof_with_summary(Some(ExecutionStreamTerminalSummary { + standardized_usage: Some(usage), + observed_finish: true, + ..ExecutionStreamTerminalSummary::default() + })); + + let raw = encode_stream_frame_ndjson(&frame).expect("frame should encode"); + let decoded = decode_stream_frame_ndjson(raw.trim_ascii_end()) + .expect("terminal usage frame should decode"); + + assert_eq!(decoded, frame); + } } From d78b5a81fb76b1b2fe2abeb95e2a07d775d02089 Mon Sep 17 00:00:00 2001 From: ZheFox <77232781+zhefox@users.noreply.github.com> Date: Thu, 3 Sep 2026 17:04:12 +0800 Subject: [PATCH 11/13] fix(gemini): preserve thought-only max-token responses --- .../src/execution_runtime/sync/execution.rs | 29 +++++++++++++++++++ .../gemini/generate_content/response.rs | 25 +++++++++++++--- 2 files changed, 50 insertions(+), 4 deletions(-) diff --git a/apps/aether-gateway/src/execution_runtime/sync/execution.rs b/apps/aether-gateway/src/execution_runtime/sync/execution.rs index 2b851598a..8ee2742ae 100644 --- a/apps/aether-gateway/src/execution_runtime/sync/execution.rs +++ b/apps/aether-gateway/src/execution_runtime/sync/execution.rs @@ -3745,6 +3745,35 @@ mod tests { assert!(message.contains("visible model output")); } + #[test] + fn invalid_gemini_provider_success_accepts_thought_only_max_tokens() { + let plan = test_gemini_chat_plan(); + let body = json!({ + "candidates": [{ + "content": { + "role": "model", + "parts": [{"text": "hidden plan", "thought": true}] + }, + "finishReason": "MAX_TOKENS" + }], + "usageMetadata": { + "promptTokenCount": 8, + "candidatesTokenCount": 0, + "thoughtsTokenCount": 24, + "totalTokenCount": 32 + } + }); + + let message = invalid_gemini_provider_success_message( + &plan, + None, + StatusCode::OK.as_u16(), + Some(&body), + ); + + assert!(message.is_none()); + } + #[test] fn invalid_gemini_provider_success_error_is_retryable_candidate_failure() { let error = invalid_gemini_provider_success_execution_error( diff --git a/crates/aether-ai/formats/src/formats/gemini/generate_content/response.rs b/crates/aether-ai/formats/src/formats/gemini/generate_content/response.rs index 149a89a99..1ca77b729 100644 --- a/crates/aether-ai/formats/src/formats/gemini/generate_content/response.rs +++ b/crates/aether-ai/formats/src/formats/gemini/generate_content/response.rs @@ -127,13 +127,15 @@ pub fn from_raw(body_json: &Value) -> Option { fn gemini_response_output_has_visible_content(output: &CanonicalResponseOutput) -> bool { output.content.iter().any(|block| match block { - CanonicalContentBlock::Text { text, .. } => !text.trim().is_empty(), + CanonicalContentBlock::Text { text, .. } | CanonicalContentBlock::Thinking { text, .. } => { + !text.trim().is_empty() + } CanonicalContentBlock::ToolUse { .. } | CanonicalContentBlock::ToolResult { .. } | CanonicalContentBlock::Image { .. } | CanonicalContentBlock::File { .. } | CanonicalContentBlock::Audio { .. } => true, - CanonicalContentBlock::Thinking { .. } | CanonicalContentBlock::Unknown { .. } => false, + CanonicalContentBlock::Unknown { .. } => false, }) } @@ -430,7 +432,7 @@ mod tests { } #[test] - fn gemini_response_with_only_thought_parts_is_not_success() { + fn gemini_response_with_only_thought_parts_is_success() { let body = json!({ "candidates": [{ "content": { @@ -443,7 +445,22 @@ mod tests { "responseId": "resp-thought-only" }); - assert!(from_raw(&body).is_none()); + let canonical = from_raw(&body).expect("thought text is representable output"); + assert!(matches!( + canonical.content.first(), + Some(CanonicalContentBlock::Thinking { text, .. }) if text == "hidden plan" + )); + assert!(matches!( + canonical.stop_reason, + Some(CanonicalStopReason::MaxTokens) + )); + + let openai = crate::canonical_to_openai_chat_response(&canonical); + assert_eq!( + openai["choices"][0]["message"]["reasoning_content"], + "hidden plan" + ); + assert_eq!(openai["choices"][0]["finish_reason"], "length"); } #[test] From f8b4382a544128534cedf2fb42e987a711d3a829 Mon Sep 17 00:00:00 2001 From: ZheFox <77232781+zhefox@users.noreply.github.com> Date: Thu, 3 Sep 2026 17:04:22 +0800 Subject: [PATCH 12/13] test(gateway): align Claude response message ids --- apps/aether-gateway/src/ai_serving/finalize/tests_sync.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/apps/aether-gateway/src/ai_serving/finalize/tests_sync.rs b/apps/aether-gateway/src/ai_serving/finalize/tests_sync.rs index 6c4acaac5..2bf0b7aa4 100644 --- a/apps/aether-gateway/src/ai_serving/finalize/tests_sync.rs +++ b/apps/aether-gateway/src/ai_serving/finalize/tests_sync.rs @@ -844,7 +844,7 @@ fn converts_claude_cli_response_to_openai_responses_response() { "output_text": "Hello Claude CLI", "output": [{ "type": "message", - "id": "msg_cli_123_msg", + "id": openai_responses_message_item_id("msg_cli_123", 0), "role": "assistant", "status": "completed", "content": [{ @@ -908,7 +908,7 @@ fn converts_claude_cli_tool_use_to_openai_responses_function_call() { "output": [ { "type": "message", - "id": "msg_cli_tool_123_msg", + "id": openai_responses_message_item_id("msg_cli_tool_123", 0), "role": "assistant", "status": "completed", "content": [{ From f6884eb8c4b8351b804bff2c47d313ccf86c1793 Mon Sep 17 00:00:00 2001 From: ZheFox <77232781+zhefox@users.noreply.github.com> Date: Thu, 3 Sep 2026 17:57:40 +0800 Subject: [PATCH 13/13] fix(gemini): accept signature-only reasoning exhaustion --- .../src/execution_runtime/sync/execution.rs | 27 ++++++ .../gemini/generate_content/response.rs | 97 ++++++++++++++++++- .../formats/gemini/generate_content/stream.rs | 53 +++++++++- .../src/formats/shared/sync_products.rs | 51 ++++++++++ .../formats/src/protocol/canonical.rs | 18 ++-- 5 files changed, 235 insertions(+), 11 deletions(-) diff --git a/apps/aether-gateway/src/execution_runtime/sync/execution.rs b/apps/aether-gateway/src/execution_runtime/sync/execution.rs index 8ee2742ae..e81c006e2 100644 --- a/apps/aether-gateway/src/execution_runtime/sync/execution.rs +++ b/apps/aether-gateway/src/execution_runtime/sync/execution.rs @@ -3774,6 +3774,33 @@ mod tests { assert!(message.is_none()); } + #[test] + fn invalid_gemini_provider_stream_success_accepts_signature_only_reasoning_exhaustion() { + let plan = test_gemini_chat_plan(); + let report_context = json!({ + "has_envelope": true, + "envelope_name": "antigravity:v1internal", + "provider_api_format": "gemini:generate_content", + }); + let body = concat!( + "data: {\"response\":{\"responseId\":\"resp_signature_only_123\",\"modelVersion\":\"gemini-3.7-flash-tiered\",", + "\"candidates\":[{\"index\":0,\"content\":{\"role\":\"model\",\"parts\":[{\"text\":\"\",\"thoughtSignature\":\"opaque-thought-signature\"}]},\"finishReason\":\"MAX_TOKENS\"}],", + "\"usageMetadata\":{\"promptTokenCount\":22,\"thoughtsTokenCount\":29,\"totalTokenCount\":51}},", + "\"traceId\":\"trace-signature-only\"}\n\n", + ); + + let message = invalid_gemini_provider_stream_success_message( + &plan, + Some(&report_context), + StatusCode::OK.as_u16(), + None, + body.as_bytes(), + true, + ); + + assert!(message.is_none()); + } + #[test] fn invalid_gemini_provider_success_error_is_retryable_candidate_failure() { let error = invalid_gemini_provider_success_execution_error( diff --git a/crates/aether-ai/formats/src/formats/gemini/generate_content/response.rs b/crates/aether-ai/formats/src/formats/gemini/generate_content/response.rs index 1ca77b729..5b4208272 100644 --- a/crates/aether-ai/formats/src/formats/gemini/generate_content/response.rs +++ b/crates/aether-ai/formats/src/formats/gemini/generate_content/response.rs @@ -26,6 +26,7 @@ pub fn from_raw(body_json: &Value) -> Option { } let candidates = body.get("candidates")?.as_array()?; + let usage = gemini_usage_to_canonical(body.get("usageMetadata")); let mut outputs = Vec::new(); for (fallback_index, candidate) in candidates.iter().enumerate() { let candidate_object = candidate.as_object()?; @@ -79,7 +80,10 @@ pub fn from_raw(body_json: &Value) -> Option { extensions, }); } - outputs.retain(gemini_response_output_has_visible_content); + outputs.retain(|output| { + gemini_response_output_has_visible_content(output) + || gemini_response_output_is_reasoning_exhausted_terminal(output, usage.as_ref()) + }); if outputs.is_empty() { return None; } @@ -106,7 +110,7 @@ pub fn from_raw(body_json: &Value) -> Option { outputs, content, stop_reason, - usage: gemini_usage_to_canonical(body.get("usageMetadata")), + usage, extensions: gemini_extensions( body, &[ @@ -139,6 +143,24 @@ fn gemini_response_output_has_visible_content(output: &CanonicalResponseOutput) }) } +fn gemini_response_output_is_reasoning_exhausted_terminal( + output: &CanonicalResponseOutput, + usage: Option<&CanonicalUsage>, +) -> bool { + matches!(output.stop_reason, Some(CanonicalStopReason::MaxTokens)) + && usage.is_some_and(|usage| usage.reasoning_tokens > 0) + && output.content.iter().any(|block| { + matches!( + block, + CanonicalContentBlock::Thinking { + text, + signature: Some(signature), + .. + } if text.trim().is_empty() && !signature.trim().is_empty() + ) + }) +} + pub fn to_raw(canonical: &CanonicalResponse, report_context: &Value) -> Option { let mut response = canonical_to_gemini_response(canonical, report_context)?; if let Some(object) = response.as_object_mut() { @@ -463,6 +485,77 @@ mod tests { assert_eq!(openai["choices"][0]["finish_reason"], "length"); } + #[test] + fn gemini_response_with_signature_only_reasoning_exhaustion_is_success() { + let body = json!({ + "candidates": [{ + "content": { + "role": "model", + "parts": [{ + "text": "", + "thoughtSignature": "opaque-thought-signature" + }] + }, + "finishReason": "MAX_TOKENS" + }], + "usageMetadata": { + "promptTokenCount": 22, + "thoughtsTokenCount": 29, + "totalTokenCount": 51 + }, + "modelVersion": "gemini-3.7-flash-tiered", + "responseId": "resp-signature-only" + }); + + let canonical = from_raw(&body).expect("reasoning exhaustion is a valid terminal"); + assert!(matches!( + canonical.content.first(), + Some(CanonicalContentBlock::Thinking { + text, + signature: Some(signature), + .. + }) if text.is_empty() && signature == "opaque-thought-signature" + )); + assert!(matches!( + canonical.stop_reason, + Some(CanonicalStopReason::MaxTokens) + )); + assert_eq!( + canonical.usage.as_ref().map(|usage| usage.reasoning_tokens), + Some(29) + ); + + let openai = crate::canonical_to_openai_chat_response(&canonical); + assert_eq!(openai["choices"][0]["finish_reason"], "length"); + assert_eq!( + openai["usage"]["completion_tokens_details"]["reasoning_tokens"], + 29 + ); + } + + #[test] + fn gemini_signature_only_terminal_without_reasoning_usage_is_not_success() { + let body = json!({ + "candidates": [{ + "content": { + "role": "model", + "parts": [{ + "text": "", + "thoughtSignature": "opaque-thought-signature" + }] + }, + "finishReason": "MAX_TOKENS" + }], + "usageMetadata": { + "promptTokenCount": 22, + "thoughtsTokenCount": 0, + "totalTokenCount": 22 + } + }); + + assert!(from_raw(&body).is_none()); + } + #[test] fn gemini_response_with_function_call_is_visible_output() { let body = json!({ diff --git a/crates/aether-ai/formats/src/formats/gemini/generate_content/stream.rs b/crates/aether-ai/formats/src/formats/gemini/generate_content/stream.rs index 8b890cacb..80d04e4af 100644 --- a/crates/aether-ai/formats/src/formats/gemini/generate_content/stream.rs +++ b/crates/aether-ai/formats/src/formats/gemini/generate_content/stream.rs @@ -129,7 +129,8 @@ impl GeminiProviderState { let is_reasoning = part_object .get("thought") .and_then(Value::as_bool) - .unwrap_or(false); + .unwrap_or(false) + || (text.trim().is_empty() && reasoning_signature.is_some()); let previous = if is_reasoning { self.reasoning_parts.entry(index).or_default() } else { @@ -960,6 +961,56 @@ mod tests { ))); } + #[test] + fn gemini_provider_state_preserves_signature_only_reasoning_terminal() { + let mut state = GeminiProviderState::default(); + let report_context = json!({}); + let frames = state + .push_line( + &report_context, + data_line(json!({ + "response": { + "responseId": "resp_signature_only_123", + "modelVersion": "gemini-3.7-flash-tiered", + "candidates": [{ + "index": 0, + "finishReason": "MAX_TOKENS", + "content": { + "role": "model", + "parts": [{ + "text": "", + "thoughtSignature": "opaque-thought-signature" + }] + } + }], + "usageMetadata": { + "promptTokenCount": 22, + "thoughtsTokenCount": 29, + "totalTokenCount": 51 + } + }, + "traceId": "trace-signature-only" + })), + ) + .expect("signature-only reasoning terminal should parse"); + + assert!(frames.iter().any(|frame| matches!( + frame.event, + CanonicalStreamEvent::ReasoningSignature(ref signature) + if signature == "opaque-thought-signature" + ))); + assert!(frames.iter().any(|frame| matches!( + frame.event, + CanonicalStreamEvent::Finish { + ref finish_reason, + usage: Some(CanonicalUsage { + reasoning_tokens: 29, + .. + }), + } if finish_reason.as_deref() == Some("length") + ))); + } + #[test] fn gemini_provider_state_parses_function_response_as_tool_result() { let mut state = GeminiProviderState::default(); diff --git a/crates/aether-ai/formats/src/formats/shared/sync_products.rs b/crates/aether-ai/formats/src/formats/shared/sync_products.rs index 8424884fa..5ad5c21ea 100644 --- a/crates/aether-ai/formats/src/formats/shared/sync_products.rs +++ b/crates/aether-ai/formats/src/formats/shared/sync_products.rs @@ -4257,6 +4257,57 @@ mod tests { assert_eq!(aggregated["usageMetadata"]["totalTokenCount"], 5); } + #[test] + fn aggregates_antigravity_signature_only_reasoning_exhaustion() { + let body = concat!( + "data: {\"response\":{\"responseId\":\"resp_signature_only_123\",\"modelVersion\":\"gemini-3.7-flash-tiered\",", + "\"candidates\":[{\"index\":0,\"content\":{\"role\":\"model\",\"parts\":[{\"text\":\"\",\"thoughtSignature\":\"opaque-thought-signature\"}]},\"finishReason\":\"MAX_TOKENS\"}],", + "\"usageMetadata\":{\"promptTokenCount\":22,\"thoughtsTokenCount\":29,\"totalTokenCount\":51}},", + "\"traceId\":\"trace-signature-only\"}\n\n", + ); + + let aggregated = aggregate_gemini_stream_sync_response(body.as_bytes()) + .expect("signature-only reasoning terminal should aggregate"); + + assert_eq!( + aggregated["candidates"][0]["content"]["parts"][0]["thought"], + true + ); + assert_eq!( + aggregated["candidates"][0]["content"]["parts"][0]["thoughtSignature"], + "opaque-thought-signature" + ); + assert_eq!(aggregated["candidates"][0]["finishReason"], "MAX_TOKENS"); + assert_eq!(aggregated["usageMetadata"]["thoughtsTokenCount"], 29); + assert!( + crate::formats::gemini::generate_content::response::from_raw(&aggregated).is_some() + ); + + let report_context = json!({ + "provider_api_format": "gemini:generate_content", + "client_api_format": "openai:chat", + "mapped_model": "gemini-3.7-flash-tiered", + }); + let product = maybe_build_standard_cross_format_sync_product_from_normalized_payload( + "openai_chat_sync_finalize", + 200, + Some(&report_context), + None, + Some(&base64::engine::general_purpose::STANDARD.encode(body)), + ) + .expect("signature-only reasoning terminal should convert") + .expect("cross-format product should exist"); + + assert_eq!( + product.client_body_json["choices"][0]["finish_reason"], + "length" + ); + assert_eq!( + product.client_body_json["usage"]["completion_tokens_details"]["reasoning_tokens"], + 29 + ); + } + #[test] fn gemini_stream_aggregation_rejects_unknown_parts() { let body = "data: {\"responseId\":\"resp_gem_unknown_123\",\"modelVersion\":\"gemini-2.5-pro\",\"candidates\":[{\"index\":0,\"content\":{\"role\":\"model\",\"parts\":[{\"futurePart\":{\"kept\":true}}]}}]}\n\n"; diff --git a/crates/aether-ai/formats/src/protocol/canonical.rs b/crates/aether-ai/formats/src/protocol/canonical.rs index bf4db8aab..f9fe69512 100644 --- a/crates/aether-ai/formats/src/protocol/canonical.rs +++ b/crates/aether-ai/formats/src/protocol/canonical.rs @@ -1249,19 +1249,21 @@ pub(crate) fn gemini_part_to_canonical_block( ) -> Option { let part_object = part.as_object()?; if let Some(text) = part_object.get("text").and_then(Value::as_str) { - if part_object + let thought_signature = part_object + .get("thoughtSignature") + .or_else(|| part_object.get("thought_signature")) + .and_then(Value::as_str) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned); + let is_thinking = part_object .get("thought") .and_then(Value::as_bool) .unwrap_or(false) - { + || (text.trim().is_empty() && thought_signature.is_some()); + if is_thinking { return Some(CanonicalContentBlock::Thinking { text: text.to_string(), - signature: part_object - .get("thoughtSignature") - .or_else(|| part_object.get("thought_signature")) - .and_then(Value::as_str) - .filter(|value| !value.is_empty()) - .map(ToOwned::to_owned), + signature: thought_signature, encrypted_content: None, extensions: gemini_extensions( part_object,