Merge pull request #789 from zhefox/codex/fix-antigravity-quota

Codex/fix antigravity quota
This commit is contained in:
ZheFox
2026-09-03 18:48:58 +08:00
committed by GitHub
57 changed files with 3224 additions and 285 deletions
@@ -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": [{
@@ -843,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": [{
@@ -907,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": [{
@@ -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": [{
@@ -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());
@@ -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,
@@ -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,
@@ -213,6 +218,21 @@ impl GatewayDataState {
self
}
#[cfg(test)]
pub(crate) fn with_cached_provider_catalog_reader_for_tests<T>(
mut self,
repository: Arc<T>,
) -> Self
where
T: ProviderCatalogReadRepository + 'static,
{
let inner: Arc<dyn ProviderCatalogReadRepository> = 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,
@@ -877,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::<StoredRoutingGroupBinding>(),
std::iter::empty::<StoredRoutingGroupVersion>(),
));
self.with_routing_group_repository_for_tests(repository)
}
#[cfg(test)]
pub(crate) fn with_auth_api_key_reader(
mut self,
@@ -318,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(
@@ -1456,13 +1458,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)
@@ -1479,7 +1498,19 @@ async fn read_pool_catalog_key_contexts_by_id(
),
)
})
.collect()
.collect::<BTreeMap<_, _>>();
// 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(
@@ -1724,7 +1755,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()
@@ -2000,7 +2045,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,
@@ -2112,6 +2157,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<_>>(),
vec!["key-ready"]
);
assert_eq!(
skipped
.iter()
.map(|item| (item.candidate.key_id.as_str(), item.skip_reason))
.collect::<Vec<_>>(),
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(
@@ -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]
@@ -3745,6 +3745,62 @@ 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_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(
@@ -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<String, String>,
proxy_override: Option<&ProxySnapshot>,
) -> Option<serde_json::Value> {
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();
@@ -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<Output = ()>,
{
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::<Vec<_>>()
})
.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",
&current_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)
);
}
}
@@ -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
@@ -50,6 +50,7 @@ pub(super) fn finalize_gateway_response(
mut response: Response<Body>,
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<AdmissionPermit>,
) -> Response<Body> {
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),
@@ -978,6 +978,7 @@ async fn proxy_request_inner(
response,
&trace_id,
&remote_addr,
client_ip,
request.method(),
request
.uri()
@@ -1014,6 +1015,7 @@ async fn proxy_request_inner(
response,
&trace_id,
&remote_addr,
client_ip,
request.method(),
request
.uri()
@@ -1054,6 +1056,7 @@ async fn proxy_request_inner(
response,
&trace_id,
&remote_addr,
client_ip,
request.method(),
request
.uri()
@@ -1114,6 +1117,7 @@ async fn proxy_request_inner(
response,
&trace_id,
&remote_addr,
client_ip,
&parts.method,
parts
.uri
@@ -1135,6 +1139,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,
@@ -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()
@@ -1089,6 +1091,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()
@@ -1143,6 +1146,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()
@@ -1209,6 +1213,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(
@@ -654,6 +654,112 @@ fn antigravity_model_quota_window_snapshot(
Some(window)
}
fn antigravity_grouped_quota_window_snapshots(
metadata: &Map<String, Value>,
observed_at_unix_secs: Option<u64>,
) -> Vec<Value> {
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<String, Value>,
fields: &[&str],
@@ -989,7 +1095,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(
@@ -1010,6 +1116,17 @@ fn build_codex_quota_status_snapshot(
.into_iter()
.flatten()
.collect::<Vec<_>>();
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()
@@ -1611,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()
@@ -1625,6 +1742,13 @@ fn build_antigravity_quota_status_snapshot(
.collect::<Vec<_>>()
})
.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;
@@ -3121,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::<Vec<_>>();
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();
@@ -4155,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();
@@ -1213,6 +1213,18 @@ fn probe_result_hard_state(item: &Value) -> Option<PoolMemberHardState> {
.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<PoolMemberHardState> {
}
}
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,
@@ -2042,6 +2042,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),
@@ -2064,6 +2073,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
@@ -3635,6 +3665,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]
@@ -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()
@@ -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;
@@ -1949,7 +1950,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"
@@ -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": [{
@@ -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,
@@ -1424,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);
@@ -1465,7 +1467,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": [{
@@ -1537,7 +1539,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,
@@ -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": [{
@@ -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);
@@ -2123,7 +2124,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,
@@ -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!({
@@ -2017,7 +2018,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"
@@ -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);
@@ -2321,7 +2322,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"
@@ -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<serde_json::Value>,
}
@@ -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::<SeenExecutionRuntimeRequest>));
let seen_execution_runtime = Arc::new(Mutex::new(Vec::<SeenExecutionRuntimeRequest>::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()])
@@ -2466,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))
@@ -2476,11 +2546,31 @@ 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))
);
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();
+24
View File
@@ -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<bool> {
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(),
+318 -9
View File
@@ -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";
const CODEX_ACTIVE_LIMIT_HEADER: &str = "x-codex-active-limit";
const CODEX_HEADER_PREFIX: &str = "x-codex-";
const CODEX_LIMIT_NAME_HEADER_SUFFIX: &str = "-limit-name";
@@ -281,10 +282,128 @@ 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,
}))
}
pub fn parse_antigravity_quota_summary_response(
value: &serde_json::Value,
) -> Option<serde_json::Value> {
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,
@@ -2022,6 +2141,119 @@ 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<String, serde_json::Value>,
) -> Vec<serde_json::Value> {
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;
};
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))
.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<String, serde_json::Value>,
) -> Option<&serde_json::Map<String, serde_json::Value>> {
@@ -2128,6 +2360,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));
@@ -3829,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,
@@ -6249,6 +6489,43 @@ 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!(
additional.is_empty(),
"the Spark windows are already normalized into the dedicated spark fields"
);
}
#[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]
@@ -6749,20 +7026,52 @@ 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")
);
}
#[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(
+1
View File
@@ -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,
@@ -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
@@ -26,6 +26,7 @@ pub fn from_raw(body_json: &Value) -> Option<CanonicalResponse> {
}
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<CanonicalResponse> {
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<CanonicalResponse> {
outputs,
content,
stop_reason,
usage: gemini_usage_to_canonical(body.get("usageMetadata")),
usage,
extensions: gemini_extensions(
body,
&[
@@ -127,16 +131,36 @@ pub fn from_raw(body_json: &Value) -> Option<CanonicalResponse> {
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,
})
}
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<Value> {
let mut response = canonical_to_gemini_response(canonical, report_context)?;
if let Some(object) = response.as_object_mut() {
@@ -430,7 +454,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,6 +467,92 @@ mod tests {
"responseId": "resp-thought-only"
});
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]
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());
}
@@ -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();
@@ -4,7 +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,
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,
@@ -2328,7 +2328,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 {
@@ -2347,7 +2347,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()
}
@@ -4357,7 +4357,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\":"));
@@ -4421,8 +4421,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]
@@ -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"] {
@@ -10,6 +10,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_";
const GEMINI_TOOL_SIGNATURE_CARRIER_PREFIX: &str = "cpa-gemini-responses-carrier-v1:";
const MAX_GEMINI_THOUGHT_SIGNATURE_LEN: usize = 32 * 1024 * 1024;
const MAX_GEMINI_THOUGHT_SIGNATURE_ENCODED_LEN: usize =
@@ -102,6 +103,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
/// `<response-id>_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
@@ -255,6 +317,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,
strip_incompatible_openai_responses_reasoning_items,
strip_incompatible_openai_responses_reasoning_items_with_policy,
@@ -348,6 +411,38 @@ 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);
@@ -7,8 +7,10 @@ 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::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,
};
use aether_ai_formats::formats::registry::{convert_response, FormatContext, FormatError};
use aether_ai_formats::{
canonical_response_unknown_block_count, canonical_to_claude_response,
@@ -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())
@@ -4236,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";
+1
View File
@@ -57,6 +57,7 @@ pub use formats::openai::responses::request::{
validate_openai_responses_request_contract, OpenAiResponsesRequestContractViolation,
};
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,
strip_incompatible_openai_responses_reasoning_items,
strip_incompatible_openai_responses_reasoning_items_with_policy,
@@ -3,6 +3,7 @@ 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,
};
@@ -1248,19 +1249,21 @@ pub(crate) fn gemini_part_to_canonical_block(
) -> Option<CanonicalContentBlock> {
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,
@@ -4102,11 +4105,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,
@@ -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<bool, DataLayerError> {
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::<bool, _>("metadata_is_object")
.map_postgres_err()?;
let namespace_exists = row
.try_get::<bool, _>("namespace_exists")
.map_postgres_err()?;
let current = row
.try_get::<Option<serde_json::Value>, _>("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<bool, DataLayerError> {
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<StoredProviderCatalogKey, DataLayerError>
#[cfg(test)]
mod tests {
use aether_data_contracts::repository::provider_catalog::ProviderCatalogKeyRuntimeMetadataUpdate;
use serde_json::json;
use super::SqlxProviderCatalogReadRepository;
use crate::{PostgresPoolConfig, PostgresPoolFactory};
@@ -3662,13 +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_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,
));
}
#[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]
@@ -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 =
@@ -8,6 +8,7 @@ pub struct PublicRequestContext<Decision> {
pub request_query_string: Option<String>,
pub request_content_type: Option<String>,
pub host_header: Option<String>,
pub client_ip: Option<String>,
pub control_decision: Option<Decision>,
}
@@ -32,6 +33,7 @@ impl<Decision> PublicRequestContext<Decision> {
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,
}
}
@@ -17,7 +17,9 @@ pub fn decode_stream_frame_ndjson(line: &[u8]) -> Result<StreamFrame, IoError> {
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);
}
}
+22
View File
@@ -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]
+72 -20
View File
@@ -1037,7 +1037,9 @@ fn parse_antigravity_models_response(body: &Value) -> Result<(Vec<Value>, 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<String, Value> {
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::<f64>().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()));
+1 -1
View File
@@ -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),
+33 -3
View File
@@ -18,6 +18,8 @@ pub struct PoolSchedulingPreset {
pub struct PoolSchedulingConfig {
pub scheduling_presets: Vec<PoolSchedulingPreset>,
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<u64>,
}
@@ -225,9 +227,14 @@ fn schedule_pool_group<Candidate>(
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<_>>(),
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)
+340 -10
View File
@@ -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,14 +28,16 @@ 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_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,
@@ -882,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();
@@ -969,6 +1032,273 @@ 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()
+12 -2
View File
@@ -7,8 +7,9 @@ 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_quota_snapshot_exhausted_decision, provider_pool_quota_usage_ratio,
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,
};
#[derive(Debug, Clone)]
@@ -16,6 +17,10 @@ pub struct ProviderPoolMemberInput<'a> {
pub provider_type: &'a str,
pub key: &'a StoredProviderCatalogKey,
pub auth_config: Option<&'a Map<String, Value>>,
/// 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>,
}
@@ -71,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)
}
@@ -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<String, String>,
) -> 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!({})));
}
}
@@ -11,8 +11,9 @@ 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_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;
@@ -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)
{
@@ -8,8 +8,9 @@ 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_metadata_bucket, provider_pool_model_quota_exhausted,
provider_pool_quota_snapshot_exhausted_decision, 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)
{
@@ -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)
{
@@ -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::{
@@ -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;
}
+501 -52
View File
@@ -32,54 +32,22 @@ pub fn provider_pool_key_quota_hard_blocked(
})
}
pub(crate) fn provider_pool_model_quota_exhausted(
/// 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,
) -> Option<bool> {
let quota_snapshot = provider_pool_member_quota_snapshot(key, provider_type)?;
let windows = quota_snapshot.get("windows")?.as_array()?;
let normalized_provider = provider_type.trim().to_ascii_lowercase();
let normalized_model = provider_model_name.trim().to_ascii_lowercase();
let matches_window = |window: &Map<String, Value>| {
let code = window
.get("code")
.and_then(Value::as_str)
.unwrap_or_default()
.trim()
.to_ascii_lowercase();
if normalized_provider == "codex" {
let spark_model = normalized_model.contains("spark");
return code.starts_with("spark_") == spark_model;
}
if normalized_provider == "antigravity" {
return window
.get("model")
.and_then(Value::as_str)
.is_some_and(|model| model.trim().eq_ignore_ascii_case(&normalized_model));
}
false
};
let matching_windows = windows
.iter()
.filter_map(Value::as_object)
.filter(|window| matches_window(window))
.collect::<Vec<_>>();
if matching_windows.is_empty() {
return None;
}
let now_unix_secs = provider_pool_current_unix_secs();
let snapshot_observed_at = provider_pool_timestamp_unix_secs(quota_snapshot.get("observed_at"))
.or_else(|| provider_pool_timestamp_unix_secs(quota_snapshot.get("updated_at")));
Some(matching_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)
})
}))
) -> 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),
})
}
pub fn provider_pool_member_quota_snapshot<'a>(
@@ -96,6 +64,421 @@ 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<bool> {
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<u64>, 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::<Vec<_>>();
if !model_matches.is_empty() {
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),
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::<Vec<_>>();
if !family_matches.is_empty() {
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),
observed_at,
)
{
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::<Vec<_>>();
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));
}
}
}
resolved.map(|(_, exhausted)| exhausted)
}
fn provider_pool_should_replace_model_quota_resolution(
previous_observed_at: Option<u64>,
next_observed_at: Option<u64>,
) -> 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<bool> {
provider_pool_model_quota_exhausted(key, provider_type, provider_model_name)
}
fn provider_pool_explicit_model_windows_exhausted(
windows: Vec<&Map<String, Value>>,
snapshot_observed_at: Option<u64>,
) -> bool {
let now_unix_secs = provider_pool_current_unix_secs();
// 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| {
provider_pool_reset_deadline_elapsed(window, snapshot_observed_at, now)
})
})
}
fn provider_pool_any_window_exhausted(
windows: Vec<&Map<String, Value>>,
snapshot_observed_at: Option<u64>,
) -> 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<String, Value>) -> 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<String, Value>,
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<String, Value>,
requested_tokens: &std::collections::BTreeSet<String>,
) -> 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<String>,
) -> 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",
"tiered",
"vendor",
]
.contains(&token)
}
fn provider_pool_window_has_explicit_model(window: &Map<String, Value>) -> 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<String> {
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<String, Value>) -> Vec<Map<String, Value>> {
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,
@@ -249,12 +632,54 @@ pub(crate) fn provider_pool_reset_deadline_elapsed(
fn provider_pool_quota_window_is_exhausted(window: &Map<String, Value>) -> 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<String, Value>) -> 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<String, Value>,
provider_type: &str,
@@ -289,6 +714,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)
@@ -317,16 +754,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::<Vec<_>>();
// 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::<f64>;
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)));
@@ -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<String, String> {
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),
@@ -45,11 +45,11 @@
<div class="py-2">
<div
v-if="items.length > 0"
v-if="displayItems.length > 0"
class="grid grid-cols-2 gap-3"
>
<div
v-for="item in items"
v-for="item in displayItems"
:key="item.model"
>
<div class="flex items-center justify-between text-[10px] mb-0.5">
@@ -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<QuotaItem[]>(() => {
const snapshotItems = buildItemsFromQuotaSnapshot(props.quotaSnapshot)
if (snapshotItems.length > 0) return snapshotItems
@@ -291,6 +333,8 @@ const rawItems = computed<QuotaItem[]>(() => {
})
const items = computed<QuotaItem[]>(() => summarizeAntigravityQuotaItems(rawItems.value))
const groupedItems = computed<QuotaItem[]>(() => buildGroupedItemsFromQuotaSnapshot(props.quotaSnapshot))
const displayItems = computed<QuotaItem[]>(() => [...groupedItems.value, ...items.value])
async function handleTestModel(modelName: string) {
if (!props.providerId || testingModel.value) return
@@ -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({