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

Merge remote-tracking branch 'origin/main' into worktree-linear-enchanting-bunny

Browse Source
This commit is contained in:
elky
2026-09-03 20:49:18 +08:00
parent 1de2e70d41 89b57464d2
commit 670d5e8d33
65 changed files with 3439 additions and 564 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/ai_serving/finalize/tests_sync.rs
+7 -6
View File
@@ -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": [{
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs
+2 -2
View File
@@ -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());
apps/aether-gateway/src/ai_serving/pure/mod.rs
+1
View File
@@ -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,
apps/aether-gateway/src/data/state/testing/mod.rs
+44
View File
@@ -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,
@@ -905,6 +925,30 @@ impl GatewayDataState {
self.with_routing_group_repository_for_tests(repository)
}
#[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,
apps/aether-gateway/src/dispatch/pool_scheduler.rs
+100 -6
View File
@@ -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(
apps/aether-gateway/src/execution_runtime/grok.rs
+7 -2
View File
@@ -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]
apps/aether-gateway/src/execution_runtime/sync/execution.rs
+56
View File
@@ -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(
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/antigravity.rs
+107 -7
View File
@@ -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();
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/shared.rs
+141 -5
View File
@@ -240,7 +240,16 @@ fn merge_upstream_metadata(
.unwrap_or_default();
if let Some(update_object) = updates.as_object() {
for (key, value) in update_object {
merged.insert(key.clone(), value.clone());
let mut next = value.clone();
if let (Some(current_namespace), Some(next_namespace)) = (
merged.get(key).and_then(serde_json::Value::as_object),
next.as_object_mut(),
) {
let mut combined = current_namespace.clone();
combined.extend(next_namespace.clone());
next = serde_json::Value::Object(combined);
}
merged.insert(key.clone(), next);
}
}
serde_json::Value::Object(merged)
@@ -1306,7 +1315,8 @@ where
F: std::future::Future<Output = ()>,
{
let Some(mut latest_key) = state
.read_provider_catalog_keys_by_ids(&[key_id.to_string()])
.app()
.list_provider_catalog_keys_by_ids_strong(&[key_id.to_string()])
.await?
.into_iter()
.next()
@@ -1350,9 +1360,18 @@ where
let metadata_updates = metadata_update
.and_then(serde_json::Value::as_object)
.map(|updates| {
let merged = latest_key
.upstream_metadata
.as_ref()
.and_then(serde_json::Value::as_object);
updates
.iter()
.map(|(namespace, value)| (namespace.clone(), value.clone()))
.keys()
.filter_map(|namespace| {
merged
.and_then(|metadata| metadata.get(namespace))
.cloned()
.map(|value| (namespace.clone(), value))
})
.collect::<Vec<_>>()
})
.unwrap_or_default();
@@ -1384,7 +1403,7 @@ where
} else {
serde_json::json!({})
};
let mut expected = observed_upstream_metadata
let expected = observed_upstream_metadata
.as_ref()
.and_then(serde_json::Value::as_object)
.and_then(|metadata| metadata.get(namespace))
@@ -2807,4 +2826,121 @@ mod tests {
json!({"remaining":4})
);
}
#[tokio::test]
async fn quota_refresh_strong_read_bypasses_stale_provider_catalog_cache() {
let mut key = StoredProviderCatalogKey::new(
"key-antigravity-stale-cache".to_string(),
"provider-antigravity-stale-cache".to_string(),
"Antigravity stale cache".to_string(),
"oauth".to_string(),
None,
true,
)
.expect("key should build");
key.upstream_metadata = Some(json!({
"antigravity": {
"project_id": "project-1",
"quota_by_model": {
"gemini-3.7-flash-tiered": {"remaining_fraction": 0.9}
}
}
}));
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![],
vec![],
vec![key],
));
let data =
GatewayDataState::with_provider_catalog_repository_for_tests(Arc::clone(&repository))
.with_cached_provider_catalog_reader_for_tests(Arc::clone(&repository));
let app = AppState::new()
.expect("app should build")
.with_data_state_for_tests(data);
let admin_state = AdminAppState::new(&app);
let key_ids = ["key-antigravity-stale-cache".to_string()];
let cached = app
.read_provider_catalog_keys_by_ids(&key_ids)
.await
.expect("initial cached read should succeed");
assert_eq!(
cached[0].upstream_metadata.as_ref().unwrap()["antigravity"]["quota_by_model"]
["gemini-3.7-flash-tiered"]["remaining_fraction"],
json!(0.9)
);
let current_namespace = json!({
"project_id": "project-1",
"model_fetch_revision": 2,
"quota_by_model": {
"gemini-3.7-flash-tiered": {"remaining_fraction": 0.7}
}
});
assert!(repository
.upsert_key_upstream_metadata_namespace(
"key-antigravity-stale-cache",
"antigravity",
&current_namespace,
None,
)
.await
.expect("out-of-band metadata update should succeed"));
let still_cached = app
.read_provider_catalog_keys_by_ids(&key_ids)
.await
.expect("stale cached read should succeed");
assert_eq!(
still_cached[0].upstream_metadata.as_ref().unwrap()["antigravity"]["quota_by_model"]
["gemini-3.7-flash-tiered"]["remaining_fraction"],
json!(0.9),
"regression setup must keep the ordinary read stale"
);
let metadata_update = json!({
"antigravity": {
"project_id": "project-1",
"quota_by_model": {
"gemini-3.7-flash-tiered": {"remaining_fraction": 0.6}
},
"quota_groups": [{
"display_name": "Gemini models",
"buckets": [{"bucket_id": "gemini-weekly", "window": "weekly"}]
}]
}
});
assert!(persist_provider_quota_refresh_state(
&admin_state,
"key-antigravity-stale-cache",
Some(&metadata_update),
None,
None,
None,
)
.await
.expect("quota refresh persistence should not error"));
let stored = repository
.list_keys_by_ids(&key_ids)
.await
.expect("key should reload")
.pop()
.expect("key should exist");
assert_eq!(
stored.upstream_metadata.as_ref().unwrap()["antigravity"]["quota_groups"][0]["buckets"]
[0]["bucket_id"],
json!("gemini-weekly")
);
assert_eq!(
stored.upstream_metadata.as_ref().unwrap()["antigravity"]["model_fetch_revision"],
json!(2),
"quota refresh must preserve fields written by another Antigravity metadata producer"
);
assert_eq!(
stored.upstream_metadata.as_ref().unwrap()["antigravity"]["quota_by_model"]
["gemini-3.7-flash-tiered"]["remaining_fraction"],
json!(0.6)
);
}
}
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/writes.rs
+42 -12
View File
@@ -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
apps/aether-gateway/src/handlers/proxy/finalize.rs
+12
View File
@@ -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),
apps/aether-gateway/src/handlers/proxy/mod.rs
+5
View File
@@ -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,
apps/aether-gateway/src/handlers/proxy/websocket/live/http.rs
+5
View File
@@ -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(
apps/aether-gateway/src/handlers/shared/catalog.rs
+234 -2
View File
@@ -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();
apps/aether-gateway/src/maintenance/runtime/pool_quota_probe.rs
+52
View File
@@ -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,
apps/aether-gateway/src/orchestration/effects.rs
+37
View File
@@ -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]
apps/aether-gateway/src/scheduler/candidate/runtime.rs
+24 -9
View File
@@ -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()
apps/aether-gateway/src/tests/ai_execute/finalize_local.rs
+3 -2
View File
@@ -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"
apps/aether-gateway/src/tests/ai_execute/finalize_local_cli/compact.rs
+2 -1
View File
@@ -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": [{
apps/aether-gateway/src/tests/ai_execute/finalize_local_cli/cross_format.rs
+5 -3
View File
@@ -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,
apps/aether-gateway/src/tests/ai_execute/finalize_local_cli/direct.rs
+2 -1
View File
@@ -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": [{
apps/aether-gateway/src/tests/ai_execute/finalize_local_provider/gemini.rs
+3 -2
View File
@@ -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,
apps/aether-gateway/src/tests/ai_execute/stream_provider_gemini/local_cli.rs
+2 -1
View File
@@ -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"
apps/aether-gateway/src/tests/ai_execute/sync/gemini/cli.rs
+3 -2
View File
@@ -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"
apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs
+135 -33
View File
@@ -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();