mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-10 21:20:20 +08:00
Merge pull request #789 from zhefox/codex/fix-antigravity-quota
Codex/fix antigravity quota
This commit is contained in:
@@ -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",
|
||||
¤t_namespace,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.expect("out-of-band metadata update should succeed"));
|
||||
let still_cached = app
|
||||
.read_provider_catalog_keys_by_ids(&key_ids)
|
||||
.await
|
||||
.expect("stale cached read should succeed");
|
||||
assert_eq!(
|
||||
still_cached[0].upstream_metadata.as_ref().unwrap()["antigravity"]["quota_by_model"]
|
||||
["gemini-3.7-flash-tiered"]["remaining_fraction"],
|
||||
json!(0.9),
|
||||
"regression setup must keep the ordinary read stale"
|
||||
);
|
||||
|
||||
let metadata_update = json!({
|
||||
"antigravity": {
|
||||
"project_id": "project-1",
|
||||
"quota_by_model": {
|
||||
"gemini-3.7-flash-tiered": {"remaining_fraction": 0.6}
|
||||
},
|
||||
"quota_groups": [{
|
||||
"display_name": "Gemini models",
|
||||
"buckets": [{"bucket_id": "gemini-weekly", "window": "weekly"}]
|
||||
}]
|
||||
}
|
||||
});
|
||||
assert!(persist_provider_quota_refresh_state(
|
||||
&admin_state,
|
||||
"key-antigravity-stale-cache",
|
||||
Some(&metadata_update),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.expect("quota refresh persistence should not error"));
|
||||
|
||||
let stored = repository
|
||||
.list_keys_by_ids(&key_ids)
|
||||
.await
|
||||
.expect("key should reload")
|
||||
.pop()
|
||||
.expect("key should exist");
|
||||
assert_eq!(
|
||||
stored.upstream_metadata.as_ref().unwrap()["antigravity"]["quota_groups"][0]["buckets"]
|
||||
[0]["bucket_id"],
|
||||
json!("gemini-weekly")
|
||||
);
|
||||
assert_eq!(
|
||||
stored.upstream_metadata.as_ref().unwrap()["antigravity"]["model_fetch_revision"],
|
||||
json!(2),
|
||||
"quota refresh must preserve fields written by another Antigravity metadata producer"
|
||||
);
|
||||
assert_eq!(
|
||||
stored.upstream_metadata.as_ref().unwrap()["antigravity"]["quota_by_model"]
|
||||
["gemini-3.7-flash-tiered"]["remaining_fraction"],
|
||||
json!(0.6)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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";
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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()));
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
+32
-2
@@ -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({
|
||||
|
||||
Reference in New Issue
Block a user