From fe8ff268df3a37eff9c2ab381f016427c1839470 Mon Sep 17 00:00:00 2001 From: ZheFox <77232781+zhefox@users.noreply.github.com> Date: Fri, 4 Sep 2026 11:34:41 +0800 Subject: [PATCH 01/13] fix oauth identity and codex reset credits --- .../admin/provider/oauth/quota/shared.rs | 66 ++++++- .../admin/provider/oauth/state/exchange.rs | 22 ++- .../src/tests/control/admin/oauth.rs | 169 ++++++++++++++++ crates/aether-admin/src/provider/quota.rs | 101 +++++++++- .../src/provider/providers/antigravity.rs | 185 +++++++++++++++++- .../src/provider/providers/mod.rs | 2 +- frontend/src/views/admin/PoolManagement.vue | 13 +- 7 files changed, 545 insertions(+), 13 deletions(-) diff --git a/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/shared.rs b/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/shared.rs index 55f3d0837..648fc7803 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/shared.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/shared.rs @@ -570,6 +570,50 @@ pub(crate) async fn reserve_codex_account_reset( Ok(None) } +fn record_locally_consumed_codex_reset_credit( + codex: &mut serde_json::Map, + observed_at_unix_secs: u64, +) { + let Some(reset_credits) = codex + .get_mut("reset_credits") + .and_then(serde_json::Value::as_object_mut) + else { + return; + }; + let Some(available_count) = reset_credits + .get("available_count") + .and_then(admin_provider_quota_pure::coerce_json_u64) + else { + return; + }; + + reset_credits.insert( + "available_count".to_string(), + serde_json::json!(available_count.saturating_sub(1)), + ); + reset_credits.insert( + "updated_at".to_string(), + serde_json::json!(observed_at_unix_secs), + ); + reset_credits.insert( + "detail_source".to_string(), + serde_json::json!("local_consume"), + ); + reset_credits.insert( + "detail_status".to_string(), + serde_json::json!("pending_refresh"), + ); + reset_credits.remove("detail_error"); + if let Some(credits) = reset_credits + .get_mut("credits") + .and_then(serde_json::Value::as_array_mut) + { + if !credits.is_empty() { + credits.remove(0); + } + } +} + pub(crate) async fn complete_codex_account_reset( state: &AdminAppState<'_>, key_id: &str, @@ -635,6 +679,9 @@ pub(crate) async fn complete_codex_account_reset( generation: reservation.generation, outcome: outcome.to_string(), }; + if outcome == "reset" { + record_locally_consumed_codex_reset_credit(&mut codex, fence_unix_ms / 1_000); + } codex_reset_write_bounded_history(&mut codex, &terminal); if codex_reset_reservation_from_object(&codex).as_ref() == Some(reservation) { codex.remove(admin_provider_quota_pure::CODEX_QUOTA_ACCOUNT_RESET_RESERVATION_KEY); @@ -1671,7 +1718,19 @@ mod tests { .expect("key should build"); key.encrypted_auth_config = Some("auth-v1".to_string()); key.upstream_metadata = Some(json!({ - "codex": {"credential_generation": "credential-v1"} + "codex": { + "credential_generation": "credential-v1", + "reset_credits": { + "available_count": 2, + "updated_at": 100u64, + "detail_source": "wham_readonly", + "detail_status": "available", + "credits": [ + {"id": "credit-1", "expires_at": 20_000u64}, + {"id": "credit-2", "expires_at": 30_000u64} + ] + } + } })); let credential = ProviderCatalogKeyOAuthCredentialFence { encrypted_api_key: None, @@ -1926,6 +1985,11 @@ mod tests { codex["account_quota_reset_history"][0]["outcome"], json!("reset") ); + assert_eq!(codex["reset_credits"]["available_count"], json!(1u64)); + assert_eq!( + codex["reset_credits"]["credits"], + json!([{"id": "credit-2", "expires_at": 30_000u64}]) + ); } } diff --git a/apps/aether-gateway/src/handlers/admin/provider/oauth/state/exchange.rs b/apps/aether-gateway/src/handlers/admin/provider/oauth/state/exchange.rs index bfda465fb..b737510f9 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/oauth/state/exchange.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/oauth/state/exchange.rs @@ -4,8 +4,9 @@ use super::super::errors::{ use crate::handlers::admin::request::{AdminAppState, AdminProviderOAuthTemplate}; use aether_contracts::ProxySnapshot; use aether_oauth::provider::providers::{ - ClaudeCodeProviderOAuthAdapter, GenericProviderOAuthAdapter, CLAUDE_CODE_PROVIDER_TYPE, - CLAUDE_CODE_TOKEN_URL, CLAUDE_CODE_WEB_BASE_URL, + AntigravityProviderOAuthAdapter, ClaudeCodeProviderOAuthAdapter, GenericProviderOAuthAdapter, + ANTIGRAVITY_USER_INFO_URL, CLAUDE_CODE_PROVIDER_TYPE, CLAUDE_CODE_TOKEN_URL, + CLAUDE_CODE_WEB_BASE_URL, }; use aether_oauth::provider::{ ProviderOAuthCookieAuthorizationInput, ProviderOAuthService, ProviderOAuthTransportContext, @@ -43,7 +44,14 @@ fn provider_oauth_exchange_context( fn provider_oauth_service_for_template( template: AdminProviderOAuthTemplate, token_url: String, + antigravity_user_info_url: String, ) -> Result> { + if template.provider_type.eq_ignore_ascii_case("antigravity") { + let adapter = AntigravityProviderOAuthAdapter::default() + .with_token_url_override(token_url) + .with_user_info_url_override(antigravity_user_info_url); + return Ok(ProviderOAuthService::new().with_adapter(Arc::new(adapter))); + } GenericProviderOAuthAdapter::for_provider_type(template.provider_type) .map(|adapter| adapter.with_token_url_override(token_url)) .map(|adapter| ProviderOAuthService::new().with_adapter(Arc::new(adapter))) @@ -75,7 +83,10 @@ pub(crate) async fn exchange_admin_provider_oauth_code( proxy: Option, ) -> Result> { let token_url = state.provider_oauth_token_url(template.provider_type, template.token_url); - let service = provider_oauth_service_for_template(template, token_url)?; + let antigravity_user_info_url = + state.provider_oauth_token_url("antigravity_user_info", ANTIGRAVITY_USER_INFO_URL); + let service = + provider_oauth_service_for_template(template, token_url, antigravity_user_info_url)?; let ctx = provider_oauth_exchange_context(template.provider_type, proxy); let executor = crate::oauth::GatewayOAuthHttpExecutor::new(*state); let result = service @@ -103,7 +114,10 @@ pub(crate) async fn exchange_admin_provider_oauth_refresh_token( proxy: Option, ) -> Result> { let token_url = state.provider_oauth_token_url(template.provider_type, template.token_url); - let service = provider_oauth_service_for_template(template, token_url)?; + let antigravity_user_info_url = + state.provider_oauth_token_url("antigravity_user_info", ANTIGRAVITY_USER_INFO_URL); + let service = + provider_oauth_service_for_template(template, token_url, antigravity_user_info_url)?; let ctx = provider_oauth_exchange_context(template.provider_type, proxy); let executor = crate::oauth::GatewayOAuthHttpExecutor::new(*state); let input = aether_oauth::provider::ProviderOAuthImportInput { diff --git a/apps/aether-gateway/src/tests/control/admin/oauth.rs b/apps/aether-gateway/src/tests/control/admin/oauth.rs index 8f330ca32..864680d45 100644 --- a/apps/aether-gateway/src/tests/control/admin/oauth.rs +++ b/apps/aether-gateway/src/tests/control/admin/oauth.rs @@ -3918,6 +3918,175 @@ async fn gateway_completes_admin_provider_oauth_provider_locally_with_trusted_ad upstream_handle.abort(); } +#[test] +fn gateway_names_new_antigravity_oauth_account_from_google_userinfo_email() { + run_admin_oauth_test( + "gateway_names_new_antigravity_oauth_account_from_google_userinfo_email", + gateway_names_new_antigravity_oauth_account_from_google_userinfo_email_impl, + ); +} + +async fn gateway_names_new_antigravity_oauth_account_from_google_userinfo_email_impl() { + let upstream_hits = Arc::new(Mutex::new(0usize)); + let upstream_hits_clone = Arc::clone(&upstream_hits); + let upstream = Router::new().fallback(any(move |_request: Request| { + let upstream_hits_inner = Arc::clone(&upstream_hits_clone); + async move { + *upstream_hits_inner.lock().expect("mutex should lock") += 1; + (StatusCode::OK, Body::from("unexpected upstream hit")) + } + })); + + let token_hits = Arc::new(Mutex::new(0usize)); + let token_hits_clone = Arc::clone(&token_hits); + let user_info_hits = Arc::new(Mutex::new(0usize)); + let user_info_hits_clone = Arc::clone(&user_info_hits); + let seen_user_info_authorization = Arc::new(Mutex::new(None::)); + let seen_user_info_authorization_clone = Arc::clone(&seen_user_info_authorization); + let google_server = Router::new() + .route( + "/oauth/token", + post(move || { + let token_hits_inner = Arc::clone(&token_hits_clone); + async move { + *token_hits_inner.lock().expect("mutex should lock") += 1; + Json(json!({ + "access_token": "antigravity-access-token", + "refresh_token": "antigravity-refresh-token", + "token_type": "Bearer", + "expires_in": 3600, + "scope": "https://www.googleapis.com/auth/userinfo.email" + })) + } + }), + ) + .route( + "/oauth/userinfo", + get(move |headers: HeaderMap| { + let user_info_hits_inner = Arc::clone(&user_info_hits_clone); + let seen_authorization_inner = Arc::clone(&seen_user_info_authorization_clone); + async move { + *user_info_hits_inner.lock().expect("mutex should lock") += 1; + *seen_authorization_inner.lock().expect("mutex should lock") = headers + .get(http::header::AUTHORIZATION) + .and_then(|value| value.to_str().ok()) + .map(ToOwned::to_owned); + Json(json!({ + "email": "new-antigravity@example.com", + "verified_email": true, + "name": "Antigravity User" + })) + } + }), + ); + + let mut provider = sample_provider("provider-antigravity", "antigravity", 10); + provider.provider_type = "antigravity".to_string(); + let endpoint = sample_endpoint( + "endpoint-antigravity", + "provider-antigravity", + "gemini:generate_content", + "https://daily-cloudcode-pa.googleapis.com", + ); + let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( + vec![provider], + vec![endpoint], + vec![], + )); + + let (upstream_url, upstream_handle) = start_server(upstream).await; + let (google_url, google_handle) = start_server(google_server).await; + let gateway = build_router_with_state( + AppState::new() + .expect("gateway should build") + .with_data_state_for_tests( + GatewayDataState::with_provider_catalog_repository_for_tests( + provider_catalog_repository.clone(), + ) + .with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY), + ) + .with_provider_oauth_state_entry_for_tests( + "nonce-antigravity-123", + json!({ + "nonce": "nonce-antigravity-123", + "key_id": "", + "provider_id": "provider-antigravity", + "provider_type": "antigravity", + "pkce_verifier": "verifier-antigravity-123", + }), + ) + .with_provider_oauth_token_url_for_tests( + "antigravity", + format!("{google_url}/oauth/token"), + ) + .with_provider_oauth_token_url_for_tests( + "antigravity_user_info", + format!("{google_url}/oauth/userinfo"), + ), + ); + let (gateway_url, gateway_handle) = start_server(gateway).await; + + let response = reqwest::Client::new() + .post(format!( + "{gateway_url}/api/admin/provider-oauth/providers/provider-antigravity/complete" + )) + .header(crate::constants::GATEWAY_HEADER, "rust-phase3b") + .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123") + .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin") + .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123") + .json(&json!({ + "callback_url": "http://localhost:51121/oauth2callback?code=antigravity-code-123&state=nonce-antigravity-123" + })) + .send() + .await + .expect("request should succeed"); + + let status = response.status(); + let payload: Value = response.json().await.expect("json body should parse"); + assert_eq!(status, StatusCode::OK, "payload={payload}"); + assert_eq!(payload["provider_type"], "antigravity"); + assert_eq!(payload["email"], "new-antigravity@example.com"); + assert_eq!(payload["replaced"], false); + assert_eq!(*token_hits.lock().expect("mutex should lock"), 1); + assert_eq!(*user_info_hits.lock().expect("mutex should lock"), 1); + assert_eq!( + seen_user_info_authorization + .lock() + .expect("mutex should lock") + .as_deref(), + Some("Bearer antigravity-access-token") + ); + assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0); + + let key_id = payload["key_id"] + .as_str() + .expect("created key id should be returned") + .to_string(); + let persisted_keys = provider_catalog_repository + .list_keys_by_ids(std::slice::from_ref(&key_id)) + .await + .expect("created key should load"); + let persisted = persisted_keys.first().expect("created key should exist"); + assert_eq!(persisted.name, "new-antigravity@example.com"); + let decrypted_auth_config = decrypt_python_fernet_ciphertext( + DEVELOPMENT_ENCRYPTION_KEY, + persisted + .encrypted_auth_config + .as_deref() + .expect("auth config should be stored"), + ) + .expect("auth config should decrypt"); + let auth_config: Value = + serde_json::from_str(&decrypted_auth_config).expect("auth config json should parse"); + assert_eq!(auth_config["email"], "new-antigravity@example.com"); + assert_eq!(auth_config["refresh_token"], "antigravity-refresh-token"); + + gateway_handle.abort(); + google_handle.abort(); + upstream_handle.abort(); + drop(upstream_url); +} + #[test] fn gateway_imports_admin_provider_oauth_refresh_token_locally_with_trusted_admin_principal() { run_admin_oauth_test( diff --git a/crates/aether-admin/src/provider/quota.rs b/crates/aether-admin/src/provider/quota.rs index 349f26855..4603f65af 100644 --- a/crates/aether-admin/src/provider/quota.rs +++ b/crates/aether-admin/src/provider/quota.rs @@ -1887,6 +1887,46 @@ fn codex_quota_is_account_status_key(key: &str) -> bool { matches!(key, "allowed" | "limit_reached") } +fn codex_quota_merge_reset_credits( + current_object: &serde_json::Map, + incoming: &serde_json::Value, +) -> serde_json::Value { + let Some(incoming_object) = incoming.as_object() else { + return incoming.clone(); + }; + let mut merged = current_object + .get("reset_credits") + .and_then(serde_json::Value::as_object) + .cloned() + .unwrap_or_default(); + let failed_detail = incoming_object + .get("detail_status") + .and_then(serde_json::Value::as_str) + .is_some_and(|status| status.trim().eq_ignore_ascii_case("failed")); + if !failed_detail { + return incoming.clone(); + } + + for (key, value) in incoming_object { + // A failed readonly-detail request contributes diagnostics, not an + // authoritative empty list. Keep the last known items and count so a + // transient 429 cannot make reset credits disappear from the UI. + if failed_detail + && key == "credits" + && value.as_array().is_some_and(|credits| credits.is_empty()) + && merged + .get("credits") + .and_then(serde_json::Value::as_array) + .is_some_and(|credits| !credits.is_empty()) + { + continue; + } + merged.insert(key.clone(), value.clone()); + } + + serde_json::Value::Object(merged) +} + /// Merge a parsed Codex quota observation into the stored flat metadata. /// /// Positive `window_minutes` values identify windows independently of the @@ -1961,7 +2001,12 @@ pub fn merge_codex_quota_metadata_snapshot( { continue; } - merged.insert(key.clone(), value.clone()); + let value = if key == "reset_credits" { + codex_quota_merge_reset_credits(¤t_object, value) + } else { + value.clone() + }; + merged.insert(key.clone(), value); } if let Some(incoming_order) = context.request_order().filter(|incoming| { codex_quota_request_order_is_newer(*incoming, stored_metadata_watermark) @@ -4777,6 +4822,60 @@ mod tests { assert_eq!(outcome.metadata["primary_reset_at"], json!(20_000u64)); } + #[test] + fn codex_quota_failed_reset_credit_detail_preserves_last_known_count_and_items() { + let current = json!({ + "reset_credits": { + "available_count": 2, + "updated_at": 100u64, + "detail_source": "wham_readonly", + "detail_status": "available", + "credits": [{ + "id": "credit-1", + "display_key": "credit", + "status": "available", + "expires_at": 20_000u64 + }] + }, + "updated_at": 100u64 + }); + let incoming = json!({ + "reset_credits": { + "updated_at": 110u64, + "detail_source": "wham_readonly", + "detail_status": "failed", + "detail_error": "HTTP 429", + "credits": [] + } + }); + + let outcome = merge_codex_quota( + Some(¤t), + &incoming, + 110, + 110_000, + CodexQuotaWindowCoverage::Patch, + ); + + assert!(outcome.changed); + assert_eq!( + outcome.metadata["reset_credits"]["available_count"], + json!(2u64) + ); + assert_eq!( + outcome.metadata["reset_credits"]["credits"][0]["id"], + json!("credit-1") + ); + assert_eq!( + outcome.metadata["reset_credits"]["detail_status"], + json!("failed") + ); + assert_eq!( + outcome.metadata["reset_credits"]["detail_error"], + json!("HTTP 429") + ); + } + #[test] fn codex_quota_explicit_reset_allows_usage_drop_with_same_deadline() { let current = json!({ diff --git a/crates/aether-oauth/src/provider/providers/antigravity.rs b/crates/aether-oauth/src/provider/providers/antigravity.rs index 9bc86c053..644a82812 100644 --- a/crates/aether-oauth/src/provider/providers/antigravity.rs +++ b/crates/aether-oauth/src/provider/providers/antigravity.rs @@ -1,11 +1,18 @@ use super::generic::{ provider_account_state_from_metadata, template_for_provider_type, GenericProviderOAuthAdapter, }; -use crate::provider::ProviderOAuthAdapter; +use crate::core::OAuthError; +use crate::network::{OAuthHttpExecutor, OAuthHttpRequest}; +use crate::provider::{ProviderOAuthAdapter, ProviderOAuthTokenSet, ProviderOAuthTransportContext}; +use serde_json::Value; +use std::collections::BTreeMap; + +pub const ANTIGRAVITY_USER_INFO_URL: &str = "https://www.googleapis.com/oauth2/v2/userinfo"; #[derive(Debug, Clone)] pub struct AntigravityProviderOAuthAdapter { inner: GenericProviderOAuthAdapter, + user_info_url: String, } impl Default for AntigravityProviderOAuthAdapter { @@ -15,10 +22,95 @@ impl Default for AntigravityProviderOAuthAdapter { template_for_provider_type("antigravity") .expect("antigravity template should exist"), ), + user_info_url: ANTIGRAVITY_USER_INFO_URL.to_string(), } } } +impl AntigravityProviderOAuthAdapter { + pub fn with_token_url_override(mut self, token_url: impl Into) -> Self { + self.inner = self.inner.with_token_url_override(token_url); + self + } + + pub fn with_user_info_url_override(mut self, user_info_url: impl Into) -> Self { + self.user_info_url = user_info_url.into(); + self + } + + async fn enrich_google_identity( + &self, + executor: &dyn OAuthHttpExecutor, + ctx: &ProviderOAuthTransportContext, + mut result: ProviderOAuthTokenSet, + ) -> Result { + if result + .auth_config + .get("email") + .and_then(Value::as_str) + .is_some_and(|email| !email.trim().is_empty()) + { + return Ok(result); + } + + let response = executor + .execute(OAuthHttpRequest { + request_id: "provider-oauth:antigravity-user-info".to_string(), + method: reqwest::Method::GET, + url: self.user_info_url.clone(), + headers: BTreeMap::from([ + ("accept".to_string(), "application/json".to_string()), + ( + "authorization".to_string(), + result.token_set.bearer_header_value(), + ), + ]), + content_type: None, + json_body: None, + body_bytes: None, + network: ctx.network.clone(), + transport_profile: None, + }) + .await?; + if !(200..300).contains(&response.status_code) { + return Err(OAuthError::HttpStatus { + status_code: response.status_code, + body_excerpt: response.body_text.trim().chars().take(500).collect(), + }); + } + + let profile = response + .json_body + .or_else(|| serde_json::from_str::(&response.body_text).ok()) + .ok_or_else(|| OAuthError::invalid_response("userinfo response is not json"))?; + if profile.get("verified_email").and_then(Value::as_bool) == Some(false) { + return Err(OAuthError::invalid_response( + "userinfo response returned an unverified email", + )); + } + let email = profile + .get("email") + .and_then(Value::as_str) + .map(str::trim) + .filter(|email| !email.is_empty()) + .ok_or_else(|| OAuthError::invalid_response("userinfo response missing email"))? + .to_string(); + + if let Some(auth_config) = result.auth_config.as_object_mut() { + auth_config.insert("email".to_string(), Value::String(email.clone())); + } + if let Some(token_payload) = result + .token_set + .raw_payload + .as_mut() + .and_then(Value::as_object_mut) + { + token_payload.insert("email".to_string(), Value::String(email)); + } + Ok(result) + } +} + #[async_trait::async_trait] impl ProviderOAuthAdapter for AntigravityProviderOAuthAdapter { fn provider_type(&self) -> &'static str { @@ -59,9 +151,11 @@ impl ProviderOAuthAdapter for AntigravityProviderOAuthAdapter { state: &str, pkce_verifier: Option<&str>, ) -> Result { - self.inner + let result = self + .inner .exchange_code(executor, ctx, code, state, pkce_verifier) - .await + .await?; + self.enrich_google_identity(executor, ctx, result).await } async fn import_credentials( @@ -111,7 +205,7 @@ impl ProviderOAuthAdapter for AntigravityProviderOAuthAdapter { #[cfg(test)] mod tests { - use super::AntigravityProviderOAuthAdapter; + use super::{AntigravityProviderOAuthAdapter, ANTIGRAVITY_USER_INFO_URL}; use crate::network::{OAuthHttpExecutor, OAuthHttpRequest, OAuthHttpResponse}; use crate::provider::{ ProviderOAuthAccount, ProviderOAuthAdapter, ProviderOAuthTransportContext, @@ -119,9 +213,15 @@ mod tests { use async_trait::async_trait; use serde_json::json; use std::collections::BTreeMap; + use std::sync::Mutex; struct UnusedExecutor; + #[derive(Default)] + struct GoogleOAuthExecutor { + requests: Mutex>, + } + fn transport_context() -> ProviderOAuthTransportContext { ProviderOAuthTransportContext { provider_id: String::new(), @@ -148,6 +248,43 @@ mod tests { } } + #[async_trait] + impl OAuthHttpExecutor for GoogleOAuthExecutor { + async fn execute( + &self, + request: OAuthHttpRequest, + ) -> Result { + let request_id = request.request_id.clone(); + self.requests + .lock() + .expect("requests should lock") + .push(request); + match request_id.as_str() { + "provider-oauth:exchange-code" => Ok(OAuthHttpResponse { + status_code: 200, + body_text: json!({ + "access_token": "google-access-token", + "refresh_token": "google-refresh-token", + "token_type": "Bearer", + "expires_in": 3600 + }) + .to_string(), + json_body: None, + }), + "provider-oauth:antigravity-user-info" => Ok(OAuthHttpResponse { + status_code: 200, + body_text: json!({ + "email": "antigravity@example.com", + "verified_email": true + }) + .to_string(), + json_body: None, + }), + other => panic!("unexpected OAuth request: {other}"), + } + } + } + #[test] fn antigravity_authorize_requests_offline_refresh_token() { let adapter = AntigravityProviderOAuthAdapter::default(); @@ -171,6 +308,46 @@ mod tests { ); } + #[tokio::test] + async fn antigravity_exchange_fetches_google_email_for_account_identity() { + let adapter = AntigravityProviderOAuthAdapter::default(); + let ctx = transport_context(); + let executor = GoogleOAuthExecutor::default(); + + let result = adapter + .exchange_code( + &executor, + &ctx, + "authorization-code", + "state-1", + Some("verifier-1"), + ) + .await + .expect("Antigravity OAuth exchange should succeed"); + + assert_eq!( + result.auth_config.get("email"), + Some(&json!("antigravity@example.com")) + ); + assert_eq!( + result + .token_set + .raw_payload + .as_ref() + .and_then(|payload| payload.get("email")), + Some(&json!("antigravity@example.com")) + ); + let requests = executor.requests.lock().expect("requests should lock"); + assert_eq!(requests.len(), 2); + assert_eq!(requests[1].url, ANTIGRAVITY_USER_INFO_URL); + assert_eq!(requests[1].method, reqwest::Method::GET); + assert_eq!( + requests[1].headers.get("authorization").map(String::as_str), + Some("Bearer google-access-token") + ); + assert_eq!(requests[1].network, ctx.network); + } + #[tokio::test] async fn antigravity_probe_marks_forbidden_metadata_invalid() { let adapter = AntigravityProviderOAuthAdapter::default(); diff --git a/crates/aether-oauth/src/provider/providers/mod.rs b/crates/aether-oauth/src/provider/providers/mod.rs index 4fe916a75..07f752f36 100644 --- a/crates/aether-oauth/src/provider/providers/mod.rs +++ b/crates/aether-oauth/src/provider/providers/mod.rs @@ -5,7 +5,7 @@ mod generic; mod kiro; mod windsurf; -pub use antigravity::AntigravityProviderOAuthAdapter; +pub use antigravity::{AntigravityProviderOAuthAdapter, ANTIGRAVITY_USER_INFO_URL}; pub use claude_code::{ ClaudeCodeProviderOAuthAdapter, CLAUDE_CODE_AUTHORIZE_URL, CLAUDE_CODE_CLIENT_ID, CLAUDE_CODE_COOKIE_SCOPE, CLAUDE_CODE_OAUTH_SCOPES, CLAUDE_CODE_PROVIDER_TYPE, diff --git a/frontend/src/views/admin/PoolManagement.vue b/frontend/src/views/admin/PoolManagement.vue index 684e7a7b8..843ab750f 100644 --- a/frontend/src/views/admin/PoolManagement.vue +++ b/frontend/src/views/admin/PoolManagement.vue @@ -1153,6 +1153,7 @@ import { getCodexResetCreditAvailableCount, getCodexResetCreditReservationIdempotencyKey, getVisibleCodexResetCreditItems, + mergeCodexQuotaDisplays, readPendingCodexResetCreditIdempotencyKey, rememberPendingCodexResetCreditIdempotencyKey, } from '@/features/providers/components/codex-reset-credit-display' @@ -2235,9 +2236,17 @@ function applyQuotaRefreshResultToCurrentPage(result: Awaited Date: Fri, 4 Sep 2026 12:06:39 +0800 Subject: [PATCH 02/13] fix(antigravity): sync discovered models into catalog --- .../admin/provider/oauth/quota/antigravity.rs | 62 +++++++++++++++++++ .../tests/control/admin/endpoints/quota.rs | 29 +++++++++ crates/aether-model-fetch/src/lib.rs | 5 +- crates/aether-model-fetch/src/strategy.rs | 10 ++- 4 files changed, 103 insertions(+), 3 deletions(-) diff --git a/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/antigravity.rs b/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/antigravity.rs index d3e26fc79..2cbd0848a 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/antigravity.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/antigravity.rs @@ -5,6 +5,7 @@ use super::shared::{ quota_key_auto_removed, quota_refresh_success_invalid_state, resolve_provider_quota_execution_timeouts, ProviderQuotaExecutionOutcome, }; +use crate::handlers::admin::provider::shared::payloads::AdminImportProviderModelsRequest; use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot}; use crate::GatewayError; use aether_admin::provider::quota::{ @@ -22,6 +23,63 @@ use std::collections::BTreeMap; use std::time::{SystemTime, UNIX_EPOCH}; use tracing::warn; +fn antigravity_discovered_model_ids(metadata_update: Option<&serde_json::Value>) -> Vec { + metadata_update + .and_then(|value| value.pointer("/antigravity/quota_by_model")) + .and_then(serde_json::Value::as_object) + .into_iter() + .flat_map(|models| models.keys()) + .map(String::as_str) + .filter(|model_id| aether_model_fetch::antigravity_model_id_is_routable(model_id)) + .map(ToOwned::to_owned) + .collect() +} + +async fn sync_antigravity_discovered_models( + state: &AdminAppState<'_>, + provider_id: &str, + metadata_update: Option<&serde_json::Value>, +) { + if !state.has_global_model_data_reader() || !state.has_global_model_data_writer() { + return; + } + let model_ids = antigravity_discovered_model_ids(metadata_update); + if model_ids.is_empty() { + return; + } + + let result = state + .build_admin_import_provider_models_payload( + provider_id, + AdminImportProviderModelsRequest { + model_ids, + tiered_pricing: None, + price_per_request: None, + }, + ) + .await; + match result { + Ok(payload) => { + let errors = payload + .get("errors") + .and_then(serde_json::Value::as_array) + .map(Vec::len) + .unwrap_or(0); + if errors > 0 { + warn!( + provider_id, + errors, "Antigravity discovered-model catalog sync completed with item errors" + ); + } + } + Err(error) => warn!( + provider_id, + error = %error, + "Antigravity discovered-model catalog sync failed" + ), + } +} + async fn execute_antigravity_quota_plan( state: &AdminAppState<'_>, transport: &AdminGatewayProviderTransportSnapshot, @@ -330,6 +388,10 @@ pub(crate) async fn refresh_antigravity_provider_quota_locally( continue; } + if status == "success" { + sync_antigravity_discovered_models(state, &provider.id, metadata_update.as_ref()).await; + } + if status == "success" { success_count += 1; } else { diff --git a/apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs b/apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs index 444d550f2..7120cf871 100644 --- a/apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs +++ b/apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs @@ -4,8 +4,12 @@ use std::sync::{Arc, Mutex}; use aether_crypto::{ decrypt_python_fernet_ciphertext, encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY, }; +use aether_data::repository::global_models::InMemoryGlobalModelReadRepository; use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository; use aether_data::repository::proxy_nodes::InMemoryProxyNodeRepository; +use aether_data_contracts::repository::global_models::{ + AdminProviderModelListQuery, GlobalModelReadRepository, +}; use aether_data_contracts::repository::provider_catalog::{ ProviderCatalogReadRepository, StoredProviderCatalogKey, StoredProviderCatalogProvider, }; @@ -2330,6 +2334,12 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru }, "gemini-2.5-pro": { "displayName": "Gemini 2.5 Pro" + }, + "gemini-3.7-flash-tiered": { + "displayName": "Gemini 3.7 Flash" + }, + "chat_23310": { + "displayName": "Internal Chat" } } }), @@ -2431,6 +2441,7 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru )], vec![key], )); + let global_model_repository = Arc::new(InMemoryGlobalModelReadRepository::default()); let (upstream_url, upstream_handle) = start_server(upstream).await; let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await; @@ -2440,6 +2451,7 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru GatewayDataState::with_provider_catalog_repository_for_tests( provider_catalog_repository.clone(), ) + .with_global_model_repository_for_tests(global_model_repository.clone()) .with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY), ), ); @@ -2541,6 +2553,23 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru .and_then(|value| value.get("remaining_fraction")), Some(&json!(0.25)) ); + let imported_provider_models = global_model_repository + .list_admin_provider_models(&AdminProviderModelListQuery { + provider_id: "provider-antigravity".to_string(), + is_active: None, + offset: 0, + limit: 100, + }) + .await + .expect("imported Antigravity provider models should read"); + let imported_model_names = imported_provider_models + .iter() + .map(|model| model.provider_model_name.as_str()) + .collect::>(); + assert!(imported_model_names.contains("claude-sonnet-4")); + assert!(imported_model_names.contains("gemini-2.5-pro")); + assert!(imported_model_names.contains("gemini-3.7-flash-tiered")); + assert!(!imported_model_names.contains("chat_23310")); assert_eq!( reloaded[0] .upstream_metadata diff --git a/crates/aether-model-fetch/src/lib.rs b/crates/aether-model-fetch/src/lib.rs index bb3432805..8ef4b5855 100644 --- a/crates/aether-model-fetch/src/lib.rs +++ b/crates/aether-model-fetch/src/lib.rs @@ -21,8 +21,9 @@ pub use logic::{ upstream_metadata_namespace_updates, ModelFetchRunSummary, ModelsFetchPage, ModelsFetchSuccess, }; pub use strategy::{ - fetch_models_from_transports, fetch_models_from_transports_for_client_version, - ModelFetchStrategy, ModelFetchStrategyKind, ModelsFetchOutcome, SelectedModelFetchStrategy, + antigravity_model_id_is_routable, fetch_models_from_transports, + fetch_models_from_transports_for_client_version, ModelFetchStrategy, ModelFetchStrategyKind, + ModelsFetchOutcome, SelectedModelFetchStrategy, }; pub use transport::{ build_antigravity_fetch_available_models_plan, build_antigravity_load_code_assist_plan, diff --git a/crates/aether-model-fetch/src/strategy.rs b/crates/aether-model-fetch/src/strategy.rs index 64951adb8..249bec490 100644 --- a/crates/aether-model-fetch/src/strategy.rs +++ b/crates/aether-model-fetch/src/strategy.rs @@ -1018,7 +1018,7 @@ fn parse_antigravity_models_response(body: &Value) -> Result<(Vec, Option let mut quota_by_model = serde_json::Map::new(); for (model_id, model_data) in models_object { let model_id = model_id.trim(); - if model_id.is_empty() || ANTIGRAVITY_BLOCKED_MODELS.contains(&model_id) { + if !antigravity_model_id_is_routable(model_id) { continue; } let model_object = model_data.as_object().cloned().unwrap_or_default(); @@ -1054,6 +1054,14 @@ fn parse_antigravity_models_response(body: &Value) -> Result<(Vec, Option Ok((models, upstream_metadata)) } +pub fn antigravity_model_id_is_routable(model_id: &str) -> bool { + let model_id = model_id.trim(); + !model_id.is_empty() + && !ANTIGRAVITY_BLOCKED_MODELS + .iter() + .any(|blocked| blocked.eq_ignore_ascii_case(model_id)) +} + fn parse_kiro_available_models_response( body: &Value, ) -> Result<(Vec, Option), String> { From c8d1ae3e7ea4fb8021f0bd3cf2c59f0484c055c6 Mon Sep 17 00:00:00 2001 From: ZheFox <77232781+zhefox@users.noreply.github.com> Date: Fri, 4 Sep 2026 13:02:22 +0800 Subject: [PATCH 03/13] test(codex): preserve reset credit fixture metadata --- .../handlers/admin/provider/oauth/quota/shared.rs | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/shared.rs b/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/shared.rs index 648fc7803..9baaa7d39 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/shared.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/shared.rs @@ -1824,6 +1824,13 @@ mod tests { let key_id = "key-codex-reset-credential-generation"; let (app, repository, credential) = codex_reset_state_machine_test_state(key_id); let admin_state = AdminAppState::new(&app); + let original_metadata = repository + .list_keys_by_ids(&[key_id.to_string()]) + .await + .expect("key should load before reservation") + .pop() + .expect("key should exist before reservation") + .upstream_metadata; let result = reserve_codex_account_reset( &admin_state, @@ -1847,10 +1854,7 @@ mod tests { .expect("key should reload") .pop() .expect("key should exist"); - assert_eq!( - stored.upstream_metadata.unwrap()["codex"], - json!({"credential_generation":"credential-v1"}) - ); + assert_eq!(stored.upstream_metadata, original_metadata); } #[tokio::test] From ba11a72214db1fd903199f0cea546fd3e561e989 Mon Sep 17 00:00:00 2001 From: ZheFox <77232781+zhefox@users.noreply.github.com> Date: Fri, 4 Sep 2026 13:54:50 +0800 Subject: [PATCH 04/13] test(gateway): align CI fixtures with routing contracts --- .../src/ai_serving/finalize/tests_sync.rs | 5 ++-- .../src/tests/architecture/admin_provider.rs | 5 ++-- .../src/tests/control/admin/models/global.rs | 9 +++---- .../src/tests/control/admin/system.rs | 9 +++---- apps/aether-gateway/src/tests/frontdoor/ai.rs | 6 ++++- .../src/tests/frontdoor/internal.rs | 6 ++++- apps/aether-gateway/src/tests/proxy.rs | 25 ++++++++++++++++--- 7 files changed, 43 insertions(+), 22 deletions(-) diff --git a/apps/aether-gateway/src/ai_serving/finalize/tests_sync.rs b/apps/aether-gateway/src/ai_serving/finalize/tests_sync.rs index 2bf0b7aa4..cd34e1417 100644 --- a/apps/aether-gateway/src/ai_serving/finalize/tests_sync.rs +++ b/apps/aether-gateway/src/ai_serving/finalize/tests_sync.rs @@ -1,6 +1,5 @@ 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; @@ -12,10 +11,10 @@ use super::{ convert_gemini_chat_response_to_openai_chat, convert_gemini_response_to_openai_responses, maybe_build_local_core_sync_finalize_response, }; -use crate::ai_serving::GatewayControlDecision; use crate::ai_serving::{ convert_openai_chat_response_to_openai_responses, - convert_openai_responses_response_to_openai_chat, + convert_openai_responses_response_to_openai_chat, openai_responses_message_item_id, + GatewayControlDecision, }; use crate::usage::GatewaySyncReportRequest; diff --git a/apps/aether-gateway/src/tests/architecture/admin_provider.rs b/apps/aether-gateway/src/tests/architecture/admin_provider.rs index b67ad922e..4173e3fe8 100644 --- a/apps/aether-gateway/src/tests/architecture/admin_provider.rs +++ b/apps/aether-gateway/src/tests/architecture/admin_provider.rs @@ -1963,8 +1963,9 @@ fn admin_provider_oauth_quota_mod_stays_thin() { "handlers/admin/provider/oauth/quota/antigravity.rs should import common quota helpers from shared.rs" ); assert!( - quota_antigravity - .contains("use aether_provider_pool::build_antigravity_pool_quota_request;"), + quota_antigravity.contains("use aether_provider_pool::{") + && quota_antigravity.contains("build_antigravity_pool_quota_request") + && quota_antigravity.contains("build_antigravity_pool_quota_summary_request"), "handlers/admin/provider/oauth/quota/antigravity.rs should delegate antigravity quota request construction to aether-provider-pool" ); let quota_chatgpt_web = read_workspace_file( diff --git a/apps/aether-gateway/src/tests/control/admin/models/global.rs b/apps/aether-gateway/src/tests/control/admin/models/global.rs index 38027db47..cec076492 100644 --- a/apps/aether-gateway/src/tests/control/admin/models/global.rs +++ b/apps/aether-gateway/src/tests/control/admin/models/global.rs @@ -850,10 +850,7 @@ async fn gateway_handles_admin_global_model_routing_locally_with_trusted_admin_p provider_catalog_repository, ) .with_global_model_repository_for_tests(global_model_repository) - .with_system_config_values_for_tests(vec![ - ("scheduling_mode".to_string(), json!("fixed_order")), - ("provider_priority_mode".to_string(), json!("global_key")), - ]), + .with_system_default_routing_group_for_tests(), ), ); let (gateway_url, gateway_handle) = start_server(gateway).await; @@ -876,8 +873,8 @@ async fn gateway_handles_admin_global_model_routing_locally_with_trusted_admin_p assert_eq!(payload["global_model_name"], "gpt-5"); assert_eq!(payload["display_name"], "GPT 5"); assert_eq!(payload["global_model_mappings"], json!(["gpt-5-upstream"])); - assert_eq!(payload["scheduling_mode"], "fixed_order"); - assert_eq!(payload["priority_mode"], "global_key"); + assert_eq!(payload["scheduling_mode"], "cache_affinity"); + assert_eq!(payload["priority_mode"], "provider"); assert_eq!(payload["total_providers"], 2); assert_eq!(payload["active_providers"], 2); diff --git a/apps/aether-gateway/src/tests/control/admin/system.rs b/apps/aether-gateway/src/tests/control/admin/system.rs index d2763f71a..b54269579 100644 --- a/apps/aether-gateway/src/tests/control/admin/system.rs +++ b/apps/aether-gateway/src/tests/control/admin/system.rs @@ -1948,7 +1948,7 @@ async fn gateway_validates_chat_pii_redaction_system_config_locally_with_trusted } #[tokio::test] -async fn gateway_handles_admin_system_provider_priority_mode_locally_with_bearer_admin_session() { +async fn gateway_rejects_removed_admin_system_provider_priority_mode_with_bearer_admin_session() { let upstream_hits = Arc::new(Mutex::new(0usize)); let upstream_hits_clone = Arc::clone(&upstream_hits); let upstream = Router::new().route( @@ -1962,7 +1962,7 @@ async fn gateway_handles_admin_system_provider_priority_mode_locally_with_bearer }), ); - let (upstream_url, upstream_handle) = start_server(upstream).await; + let (_upstream_url, upstream_handle) = start_server(upstream).await; let state = AppState::new().expect("gateway should build"); let access_token = issue_test_admin_access_token(&state, "device-admin-config").await; let gateway = build_router_with_state(state); @@ -1978,10 +1978,7 @@ async fn gateway_handles_admin_system_provider_priority_mode_locally_with_bearer .await .expect("request should succeed"); - assert_eq!(response.status(), StatusCode::OK); - let payload: serde_json::Value = response.json().await.expect("json body should parse"); - assert_eq!(payload["key"], "provider_priority_mode"); - assert_eq!(payload["value"], "provider"); + assert_eq!(response.status(), StatusCode::NOT_FOUND); assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0); gateway_handle.abort(); diff --git a/apps/aether-gateway/src/tests/frontdoor/ai.rs b/apps/aether-gateway/src/tests/frontdoor/ai.rs index 13bcd07e4..93bcc450a 100644 --- a/apps/aether-gateway/src/tests/frontdoor/ai.rs +++ b/apps/aether-gateway/src/tests/frontdoor/ai.rs @@ -5,6 +5,7 @@ use super::{ InMemoryVideoTaskRepository, StoredAuthApiKeySnapshot, UpsertVideoTask, VideoTaskLookupKey, VideoTaskReadRepository, VideoTaskStatus, VideoTaskWriteRepository, DEVELOPMENT_ENCRYPTION_KEY, }; +use crate::data::GatewayDataState; use crate::image_capabilities::openai_image_gateway_max_generation_count; use crate::tests::{ any, build_router_with_state, build_state_with_execution_runtime_override, json, start_server, @@ -3477,7 +3478,10 @@ async fn gateway_does_not_locally_reject_image_model_name_on_chat_completions() let gateway = build_router_with_state( AppState::new() .expect("gateway should build") - .with_auth_api_key_data_reader_for_tests(auth_repository), + .with_data_state_for_tests( + GatewayDataState::with_auth_api_key_reader_for_tests(auth_repository) + .with_system_default_routing_group_for_tests(), + ), ); let (gateway_url, gateway_handle) = start_server(gateway).await; diff --git a/apps/aether-gateway/src/tests/frontdoor/internal.rs b/apps/aether-gateway/src/tests/frontdoor/internal.rs index c2fba96f9..21313518b 100644 --- a/apps/aether-gateway/src/tests/frontdoor/internal.rs +++ b/apps/aether-gateway/src/tests/frontdoor/internal.rs @@ -5,6 +5,7 @@ use super::{ InMemoryProviderCatalogReadRepository, InMemoryRequestCandidateRepository, DEVELOPMENT_ENCRYPTION_KEY, }; +use crate::data::GatewayDataState; use crate::tests::{ any, build_router, build_router_with_state, build_state_with_execution_runtime_override, json, start_server, strip_sse_keepalive_comments, AppState, Arc, Body, HeaderValue, Json, Mutex, @@ -1301,7 +1302,10 @@ async fn gateway_returns_internal_gateway_decision_sync_fallback_with_resolved_a let gateway = build_router_with_state( AppState::new() .expect("gateway should build") - .with_auth_api_key_data_reader_for_tests(auth_repository), + .with_data_state_for_tests( + GatewayDataState::with_auth_api_key_reader_for_tests(auth_repository) + .with_system_default_routing_group_for_tests(), + ), ); let (gateway_url, gateway_handle) = start_server(gateway).await; diff --git a/apps/aether-gateway/src/tests/proxy.rs b/apps/aether-gateway/src/tests/proxy.rs index 81c9e6dda..f4a798aba 100644 --- a/apps/aether-gateway/src/tests/proxy.rs +++ b/apps/aether-gateway/src/tests/proxy.rs @@ -18,6 +18,10 @@ use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadReposi use aether_data_contracts::repository::provider_catalog::{ StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider, }; +use aether_scheduler_core::{ + build_scheduler_affinity_cache_key_for_api_key_id_with_client_session_and_scope, + SchedulerAffinityScope, +}; use sha2::{Digest, Sha256}; fn hash_api_key(value: &str) -> String { @@ -26,6 +30,18 @@ fn hash_api_key(value: &str) -> String { format!("{:x}", hasher.finalize()) } +fn system_default_affinity_cache_key(api_key_id: &str, api_format: &str, model: &str) -> String { + let scope = SchedulerAffinityScope::new("system-default", Some(1)); + build_scheduler_affinity_cache_key_for_api_key_id_with_client_session_and_scope( + api_key_id, + api_format, + model, + None, + Some(&scope), + ) + .expect("system-default affinity cache key should build") +} + fn sample_auth_snapshot( api_key_id: &str, user_id: &str, @@ -626,6 +642,7 @@ async fn gateway_forwards_public_request_to_remote_tunnel_owner_before_fallback_ "development-key", ) .with_auth_api_key_reader(auth_repository) + .with_system_default_routing_group_for_tests() .with_system_config_values_for_tests(vec![( tunnel_attachment_key("node-owner"), serde_json::to_value(crate::tunnel::TunnelAttachmentRecord { @@ -648,7 +665,7 @@ async fn gateway_forwards_public_request_to_remote_tunnel_owner_before_fallback_ state.client = short_timeout_client.clone(); state.owner_forward_client = short_timeout_client; state.remember_scheduler_affinity_target( - "scheduler_affinity:api-key-affinity-1:openai:chat:gpt-4.1", + &system_default_affinity_cache_key("api-key-affinity-1", "openai:chat", "gpt-4.1"), crate::cache::SchedulerAffinityTarget { provider_id: "provider-owner".to_string(), endpoint_id: "endpoint-owner".to_string(), @@ -853,6 +870,7 @@ async fn gateway_aggregates_sync_sse_from_remote_tunnel_owner_before_returning_t "development-key", ) .with_auth_api_key_reader(auth_repository) + .with_system_default_routing_group_for_tests() .with_system_config_values_for_tests(vec![( tunnel_attachment_key("node-cli-owner"), serde_json::to_value(crate::tunnel::TunnelAttachmentRecord { @@ -869,7 +887,7 @@ async fn gateway_aggregates_sync_sse_from_remote_tunnel_owner_before_returning_t .with_data_state_for_tests(data_state) .with_tunnel_identity_for_tests("gateway-a", Some("http://gateway-a:8080")); state.remember_scheduler_affinity_target( - "scheduler_affinity:api-key-affinity-cli-1:openai:responses:gpt-5.4", + &system_default_affinity_cache_key("api-key-affinity-cli-1", "openai:responses", "gpt-5.4"), crate::cache::SchedulerAffinityTarget { provider_id: "provider-cli-owner".to_string(), endpoint_id: "endpoint-cli-owner".to_string(), @@ -1100,6 +1118,7 @@ async fn gateway_streamifies_sync_json_from_remote_tunnel_owner_before_returning "development-key", ) .with_auth_api_key_reader(auth_repository) + .with_system_default_routing_group_for_tests() .with_system_config_values_for_tests(vec![( tunnel_attachment_key("node-cli-owner"), serde_json::to_value(crate::tunnel::TunnelAttachmentRecord { @@ -1120,7 +1139,7 @@ async fn gateway_streamifies_sync_json_from_remote_tunnel_owner_before_returning .build() .expect("short shared client should build"); state.remember_scheduler_affinity_target( - "scheduler_affinity:api-key-affinity-cli-1:openai:responses:gpt-5.4", + &system_default_affinity_cache_key("api-key-affinity-cli-1", "openai:responses", "gpt-5.4"), crate::cache::SchedulerAffinityTarget { provider_id: "provider-cli-owner".to_string(), endpoint_id: "endpoint-cli-owner".to_string(), From dabaeb8dfaa46bf21fa6881059d55c685608a72e Mon Sep 17 00:00:00 2001 From: AAEE86 Date: Fri, 4 Sep 2026 10:50:23 +0800 Subject: [PATCH 05/13] ci: derive nightly image owner from repository --- .github/workflows/nightly.yml | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml index fa99fe576..5de9df839 100644 --- a/.github/workflows/nightly.yml +++ b/.github/workflows/nightly.yml @@ -25,7 +25,6 @@ env: CARGO_PROFILE_TEST_DEBUG: '0' CARGO_TERM_COLOR: always RUST_BACKTRACE: '1' - GHCR_IMAGE: ghcr.io/fawney19/aether jobs: source: @@ -36,6 +35,7 @@ jobs: sha: ${{ steps.snapshot.outputs.sha }} short_sha: ${{ steps.snapshot.outputs.short_sha }} date: ${{ steps.snapshot.outputs.date }} + ghcr_image: ${{ steps.snapshot.outputs.ghcr_image }} steps: - name: Require main branch id: snapshot @@ -49,9 +49,13 @@ jobs: fi sha="${GITHUB_SHA}" + # Docker 镜像仓库名必须全小写；GitHub owner 可能保留大写，先统一规范化。 + repository_owner="${GITHUB_REPOSITORY%%/*}" + repository_owner="${repository_owner,,}" echo "sha=${sha}" >> "${GITHUB_OUTPUT}" echo "short_sha=${sha:0:7}" >> "${GITHUB_OUTPUT}" echo "date=$(date -u +'%Y-%m-%d')" >> "${GITHUB_OUTPUT}" + echo "ghcr_image=ghcr.io/${repository_owner}/aether" >> "${GITHUB_OUTPUT}" echo "Building main at ${sha}." # Keep the scheduled backend coverage in one place so it cannot drift from PR CI. @@ -308,6 +312,8 @@ jobs: needs: [source, checks, build] if: ${{ needs.checks.result == 'success' && needs.build.result == 'success' }} runs-on: ubuntu-latest + env: + GHCR_IMAGE: ${{ needs.source.outputs.ghcr_image }} permissions: actions: read contents: read @@ -466,6 +472,7 @@ jobs: SOURCE_SHA: ${{ needs.source.outputs.sha }} SOURCE_SHORT_SHA: ${{ needs.source.outputs.short_sha }} RELEASE_DATE: ${{ needs.source.outputs.date }} + GHCR_IMAGE: ${{ needs.source.outputs.ghcr_image }} run: | set -euo pipefail From 1eb2d10decef56150ec51c7ad9c14d183ae7ef44 Mon Sep 17 00:00:00 2001 From: AAEE86 Date: Fri, 4 Sep 2026 15:24:10 +0800 Subject: [PATCH 06/13] fix(routing): filter providers by selected model --- .../routing/__tests__/providerQuery.spec.ts | 31 +++++++++++++++ .../RoutingPriorityPolicyEditor.vue | 28 ++++++++++++- .../features/routing/utils/providerQuery.ts | 39 +++++++++++++++++++ frontend/src/views/admin/RoutingProfiles.vue | 6 +++ 4 files changed, 102 insertions(+), 2 deletions(-) create mode 100644 frontend/src/features/routing/__tests__/providerQuery.spec.ts create mode 100644 frontend/src/features/routing/utils/providerQuery.ts diff --git a/frontend/src/features/routing/__tests__/providerQuery.spec.ts b/frontend/src/features/routing/__tests__/providerQuery.spec.ts new file mode 100644 index 000000000..651fd59b7 --- /dev/null +++ b/frontend/src/features/routing/__tests__/providerQuery.spec.ts @@ -0,0 +1,31 @@ +import { describe, expect, it } from 'vitest' + +import { DEFAULT_ROUTING_POLICY_MODEL } from '../utils/routingPolicy' +import { buildRoutingProviderSummaryQuery } from '../utils/providerQuery' + +describe('routing provider query', () => { + it('filters provider sorting by the selected GlobalModel ID', () => { + expect(buildRoutingProviderSummaryQuery( + 'claude-fable-5-1', + 'global-claude-fable-5-1', + 'provider', + )).toEqual({ + page: 1, + page_size: 9999, + model_id: 'global-claude-fable-5-1', + }) + }) + + it('keeps unified sorting and key sorting unfiltered', () => { + const expected = { page: 1, page_size: 9999 } + + expect(buildRoutingProviderSummaryQuery(DEFAULT_ROUTING_POLICY_MODEL, undefined, 'provider')) + .toEqual(expected) + expect(buildRoutingProviderSummaryQuery('claude-fable-5-1', 'global-claude-fable-5-1', 'global_key')) + .toEqual(expected) + }) + + it('does not fall back to all providers before a model ID is available', () => { + expect(buildRoutingProviderSummaryQuery('claude-fable-5-1', undefined, 'provider')).toBeNull() + }) +}) diff --git a/frontend/src/features/routing/components/RoutingPriorityPolicyEditor.vue b/frontend/src/features/routing/components/RoutingPriorityPolicyEditor.vue index e27b80c27..a28d65b68 100644 --- a/frontend/src/features/routing/components/RoutingPriorityPolicyEditor.vue +++ b/frontend/src/features/routing/components/RoutingPriorityPolicyEditor.vue @@ -324,6 +324,7 @@ import { type RoutingPriorityMode, type RoutingSchedulingMode, } from '../utils/routingPolicy' +import { buildRoutingProviderSummaryQuery } from '../utils/providerQuery' interface ProviderPriorityRow { id: string @@ -367,6 +368,7 @@ interface GlobalKeySource { const props = defineProps<{ config: RoutingGroupConfig model?: string + modelId?: string priorityMode?: RoutingPriorityMode schedulingMode?: RoutingSchedulingMode showPriorityMode?: boolean @@ -398,6 +400,7 @@ const draggedKeyId = ref(null) const dragOverKeyId = ref(null) const providerMultiSelectEnabled = ref(false) const selectedProviderIds = ref>(new Set()) +let providerLoadRequestId = 0 const config = computed(() => normalizeRoutingGroupConfig(props.config)) const targetModel = computed(() => props.model?.trim() || DEFAULT_ROUTING_POLICY_MODEL) @@ -505,6 +508,12 @@ watch(effectivePriorityMode, mode => { providerMultiSelectEnabled.value = false selectedProviderIds.value = new Set() } + void loadProviders() +}) + +// 父组件异步解析全局模型 ID 后，重新加载对应模型的提供商列表。 +watch([targetModel, () => props.modelId], () => { + void loadProviders() }) watch(providerRows, rows => { @@ -561,16 +570,31 @@ function updateSchedulingMode(mode: RoutingSchedulingMode): void { } async function loadProviders(): Promise { + const requestId = ++providerLoadRequestId loadingProviders.value = true loadError.value = null try { - const response = await getProvidersSummary({ page: 1, page_size: 9999 }) + const query = buildRoutingProviderSummaryQuery( + targetModel.value, + props.modelId, + effectivePriorityMode.value, + ) + if (!query) { + providers.value = [] + return + } + + const response = await getProvidersSummary(query) + if (requestId !== providerLoadRequestId) return providers.value = response.items } catch (err) { + if (requestId !== providerLoadRequestId) return loadError.value = parseApiError(err, '加载 Provider 失败') providers.value = [] } finally { - loadingProviders.value = false + if (requestId === providerLoadRequestId) { + loadingProviders.value = false + } } } diff --git a/frontend/src/features/routing/utils/providerQuery.ts b/frontend/src/features/routing/utils/providerQuery.ts new file mode 100644 index 000000000..43fe678b8 --- /dev/null +++ b/frontend/src/features/routing/utils/providerQuery.ts @@ -0,0 +1,39 @@ +import type { ProviderSummaryQuery } from '@/api/endpoints/providers' + +import { + DEFAULT_ROUTING_POLICY_MODEL, + type RoutingPriorityMode, +} from './routingPolicy' + +const PROVIDER_SUMMARY_PAGE_SIZE = 9999 + +/** + * 构造路由排序编辑器的提供商查询参数。 + * 按模型配置时必须使用 GlobalModel ID 过滤，避免把模型名称误当成 ID；统一调度和 + * Key 排序仍需要完整的提供商集合。 + */ +export function buildRoutingProviderSummaryQuery( + model: string | undefined, + modelId: string | undefined, + priorityMode: RoutingPriorityMode, +): ProviderSummaryQuery | null { + const query: ProviderSummaryQuery = { + page: 1, + page_size: PROVIDER_SUMMARY_PAGE_SIZE, + } + const isModelScoped = (model?.trim() || DEFAULT_ROUTING_POLICY_MODEL) + !== DEFAULT_ROUTING_POLICY_MODEL + + if (!isModelScoped || priorityMode !== 'provider') { + return query + } + + const normalizedModelId = modelId?.trim() + // 模型 ID 尚未由父组件解析出来时，不能回退到全量列表，否则会短暂显示错误的提供商。 + if (!normalizedModelId) return null + + return { + ...query, + model_id: normalizedModelId, + } +} diff --git a/frontend/src/views/admin/RoutingProfiles.vue b/frontend/src/views/admin/RoutingProfiles.vue index 5d0794bed..c9d7fdfdf 100644 --- a/frontend/src/views/admin/RoutingProfiles.vue +++ b/frontend/src/views/admin/RoutingProfiles.vue @@ -745,6 +745,7 @@ item.name.trim() === normalizedName)?.id +} + function replaceGroup(group: RoutingGroupRecord, select = true): void { const normalized = normalizeRecord(group) const index = groups.value.findIndex(item => item.id === normalized.id) From 66d6c17d2d80ddebf05c682f06793818f3b8158d Mon Sep 17 00:00:00 2001 From: ZheFox <77232781+zhefox@users.noreply.github.com> Date: Fri, 4 Sep 2026 16:36:40 +0800 Subject: [PATCH 07/13] fix antigravity reasoning streaming and terminal errors --- .../execution_runtime/stream/commit_policy.rs | 14 +- .../src/execution_runtime/stream/execution.rs | 33 +++- .../formats/gemini/generate_content/stream.rs | 131 ++++++++++++++ .../shared/stream_core/format_matrix.rs | 162 +++++++++++++++++- 4 files changed, 319 insertions(+), 21 deletions(-) diff --git a/apps/aether-gateway/src/execution_runtime/stream/commit_policy.rs b/apps/aether-gateway/src/execution_runtime/stream/commit_policy.rs index 3cbf9f641..9cfd0cae1 100644 --- a/apps/aether-gateway/src/execution_runtime/stream/commit_policy.rs +++ b/apps/aether-gateway/src/execution_runtime/stream/commit_policy.rs @@ -455,7 +455,10 @@ fn gemini_part_is_client_semantic(part: &Value) -> bool { return true; } if part.get("thought").and_then(Value::as_bool) == Some(true) { - return false; + return part + .get("text") + .and_then(Value::as_str) + .is_some_and(|text| !text.is_empty()); } if part.keys().all(|key| key == "thoughtSignature") { return false; @@ -584,17 +587,12 @@ mod tests { } #[test] - fn gemini_gate_waits_through_thought_and_commits_on_text() { + fn gemini_gate_commits_on_first_nonempty_thought() { let mut gate = StreamCommitGate::new(gemini_policy()); let thought = b"data: {\"response\":{\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"thought\":true,\"text\":\"checking\"}]}}]}}\n\n"; - let text = b"data: {\"response\":{\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"text\":\"answer\"}]}}]}}\n\n"; assert_eq!( gate.observe_provider_bytes(thought), - StreamPrecommitObservation::Pending - ); - assert_eq!( - gate.observe_provider_bytes(text), StreamPrecommitObservation::Commit ); assert_eq!(gate.state(), StreamCommitState::Committed); @@ -615,7 +613,7 @@ mod tests { #[test] fn gemini_gate_rejects_malformed_function_call_before_commit() { let mut gate = StreamCommitGate::new(gemini_policy()); - let thought = b"data: {\"response\":{\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"thought\":true,\"text\":\"calling\"}]}}]}}\n\n"; + let thought = b"data: {\"response\":{\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"thoughtSignature\":\"signature\",\"text\":\"\"}]}}]}}\n\n"; let malformed = b"data: {\"response\":{\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"thoughtSignature\":\"signature\",\"text\":\"\"}]},\"finishReason\":\"MALFORMED_FUNCTION_CALL\",\"finishMessage\":\"Malformed function call: Function call is empty - no input to parse.\"}]}}\n\n"; assert_eq!( diff --git a/apps/aether-gateway/src/execution_runtime/stream/execution.rs b/apps/aether-gateway/src/execution_runtime/stream/execution.rs index 800c09d12..536326765 100644 --- a/apps/aether-gateway/src/execution_runtime/stream/execution.rs +++ b/apps/aether-gateway/src/execution_runtime/stream/execution.rs @@ -10561,7 +10561,7 @@ mod tests { } #[tokio::test] - async fn malformed_antigravity_function_call_retries_before_stream_commit() { + async fn malformed_antigravity_function_call_streams_thought_then_fails_in_band() { let request_id = "req-antigravity-malformed-function-call"; let plan = antigravity_gemini_stream_plan(request_id); let provider_catalog = provider_catalog_for_plan( @@ -10640,10 +10640,35 @@ mod tests { None, ) .await - .expect("malformed Antigravity stream should resolve through failover"); + .expect("malformed Antigravity stream should return a client stream") + .expect("the first reasoning delta should commit the selected candidate"); - assert!(response.is_none()); - assert_eq!(retry_scope, AiAttemptRetryScope::Candidate); + assert_eq!(response.status(), StatusCode::OK); + let body = to_bytes(response.into_body(), usize::MAX) + .await + .expect("response body should read"); + let body = String::from_utf8(body.to_vec()).expect("response body should be utf8"); + assert!( + body.contains("event: response.reasoning_summary_text.delta\n"), + "{body}" + ); + assert!( + body.contains("\"delta\":\"Validating the document.\""), + "{body}" + ); + assert!(body.contains("event: response.failed\n"), "{body}"); + assert!( + body.contains("\"code\":\"MALFORMED_FUNCTION_CALL\""), + "{body}" + ); + assert!( + body.contains( + "\"message\":\"Malformed function call: Function call is empty - no input to parse.\"" + ), + "{body}" + ); + assert!(!body.contains("unsupported_finish_reason"), "{body}"); + assert_eq!(retry_scope, AiAttemptRetryScope::Provider); } fn tunnel_proxy_snapshot(base_url: String) -> aether_contracts::ProxySnapshot { diff --git a/crates/aether-ai/formats/src/formats/gemini/generate_content/stream.rs b/crates/aether-ai/formats/src/formats/gemini/generate_content/stream.rs index 80d04e4af..5f5875ab2 100644 --- a/crates/aether-ai/formats/src/formats/gemini/generate_content/stream.rs +++ b/crates/aether-ai/formats/src/formats/gemini/generate_content/stream.rs @@ -104,10 +104,25 @@ impl GeminiProviderState { let Some(candidate_object) = candidate.as_object() else { continue; }; + let (response_id, response_model) = self.identity(report_context); + let terminal_error = gemini_stream_terminal_error_payload( + candidate_object, + response_id.as_str(), + response_model.as_str(), + event_object.get("usageMetadata"), + ); let Some(content) = candidate_object.get("content").and_then(Value::as_object) else { + if let Some(payload) = terminal_error { + out.push(self.unknown_frame(report_context, payload)); + self.finished = true; + } continue; }; let Some(parts) = content.get("parts").and_then(Value::as_array) else { + if let Some(payload) = terminal_error { + out.push(self.unknown_frame(report_context, payload)); + self.finished = true; + } continue; }; if !parts.is_empty() { @@ -306,6 +321,11 @@ impl GeminiProviderState { }); } } + if let Some(payload) = terminal_error { + out.push(self.unknown_frame(report_context, payload)); + self.finished = true; + continue; + } if let Some(finish_reason) = candidate_object.get("finishReason").and_then(Value::as_str) { @@ -350,6 +370,60 @@ impl GeminiProviderState { } } +fn gemini_stream_terminal_error_payload( + candidate: &Map, + response_id: &str, + model: &str, + usage_metadata: Option<&Value>, +) -> Option { + let finish_reason = candidate + .get("finishReason") + .or_else(|| candidate.get("finish_reason")) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| { + matches!( + *value, + "MALFORMED_FUNCTION_CALL" + | "UNEXPECTED_TOOL_CALL" + | "TOO_MANY_TOOL_CALLS" + | "MISSING_THOUGHT_SIGNATURE" + | "MALFORMED_RESPONSE" + ) + })?; + let message = candidate + .get("finishMessage") + .or_else(|| candidate.get("finish_message")) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) + .unwrap_or_else(|| format!("Gemini stream ended with {finish_reason}")); + + let mut response = json!({ + "id": response_id, + "object": "response", + "model": model, + "status": "failed", + "error": { + "type": "upstream_gemini_finish_error", + "code": finish_reason, + "message": message, + "upstream_status": 200 + } + }); + if let Some(usage) = canonical_usage_from_gemini_usage(usage_metadata) + .map(|usage| openai_responses_usage_from_usage(&usage)) + { + response["usage"] = usage; + } + + Some(json!({ + "type": "response.failed", + "response": response + })) +} + fn map_gemini_stream_finish_reason(value: &str) -> Option<&str> { match value { "STOP" => Some("stop"), @@ -1011,6 +1085,63 @@ mod tests { ))); } + #[test] + fn gemini_provider_state_emits_terminal_error_for_malformed_function_call() { + let mut state = GeminiProviderState::default(); + let report_context = json!({}); + let frames = state + .push_line( + &report_context, + data_line(json!({ + "response": { + "responseId": "resp_malformed_tool_call", + "modelVersion": "gemini-3.7-flash-tiered", + "candidates": [{ + "index": 0, + "content": { + "role": "model", + "parts": [{ + "text": "", + "thoughtSignature": "opaque-thought-signature" + }] + }, + "finishReason": "MALFORMED_FUNCTION_CALL", + "finishMessage": "Malformed function call: Function call is empty - no input to parse." + }], + "usageMetadata": { + "promptTokenCount": 206744, + "cachedContentTokenCount": 203947, + "thoughtsTokenCount": 1130, + "totalTokenCount": 207874 + } + } + })), + ) + .expect("malformed function call terminal should parse"); + + assert!(frames.iter().any(|frame| matches!( + &frame.event, + CanonicalStreamEvent::UnknownEvent(payload) + if payload["type"] == "response.failed" + && payload["response"]["status"] == "failed" + && payload["response"]["id"] == "resp_malformed_tool_call" + && payload["response"]["model"] == "gemini-3.7-flash-tiered" + && payload["response"]["error"]["code"] == "MALFORMED_FUNCTION_CALL" + && payload["response"]["error"]["message"] + == "Malformed function call: Function call is empty - no input to parse." + && payload["response"]["usage"]["input_tokens"] == 206744 + && payload["response"]["usage"]["output_tokens"] == 1130 + && payload["response"]["usage"]["total_tokens"] == 207874 + ))); + assert!(!frames + .iter() + .any(|frame| matches!(frame.event, CanonicalStreamEvent::Finish { .. }))); + assert!(state + .finish(&report_context) + .expect("finished error stream should not synthesize success") + .is_empty()); + } + #[test] fn gemini_provider_state_parses_function_response_as_tool_result() { let mut state = GeminiProviderState::default(); diff --git a/crates/aether-ai/formats/src/formats/shared/stream_core/format_matrix.rs b/crates/aether-ai/formats/src/formats/shared/stream_core/format_matrix.rs index 4f5e86d6f..6983f7e94 100644 --- a/crates/aether-ai/formats/src/formats/shared/stream_core/format_matrix.rs +++ b/crates/aether-ai/formats/src/formats/shared/stream_core/format_matrix.rs @@ -17,8 +17,9 @@ use crate::formats::shared::error_body::{ }; use crate::formats::shared::sse::encode_json_sse; use crate::formats::shared::stream_core::common::{ - decode_json_data_line, openai_stream_terminal_error_body, openai_stream_terminal_error_message, - unsupported_stream_event_message, CanonicalStreamEvent, CanonicalStreamFrame, CanonicalUsage, + canonical_usage_from_openai_usage, decode_json_data_line, openai_stream_terminal_error_body, + openai_stream_terminal_error_message, unsupported_stream_event_message, CanonicalStreamEvent, + CanonicalStreamFrame, CanonicalUsage, }; use crate::formats::shared::AiSurfaceFinalizeError; @@ -134,7 +135,11 @@ impl StreamingStandardFormatMatrix { } if let CanonicalStreamEvent::UnknownEvent(payload) = &frame.event { self.terminated = true; - out.extend(client.emit_unknown_event(payload)?); + if openai_stream_terminal_error_body(payload).is_some() { + out.extend(client.emit_terminal_error_frame(frame)?); + } else { + out.extend(client.emit_unknown_event(payload)?); + } break; } if let CanonicalStreamEvent::OpenAiResponsesOutputItem { raw_event, .. } = &frame.event @@ -368,6 +373,10 @@ impl StreamingStandardTerminalObserver { summary.observed_finish = true; summary.finish_reason = Some("error".to_string()); summary.parser_error = openai_stream_terminal_error_message(&payload); + summary.standardized_usage = payload + .pointer("/response/usage") + .and_then(|usage| canonical_usage_from_openai_usage(Some(usage))) + .map(standardized_usage_from_canonical); } CanonicalStreamEvent::UnknownEvent(_) => { summary.unknown_event_count = summary.unknown_event_count.saturating_add(1); @@ -613,6 +622,44 @@ impl ClientStreamEmitter { self.emit_error(error_body) } + fn emit_terminal_error_frame( + &mut self, + frame: CanonicalStreamFrame, + ) -> Result, AiSurfaceFinalizeError> { + if matches!( + self, + ClientStreamEmitter::OpenAIChat(_) | ClientStreamEmitter::OpenAIResponses(_) + ) { + return self.emit(frame); + } + let CanonicalStreamEvent::UnknownEvent(payload) = frame.event else { + return self.emit(frame); + }; + let Some(source_error_body) = openai_stream_terminal_error_body(&payload) else { + return self.emit_unknown_event(&payload); + }; + let Some(error) = source_error_body.get("error") else { + return self.emit_unknown_event(&payload); + }; + let message = error + .get("message") + .and_then(Value::as_str) + .unwrap_or("Upstream stream ended with an error"); + let code = error.get("code").and_then(|value| match value { + Value::String(value) => Some(value.as_str()), + _ => None, + }); + let Some(error_body) = build_core_error_body_for_client_format( + self.api_format(), + message, + code, + LocalCoreSyncErrorKind::ServerError, + ) else { + return Ok(Vec::new()); + }; + self.emit_error(error_body) + } + fn emit_unsupported_finish_reason( &mut self, finish_reason: &str, @@ -811,7 +858,13 @@ mod tests { }, "finishReason": "MALFORMED_FUNCTION_CALL", "finishMessage": "Malformed function call: Function call is empty - no input to parse." - }] + }], + "usageMetadata": { + "promptTokenCount": 206744, + "cachedContentTokenCount": 203947, + "thoughtsTokenCount": 1130, + "totalTokenCount": 207874 + } }, "responseId": "resp_malformed_tool_call" })), @@ -824,14 +877,105 @@ mod tests { .expect("Gemini terminal frame should produce a summary"); assert!(summary.observed_finish); - assert_eq!( - summary.finish_reason.as_deref(), - Some("MALFORMED_FUNCTION_CALL") - ); + assert_eq!(summary.finish_reason.as_deref(), Some("error")); assert_eq!( summary.parser_error.as_deref(), - Some("unsupported provider stream finish reason: MALFORMED_FUNCTION_CALL") + Some("Malformed function call: Function call is empty - no input to parse.") ); + let usage = summary + .standardized_usage + .expect("failed Gemini terminal should preserve usage"); + assert_eq!(usage.input_tokens, 206744); + assert_eq!(usage.output_tokens, 1130); + assert_eq!(usage.cache_read_tokens, 203947); + } + + #[test] + fn streams_gemini_thought_text_to_openai_responses_immediately() { + let context = report_context("gemini:generate_content", "openai:responses"); + let mut matrix = StreamingStandardFormatMatrix::default(); + let output = matrix + .transform_line( + &context, + data_line(json!({ + "response": { + "responseId": "resp_reasoning_123", + "modelVersion": "gemini-3.7-flash-tiered", + "candidates": [{ + "index": 0, + "content": { + "role": "model", + "parts": [{"thought": true, "text": "checking"}] + } + }] + } + })), + ) + .expect("first Gemini thought chunk should transform"); + let sse = String::from_utf8(output).expect("reasoning SSE should be utf8"); + + assert!( + sse.contains("event: response.reasoning_summary_text.delta\n"), + "{sse}" + ); + assert!(sse.contains("\"delta\":\"checking\""), "{sse}"); + } + + #[test] + fn transforms_malformed_gemini_function_call_to_responses_failed() { + let context = report_context("gemini:generate_content", "openai:responses"); + let mut matrix = StreamingStandardFormatMatrix::default(); + let output = matrix + .transform_line( + &context, + data_line(json!({ + "response": { + "responseId": "resp_malformed_tool_call", + "modelVersion": "gemini-3.7-flash-tiered", + "candidates": [{ + "index": 0, + "content": { + "role": "model", + "parts": [{ + "text": "", + "thoughtSignature": "opaque-thought-signature" + }] + }, + "finishReason": "MALFORMED_FUNCTION_CALL", + "finishMessage": "Malformed function call: Function call is empty - no input to parse." + }], + "usageMetadata": { + "promptTokenCount": 206744, + "cachedContentTokenCount": 203947, + "thoughtsTokenCount": 1130, + "totalTokenCount": 207874 + } + } + })), + ) + .expect("malformed Gemini terminal should transform to a stream error"); + let sse = String::from_utf8(output).expect("failed response SSE should be utf8"); + + assert!(sse.contains("event: response.failed\n"), "{sse}"); + assert!(sse.contains("\"type\":\"response.failed\""), "{sse}"); + assert!( + sse.contains("\"code\":\"MALFORMED_FUNCTION_CALL\""), + "{sse}" + ); + assert!( + sse.contains( + "\"message\":\"Malformed function call: Function call is empty - no input to parse.\"" + ), + "{sse}" + ); + assert!(sse.contains("\"input_tokens\":206744"), "{sse}"); + assert!(sse.contains("\"output_tokens\":1130"), "{sse}"); + assert!(sse.contains("\"cached_tokens\":203947"), "{sse}"); + assert!(!sse.contains("unsupported_finish_reason"), "{sse}"); + assert!(matrix + .finish(&context) + .expect("failed matrix should stay terminated") + .is_empty()); } #[test] From 14744abd573cbf3cacaaf5c8fda41f9ebbe46db8 Mon Sep 17 00:00:00 2001 From: ZheFox <77232781+zhefox@users.noreply.github.com> Date: Fri, 4 Sep 2026 16:54:42 +0800 Subject: [PATCH 08/13] fix responses replay across Gemini and Codex --- .../src/formats/openai/request_contract.rs | 11 ++++ .../src/formats/openai/responses/mod.rs | 51 +++++++++++++++++-- 2 files changed, 59 insertions(+), 3 deletions(-) diff --git a/crates/aether-ai/formats/src/formats/openai/request_contract.rs b/crates/aether-ai/formats/src/formats/openai/request_contract.rs index 7eb7107bd..b145d7c69 100644 --- a/crates/aether-ai/formats/src/formats/openai/request_contract.rs +++ b/crates/aether-ai/formats/src/formats/openai/request_contract.rs @@ -381,10 +381,21 @@ mod tests { #[test] fn finalization_strips_non_replayable_responses_reasoning_history() { + let gemini_carrier = + crate::formats::openai::responses::encode_gemini_tool_signature_carrier( + "opaque-gemini-thought-signature", + ) + .expect("Gemini signature carrier"); let mut body = json!({ "model": "gpt-5.4", "input": [ {"type": "reasoning", "id": "rs_provider_123", "summary": []}, + { + "type": "reasoning", + "id": "rs_aether_55070860f6d45c6b8f6fa11efd9dff8a", + "summary": [], + "encrypted_content": gemini_carrier + }, { "type": "reasoning", "id": "item_72d3bd8d367d01977ace23f1", diff --git a/crates/aether-ai/formats/src/formats/openai/responses/mod.rs b/crates/aether-ai/formats/src/formats/openai/responses/mod.rs index 9a9328ff8..36053f67e 100644 --- a/crates/aether-ai/formats/src/formats/openai/responses/mod.rs +++ b/crates/aether-ai/formats/src/formats/openai/responses/mod.rs @@ -167,9 +167,10 @@ pub fn normalize_openai_responses_message_item_ids(body: &mut Value) -> usize { /// 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 -/// prefix. Foreign IDs (for example `item_...`) are therefore removed. Aether-synthesized -/// reasoning summaries are also removed unless they carry encrypted reasoning state that can be -/// replayed statelessly. +/// prefix. Foreign IDs (for example `item_...`) are therefore removed. Aether's Gemini signature +/// carriers are also removed: they are intentionally transported through the Responses +/// `encrypted_content` field so they can be restored on a later Gemini tool turn, but they are not +/// OpenAI ciphertext and must never be replayed to an OpenAI/Codex backend. pub fn strip_incompatible_openai_responses_reasoning_items( body: &mut Value, provider_api_format: &str, @@ -221,6 +222,13 @@ fn openai_responses_reasoning_item_is_replayable( if object.get("type").and_then(Value::as_str) != Some("reasoning") { return true; } + if object + .get("encrypted_content") + .and_then(Value::as_str) + .is_some_and(|value| value.starts_with(GEMINI_TOOL_SIGNATURE_CARRIER_PREFIX)) + { + return false; + } if policy == OpenAiResponsesReasoningReplayPolicy::DeepSeekOpaque && deepseek_opaque_reasoning_item_is_replayable(object) { @@ -475,6 +483,43 @@ mod tests { assert_eq!(input[2]["id"], "item_message_123"); } + #[test] + fn strips_gemini_signature_carriers_before_openai_replay() { + let gemini_item_id = openai_responses_synthetic_reasoning_item_id("resp_gemini", 0); + let openai_item_id = openai_responses_synthetic_reasoning_item_id("resp_openai", 0); + let carrier = encode_gemini_tool_signature_carrier_with_direction( + "opaque-gemini-thought-signature", + GeminiToolSignatureCarrierDirection::Next, + ) + .expect("Gemini signature carrier"); + let mut body = json!({ + "input": [ + { + "type": "reasoning", + "id": gemini_item_id, + "summary": [], + "encrypted_content": carrier + }, + { + "type": "reasoning", + "id": openai_item_id, + "summary": [], + "encrypted_content": "provider-encrypted-state" + }, + {"type": "reasoning", "id": "rs_provider_123", "summary": []} + ] + }); + + assert_eq!( + strip_incompatible_openai_responses_reasoning_items(&mut body, "openai:responses"), + 1 + ); + let input = body["input"].as_array().expect("input array"); + assert_eq!(input.len(), 2); + assert_eq!(input[0]["encrypted_content"], "provider-encrypted-state"); + assert_eq!(input[1]["id"], "rs_provider_123"); + } + #[test] fn reasoning_item_sanitizer_is_scoped_to_responses_targets() { let mut body = json!({ From 9282cce1d643f8dbbf04d679b0fbbfb35066d920 Mon Sep 17 00:00:00 2001 From: stabey <36232531+stabey@users.noreply.github.com> Date: Wed, 2 Sep 2026 02:07:26 +0800 Subject: [PATCH 09/13] fix(gateway): settle stream attempts dropped before first byte A local stream attempt writes its `usage` row and its `request_candidates` slot as `pending` in `execute_execution_runtime_stream_inner`, then awaits the provider's response headers. Everything after that point runs inside the downstream request future, so a client disconnect drops it: the dispatch `.await` never resumes and nothing settles either row. The stream finalizer that already covers this only exists once upstream headers have arrived, so the pre-first-byte window has no owner at all. Both rows stay `pending` until the maintenance sweeper rewrites them as a 504 timeout ten minutes later, losing the real outcome, the real latency, and the 499. `AttemptCancellationGuard` takes that window. It is created disarmed, so an attempt dropped before it owns any row does not grow a settlement row it never had; it is armed as soon as the attempt owns its non-terminal rows, and the stream wrappers disarm it the moment the attempt returns, from where settlement belongs to the transport. On a cancelling drop it settles the candidate slot through the same snapshot writer the `pending` write above it uses, and the usage row through a terminal `Cancelled` event. The guard outlives the request future, so what it captures is retained for the whole attempt. It therefore holds no request body: the plan carries the provider request body and the report context carries the client request body, and keeping both would double the request-body residency of every in-flight stream attempt to serve a path that almost never runs. Simply omitting them is not safe either, because a terminal write is body-capture-authoritative: with both absent the seed carries the typed `none` marker, which clears the stored capture rather than leaving it alone. `build_usage_event_data_seed_describing_request_bodies` is the third option -- it derives every capture state, body reference, request type and derived request fact from the real plan and report context, and leaves out only the two body values -- so the guard's snapshot is small and its terminal write preserves the capture the `pending` write recorded. The stream candidate first-byte watchdog also drops the attempt future, but it settles the attempt itself through `build_transport_error_stop_response`. It now marks the attempt abandoned before returning so the guard stands down instead of racing a 499 against the watchdog's 504. Co-Authored-By: Claude Opus 5 --- .../execution_runtime/attempt_cancellation.rs | 568 ++++++++++++++++++ .../src/execution_runtime/mod.rs | 1 + .../src/execution_runtime/stream/execution.rs | 67 ++- .../execution_runtime/transport_failure.rs | 19 + .../src/executor/candidate_loop.rs | 4 + .../src/tests/ai_execute/lifecycle.rs | 109 ++++ crates/aether-usage/runtime/src/lib.rs | 3 +- crates/aether-usage/runtime/src/write.rs | 198 +++++- 8 files changed, 943 insertions(+), 26 deletions(-) create mode 100644 apps/aether-gateway/src/execution_runtime/attempt_cancellation.rs diff --git a/apps/aether-gateway/src/execution_runtime/attempt_cancellation.rs b/apps/aether-gateway/src/execution_runtime/attempt_cancellation.rs new file mode 100644 index 000000000..f538810e6 --- /dev/null +++ b/apps/aether-gateway/src/execution_runtime/attempt_cancellation.rs @@ -0,0 +1,568 @@ +//! Terminal settlement for a local stream attempt whose future is dropped +//! mid-flight. +//! +//! A local stream attempt writes its `usage` row and its `request_candidates` +//! slot as `pending` before it dispatches to the provider, then keeps running +//! inside the downstream request future. When the client disconnects, axum drops +//! that future: the remaining `.await`s never resume and nothing settles either +//! row. They stay `pending` until the maintenance sweeper rewrites them as a 504 +//! timeout roughly ten minutes later, which loses the real outcome and the real +//! latency. +//! +//! The stream transport therefore keeps a guard alive across the window between +//! the `pending` write and terminal settlement, and settles the attempt from +//! `Drop` when that window is left by cancellation instead of by a terminal +//! state. + +use std::sync::Arc; +use std::time::Instant; + +use aether_contracts::ExecutionPlan; +use aether_data_contracts::repository::candidates::RequestCandidateStatus; +use aether_scheduler_core::SchedulerRequestCandidateStatusUpdate; +use aether_usage_runtime::{ + build_usage_event_data_seed_describing_request_bodies, UsageEvent, UsageEventData, + UsageEventType, +}; +use serde_json::{json, Value}; +use tracing::warn; + +use crate::clock::current_unix_ms as current_request_candidate_unix_ms; +use crate::execution_runtime::attempt_lifecycle::CLIENT_CANCELLED_STATUS_CODE; +use crate::execution_runtime::transport_failure::StreamCandidateWatchdogProgress; +use crate::log_ids::short_request_id; +use crate::request_candidate_runtime::{ + record_local_request_candidate_status_snapshot, LocalRequestCandidateStatusSnapshot, +}; +use crate::request_diagnostics::{ + attach_request_diagnostics_to_report_context, current_request_diagnostics, RequestDiagnostics, +}; +use crate::AppState; + +fn elapsed_ms_since(started_at: Instant) -> u64 { + started_at.elapsed().as_millis().min(u128::from(u64::MAX)) as u64 +} + +/// The facts the guard needs to settle the attempt it is watching. +/// +/// This is held for the whole attempt, so it is deliberately free of request +/// bodies. A request body can be megabytes, and holding one per in-flight +/// attempt would cost far more than the row it settles: the usage seed is built +/// with [`build_usage_event_data_seed_describing_request_bodies`], which derives +/// every capture state, body reference and derived request fact from the real +/// plan and report context but keeps neither body. The terminal write it +/// produces therefore preserves the capture the `pending` write recorded instead +/// of clearing it. +struct ArmedAttempt { + request_id: String, + candidate_id: Option, + candidate: Option, + // Boxed: the guard lives inside the stream request future, which is already + // very large, and `UsageEventData` is a wide struct. + usage_seed: Option>, + request_diagnostics: Option>, + candidate_started_unix_ms: u64, + candidate_started_at: Instant, +} + +/// Settles an attempt as cancelled when its future is dropped before the +/// transport reaches a terminal state. +/// +/// The guard is created disarmed and stays inert until [`Self::arm`] is called, +/// so an attempt that is dropped before it owns any `pending` row does not grow +/// a settlement row it never had. The owner disarms it as soon as the attempt +/// completes, whichever way it completes: from that point terminal settlement +/// belongs to the transport (for streams, to the stream finalizer that lives in +/// the response body), and the guard must not write a second terminal state. +/// +/// A stream candidate also runs under a first-byte watchdog that drops the +/// attempt future when it gives up. That drop is not a client disconnect and the +/// watchdog settles the attempt itself, so the guard stands down for it. +pub(crate) struct AttemptCancellationGuard { + state: AppState, + error_type: &'static str, + error_message: &'static str, + watchdog: Option>, + armed: Option, +} + +impl AttemptCancellationGuard { + pub(crate) fn disarmed( + state: &AppState, + error_type: &'static str, + error_message: &'static str, + ) -> Self { + Self { + state: state.clone(), + error_type, + error_message, + watchdog: StreamCandidateWatchdogProgress::current(), + armed: None, + } + } + + /// Takes ownership of the attempt's settlement until it is disarmed. + pub(crate) fn arm( + &mut self, + plan: &ExecutionPlan, + report_context: Option<&Value>, + candidate: Option<&LocalRequestCandidateStatusSnapshot>, + candidate_started_unix_ms: u64, + candidate_started_at: Instant, + ) { + let usage_seed = self.state.usage_runtime.is_enabled().then(|| { + Box::new(build_usage_event_data_seed_describing_request_bodies( + plan, + report_context, + )) + }); + self.armed = Some(ArmedAttempt { + request_id: plan.request_id.clone(), + candidate_id: plan.candidate_id.clone(), + candidate: candidate.cloned(), + usage_seed, + request_diagnostics: current_request_diagnostics(), + candidate_started_unix_ms, + candidate_started_at, + }); + } + + pub(crate) fn disarm(&mut self) { + self.armed = None; + } +} + +/// Writes the candidate terminal row and the terminal usage event for an attempt +/// that never reached its own terminal path. +async fn settle_cancelled_attempt( + state: AppState, + armed: ArmedAttempt, + error_type: &'static str, + error_message: &'static str, +) { + let ArmedAttempt { + request_id, + candidate_id: _, + candidate, + usage_seed, + request_diagnostics, + candidate_started_unix_ms, + candidate_started_at, + } = armed; + let terminal_unix_ms = current_request_candidate_unix_ms(); + let latency_ms = elapsed_ms_since(candidate_started_at); + + if let Some(candidate) = candidate.as_ref() { + record_local_request_candidate_status_snapshot( + &state, + candidate, + SchedulerRequestCandidateStatusUpdate { + status: RequestCandidateStatus::Cancelled, + status_code: Some(CLIENT_CANCELLED_STATUS_CODE), + error_type: Some(error_type.to_string()), + error_message: Some(error_message.to_string()), + latency_ms: Some(latency_ms), + started_at_unix_ms: Some(candidate_started_unix_ms), + finished_at_unix_ms: Some(terminal_unix_ms), + }, + ) + .await; + } + + let Some(usage_data) = usage_seed else { + return; + }; + let mut usage_data = *usage_data; + // The seed was built when the attempt was armed, so it predates the + // diagnostics it should carry. Attaching them to the seed's metadata is the + // same write the report context would have carried into a seed built here: + // both land the same keys in the same object. + usage_data.request_metadata = attach_request_diagnostics_to_report_context( + usage_data.request_metadata.take(), + request_diagnostics.as_ref(), + ); + usage_data.status_code = Some(CLIENT_CANCELLED_STATUS_CODE); + usage_data.error_message = Some(error_message.to_string()); + usage_data.error_category = Some("cancelled".to_string()); + usage_data.response_time_ms = Some(latency_ms); + let error_body = json!({ + "error": { + "type": error_type, + "message": error_message, + "code": CLIENT_CANCELLED_STATUS_CODE + } + }); + usage_data.response_headers = Some(json!({"content-type": "application/json"})); + usage_data.response_body = Some(error_body.clone()); + usage_data.client_response_headers = Some(json!({"content-type": "application/json"})); + usage_data.client_response_body = Some(error_body); + + state + .usage_runtime + .record_terminal_event_direct( + state.usage_lifecycle_data_state().as_ref(), + UsageEvent::new(UsageEventType::Cancelled, request_id, usage_data), + ) + .await; +} + +impl Drop for AttemptCancellationGuard { + fn drop(&mut self) { + let Some(armed) = self.armed.take() else { + return; + }; + if self + .watchdog + .as_ref() + .is_some_and(|watchdog| watchdog.abandoned()) + { + return; + } + let state = self.state.clone(); + let error_type = self.error_type; + let error_message = self.error_message; + // `Drop` cannot await, and the settlement writes touch the database. + // Hand them to the runtime so they survive the dropped request future. + let Ok(handle) = tokio::runtime::Handle::try_current() else { + warn!( + event_name = "local_attempt_cancellation_guard_no_runtime", + log_type = "ops", + request_id = %short_request_id(armed.request_id.as_str()), + candidate_id = ?armed.candidate_id, + error_type, + "gateway could not settle dropped local attempt because no Tokio runtime is available" + ); + return; + }; + handle.spawn(async move { + settle_cancelled_attempt(state, armed, error_type, error_message).await; + }); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use aether_contracts::RequestBody; + use aether_data::repository::candidates::InMemoryRequestCandidateRepository; + use aether_data::repository::usage::InMemoryUsageReadRepository; + use aether_data_contracts::repository::candidates::RequestCandidateReadRepository; + use aether_data_contracts::repository::usage::{ + StoredRequestUsageAudit, UsageBodyCaptureState, UsageReadRepository, UsageWriteRepository, + }; + use aether_usage_runtime::{ + build_lifecycle_usage_seed, build_pending_usage_record, UsageRuntimeConfig, + }; + use std::collections::BTreeMap; + use std::time::Duration; + + use crate::request_candidate_runtime::{ + ensure_execution_request_candidate_slot, snapshot_local_request_candidate_status, + }; + + const TEST_ERROR_TYPE: &str = "local_stream_attempt_cancelled"; + const TEST_ERROR_MESSAGE: &str = + "Local stream attempt was dropped before terminal finalization."; + + fn test_stream_plan(request_id: &str) -> ExecutionPlan { + ExecutionPlan { + request_id: request_id.to_string(), + candidate_id: None, + provider_name: Some("Anthropic".to_string()), + provider_id: "provider-1".to_string(), + endpoint_id: "endpoint-1".to_string(), + key_id: "key-1".to_string(), + method: "POST".to_string(), + url: "https://example.test/v1/messages".to_string(), + headers: BTreeMap::new(), + content_type: Some("application/json".to_string()), + content_encoding: None, + body: RequestBody::from_json(json!({"stream": true, "service_tier": "priority"})), + stream: true, + client_api_format: "claude:messages".to_string(), + provider_api_format: "claude:messages".to_string(), + model_name: Some("claude-sonnet-4-5".to_string()), + proxy: None, + transport_profile: None, + timeouts: None, + } + } + + fn test_report_context() -> Option { + Some(json!({ + "candidate_index": 0, + "retry_index": 0, + "user_id": "user-cancel", + "api_key_id": "api-key-cancel", + "client_api_format": "claude:messages", + "provider_api_format": "claude:messages", + "request_path": "/v1/messages", + "request_path_and_query": "/v1/messages?beta=true", + "upstream_url": "https://example.test/v1/messages", + "mapped_model": "claude-sonnet-4-5", + "original_request_body": {"stream": true, "messages": []}, + })) + } + + fn test_state( + usage_repository: &Arc, + request_candidate_repository: &Arc, + ) -> AppState { + AppState::new() + .expect("gateway state should build") + .with_data_state_for_tests( + crate::data::GatewayDataState::with_request_candidate_and_usage_repository_for_tests( + Arc::clone(request_candidate_repository), + Arc::clone(usage_repository), + ), + ) + .with_usage_runtime_for_tests(UsageRuntimeConfig { + enabled: true, + ..UsageRuntimeConfig::default() + }) + } + + /// Writes the `pending` rows the same way a stream attempt does before it + /// dispatches to the provider, and returns the candidate slot snapshot the + /// attempt owns from that point on. + async fn record_pending_attempt( + state: &AppState, + plan: &mut ExecutionPlan, + report_context: &mut Option, + candidate_started_unix_ms: u64, + ) -> LocalRequestCandidateStatusSnapshot { + ensure_execution_request_candidate_slot(state, plan, report_context).await; + state.usage_runtime.record_pending( + state.usage_lifecycle_data_state().as_ref(), + build_lifecycle_usage_seed(plan, report_context.as_ref()), + ); + let snapshot = snapshot_local_request_candidate_status(plan, report_context.as_ref()) + .expect("attempt should own a candidate slot"); + record_local_request_candidate_status_snapshot( + state, + &snapshot, + SchedulerRequestCandidateStatusUpdate { + status: RequestCandidateStatus::Pending, + status_code: None, + error_type: None, + error_message: None, + latency_ms: None, + started_at_unix_ms: Some(candidate_started_unix_ms), + finished_at_unix_ms: None, + }, + ) + .await; + snapshot + } + + async fn wait_for_usage_status( + usage_repository: &InMemoryUsageReadRepository, + request_id: &str, + status: &str, + ) -> Option { + for _ in 0..50 { + if let Some(usage) = usage_repository + .find_by_request_id(request_id) + .await + .expect("usage should read") + { + if usage.status == status { + return Some(usage); + } + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + None + } + + #[tokio::test] + async fn armed_guard_settles_a_dropped_attempt_as_cancelled() { + let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); + let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); + let state = test_state(&usage_repository, &request_candidate_repository); + let mut plan = test_stream_plan("stream-cancel-guard-request"); + let mut report_context = test_report_context(); + let candidate_started_unix_ms = current_request_candidate_unix_ms(); + let snapshot = record_pending_attempt( + &state, + &mut plan, + &mut report_context, + candidate_started_unix_ms, + ) + .await; + + { + let mut guard = + AttemptCancellationGuard::disarmed(&state, TEST_ERROR_TYPE, TEST_ERROR_MESSAGE); + guard.arm( + &plan, + report_context.as_ref(), + Some(&snapshot), + candidate_started_unix_ms, + Instant::now(), + ); + } + + let usage = wait_for_usage_status( + usage_repository.as_ref(), + "stream-cancel-guard-request", + "cancelled", + ) + .await + .expect("cancelled usage should be recorded"); + assert_eq!(usage.billing_status, "void"); + assert_eq!(usage.status_code, Some(CLIENT_CANCELLED_STATUS_CODE)); + assert_eq!(usage.error_category.as_deref(), Some("cancelled")); + assert!(usage.response_time_ms.is_some()); + + let candidates = request_candidate_repository + .list_by_request_id("stream-cancel-guard-request") + .await + .expect("candidates should read"); + let candidate = candidates.first().expect("candidate row should exist"); + assert_eq!(candidate.status, RequestCandidateStatus::Cancelled); + assert_eq!(candidate.status_code, Some(CLIENT_CANCELLED_STATUS_CODE)); + assert_eq!(candidate.error_type.as_deref(), Some(TEST_ERROR_TYPE)); + assert!(candidate.finished_at_unix_ms.is_some()); + } + + /// The guard holds no request body, so its settlement write must describe the + /// capture rather than deny it: a typed `none` capture state would clear the + /// stored request body instead of leaving it alone. + #[tokio::test] + async fn settling_a_dropped_attempt_leaves_the_captured_request_body_alone() { + let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); + let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); + let state = test_state(&usage_repository, &request_candidate_repository); + let mut plan = test_stream_plan("stream-cancel-guard-capture"); + let mut report_context = test_report_context(); + let candidate_started_unix_ms = current_request_candidate_unix_ms(); + let snapshot = record_pending_attempt( + &state, + &mut plan, + &mut report_context, + candidate_started_unix_ms, + ) + .await; + // Stand in for a write that already captured this request's body. + let captured_body = json!({"stream": true, "service_tier": "priority"}); + let mut capture = build_pending_usage_record( + &plan, + report_context.as_ref(), + current_request_candidate_unix_ms() / 1_000, + ) + .expect("pending usage record should build"); + capture.provider_request_body = Some(captured_body.clone()); + capture.provider_request_body_state = Some(UsageBodyCaptureState::Inline); + usage_repository + .upsert(capture) + .await + .expect("captured request body should upsert"); + + { + let mut guard = + AttemptCancellationGuard::disarmed(&state, TEST_ERROR_TYPE, TEST_ERROR_MESSAGE); + guard.arm( + &plan, + report_context.as_ref(), + Some(&snapshot), + candidate_started_unix_ms, + Instant::now(), + ); + } + + let usage = wait_for_usage_status( + usage_repository.as_ref(), + "stream-cancel-guard-capture", + "cancelled", + ) + .await + .expect("cancelled usage should be recorded"); + assert_eq!(usage.provider_request_body, Some(captured_body)); + assert_ne!( + usage.provider_request_body_state, + Some(UsageBodyCaptureState::None) + ); + } + + #[tokio::test] + async fn guard_stands_down_when_the_watchdog_abandons_the_attempt() { + let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); + let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); + let state = test_state(&usage_repository, &request_candidate_repository); + let mut plan = test_stream_plan("stream-watchdog-guard-request"); + let mut report_context = test_report_context(); + let candidate_started_unix_ms = current_request_candidate_unix_ms(); + let snapshot = record_pending_attempt( + &state, + &mut plan, + &mut report_context, + candidate_started_unix_ms, + ) + .await; + + let watchdog = StreamCandidateWatchdogProgress::shared(); + Arc::clone(&watchdog) + .scope(async { + let mut guard = + AttemptCancellationGuard::disarmed(&state, TEST_ERROR_TYPE, TEST_ERROR_MESSAGE); + guard.arm( + &plan, + report_context.as_ref(), + Some(&snapshot), + candidate_started_unix_ms, + Instant::now(), + ); + // The watchdog gives up and takes over settlement before the + // abandoned attempt is dropped. + watchdog.mark_abandoned(); + }) + .await; + + assert!(wait_for_usage_status( + usage_repository.as_ref(), + "stream-watchdog-guard-request", + "cancelled", + ) + .await + .is_none()); + } + + #[tokio::test] + async fn disarmed_guard_leaves_the_attempt_pending() { + let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); + let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); + let state = test_state(&usage_repository, &request_candidate_repository); + let mut plan = test_stream_plan("stream-disarmed-guard-request"); + let mut report_context = test_report_context(); + let candidate_started_unix_ms = current_request_candidate_unix_ms(); + let snapshot = record_pending_attempt( + &state, + &mut plan, + &mut report_context, + candidate_started_unix_ms, + ) + .await; + + { + let mut guard = + AttemptCancellationGuard::disarmed(&state, TEST_ERROR_TYPE, TEST_ERROR_MESSAGE); + guard.arm( + &plan, + report_context.as_ref(), + Some(&snapshot), + candidate_started_unix_ms, + Instant::now(), + ); + guard.disarm(); + } + + assert!(wait_for_usage_status( + usage_repository.as_ref(), + "stream-disarmed-guard-request", + "cancelled", + ) + .await + .is_none()); + } +} diff --git a/apps/aether-gateway/src/execution_runtime/mod.rs b/apps/aether-gateway/src/execution_runtime/mod.rs index 89fbf1653..c8596db2a 100644 --- a/apps/aether-gateway/src/execution_runtime/mod.rs +++ b/apps/aether-gateway/src/execution_runtime/mod.rs @@ -4,6 +4,7 @@ use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; pub(crate) mod admission; +pub(crate) mod attempt_cancellation; pub(crate) mod attempt_lifecycle; mod chatgpt_web_image; mod constants; diff --git a/apps/aether-gateway/src/execution_runtime/stream/execution.rs b/apps/aether-gateway/src/execution_runtime/stream/execution.rs index 536326765..ba3369701 100644 --- a/apps/aether-gateway/src/execution_runtime/stream/execution.rs +++ b/apps/aether-gateway/src/execution_runtime/stream/execution.rs @@ -78,6 +78,7 @@ use crate::api::response::{ use crate::clock::current_unix_ms as current_request_candidate_unix_ms; use crate::constants::{CONTROL_CANDIDATE_ID_HEADER, CONTROL_REQUEST_ID_HEADER}; use crate::control::GatewayControlDecision; +use crate::execution_runtime::attempt_cancellation::AttemptCancellationGuard; use crate::execution_runtime::build_direct_execution_frame_stream; use crate::execution_runtime::chatgpt_web_image::maybe_execute_chatgpt_web_image_stream; use crate::execution_runtime::grok::maybe_execute_grok_stream; @@ -149,6 +150,11 @@ use crate::{ AppState, GatewayError, GEMINI_FILES_DOWNLOAD_PLAN_KIND, OPENAI_VIDEO_CONTENT_PLAN_KIND, }; +/// Settlement labels for a stream attempt whose future is dropped before the +/// transport reaches a terminal state. +const STREAM_ATTEMPT_CANCELLED_ERROR_TYPE: &str = "local_stream_attempt_cancelled"; +const STREAM_ATTEMPT_CANCELLED_ERROR_MESSAGE: &str = "Local stream attempt was dropped before terminal finalization, usually because the client disconnected or the request task was cancelled."; + const SSE_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(15); const SSE_KEEPALIVE_BYTES: &[u8] = b": aether-keepalive\n\n"; const SSE_CONTROL_FILTER_MAX_BUFFER_BYTES: usize = 1024 * 1024; @@ -3655,17 +3661,30 @@ pub(crate) fn execute_execution_runtime_stream<'a>( report_kind: Option, report_context: Option, ) -> Pin>, GatewayError>> + Send + 'a>> { - Box::pin(execute_execution_runtime_stream_inner( - state, - plan, - trace_id, - decision, - plan_kind, - report_kind, - report_context, - None, - None, - )) + Box::pin(async move { + let mut cancellation_guard = AttemptCancellationGuard::disarmed( + state, + STREAM_ATTEMPT_CANCELLED_ERROR_TYPE, + STREAM_ATTEMPT_CANCELLED_ERROR_MESSAGE, + ); + let result = execute_execution_runtime_stream_inner( + state, + plan, + trace_id, + decision, + plan_kind, + report_kind, + report_context, + None, + None, + &mut cancellation_guard, + ) + .await; + // The attempt reached its own terminal path, or handed settlement to the + // stream finalizer that now lives in the response body. + cancellation_guard.disarm(); + result + }) } #[allow(clippy::too_many_arguments)] @@ -3687,7 +3706,12 @@ pub(crate) fn execute_execution_runtime_stream_with_retry_scope<'a>( Box::pin(async move { let mut retry_scope = AiAttemptRetryScope::Candidate; let mut fallback_response = None; - let response = execute_execution_runtime_stream_inner( + let mut cancellation_guard = AttemptCancellationGuard::disarmed( + state, + STREAM_ATTEMPT_CANCELLED_ERROR_TYPE, + STREAM_ATTEMPT_CANCELLED_ERROR_MESSAGE, + ); + let result = execute_execution_runtime_stream_inner( state, plan, trace_id, @@ -3697,8 +3721,13 @@ pub(crate) fn execute_execution_runtime_stream_with_retry_scope<'a>( report_context, Some(&mut retry_scope), Some(&mut fallback_response), + &mut cancellation_guard, ) - .await?; + .await; + // The attempt reached its own terminal path, or handed settlement to the + // stream finalizer that now lives in the response body. + cancellation_guard.disarm(); + let response = result?; Ok(match response { Some(response) => AiAttemptExecutionOutcome::Responded(response), None => AiAttemptExecutionOutcome::Retry { @@ -3744,6 +3773,7 @@ async fn maybe_build_stream_transport_error_stop_response( .map(Some) } +#[allow(clippy::too_many_arguments)] // internal function, grouping would add unnecessary indirection async fn execute_execution_runtime_stream_inner( state: &AppState, mut plan: ExecutionPlan, @@ -3754,6 +3784,7 @@ async fn execute_execution_runtime_stream_inner( mut report_context: Option, mut retry_scope_out: Option<&mut AiAttemptRetryScope>, mut retry_fallback_out: Option<&mut Option>>, + cancellation_guard: &mut AttemptCancellationGuard, ) -> Result>, GatewayError> { let stream_started_at = Instant::now(); let mut stage_trace = RequestStageTrace::from_env(); @@ -3837,6 +3868,16 @@ async fn execute_execution_runtime_stream_inner( ) .await; } + // From here the attempt owns non-terminal rows, and everything that could + // settle them runs inside the downstream request future. Arm the guard so a + // client disconnect before the stream finalizer exists still settles them. + cancellation_guard.arm( + &plan, + report_context.as_ref(), + request_candidate_status_snapshot.as_ref(), + candidate_started_unix_secs, + stream_started_at, + ); let plan_request_id_for_log = short_request_id(plan.request_id.as_str()); let provider_name = plan .provider_name diff --git a/apps/aether-gateway/src/execution_runtime/transport_failure.rs b/apps/aether-gateway/src/execution_runtime/transport_failure.rs index ccd532adf..c3f1885fb 100644 --- a/apps/aether-gateway/src/execution_runtime/transport_failure.rs +++ b/apps/aether-gateway/src/execution_runtime/transport_failure.rs @@ -22,6 +22,7 @@ const TRANSPORT_ERROR_CLIENT_MESSAGE: &str = #[derive(Debug, Default)] pub(crate) struct StreamCandidateWatchdogProgress { terminal_started: AtomicBool, + abandoned: AtomicBool, } tokio::task_local! { @@ -37,6 +38,24 @@ impl StreamCandidateWatchdogProgress { self.terminal_started.load(Ordering::Acquire) } + /// The watchdog gave up waiting and settles this attempt itself. + /// + /// The attempt future is dropped once the watchdog returns, so its own + /// cancellation guard must stay out of the way instead of racing the + /// watchdog's terminal rows with a cancellation. + pub(crate) fn mark_abandoned(&self) { + self.abandoned.store(true, Ordering::Release); + } + + pub(crate) fn abandoned(&self) -> bool { + self.abandoned.load(Ordering::Acquire) + } + + /// The watchdog watching the attempt on this task, if it runs under one. + pub(crate) fn current() -> Option> { + STREAM_CANDIDATE_WATCHDOG_PROGRESS.try_with(Arc::clone).ok() + } + pub(crate) async fn scope(self: Arc, future: F) -> F::Output where F: Future, diff --git a/apps/aether-gateway/src/executor/candidate_loop.rs b/apps/aether-gateway/src/executor/candidate_loop.rs index 9a689089d..42f262a78 100644 --- a/apps/aether-gateway/src/executor/candidate_loop.rs +++ b/apps/aether-gateway/src/executor/candidate_loop.rs @@ -1418,6 +1418,10 @@ where let outcome = match execution_result { Some(result) => result.map(StreamCandidateWatchdogOutcome::Executed), None => { + // The abandoned attempt is dropped when this function returns. + // Claim its settlement before that so its cancellation guard does + // not race the watchdog rows written just below. + watchdog_progress.mark_abandoned(); let finished_at_unix_ms = current_unix_ms(); let request_id = short_request_id(plan.request_id.as_str()); let provider_name = plan.provider_name.as_deref().unwrap_or("-"); diff --git a/apps/aether-gateway/src/tests/ai_execute/lifecycle.rs b/apps/aether-gateway/src/tests/ai_execute/lifecycle.rs index c3cae80b1..6e26a8976 100644 --- a/apps/aether-gateway/src/tests/ai_execute/lifecycle.rs +++ b/apps/aether-gateway/src/tests/ai_execute/lifecycle.rs @@ -14,6 +14,9 @@ use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadReposi use aether_data_contracts::repository::candidate_selection::{ StoredMinimalCandidateSelectionRow, StoredProviderModelMapping, }; +use aether_data_contracts::repository::candidates::{ + RequestCandidateReadRepository, RequestCandidateStatus, +}; use aether_data_contracts::repository::provider_catalog::{ StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider, }; @@ -427,6 +430,112 @@ async fn gateway_stops_execution_runtime_stream_when_client_disconnects_impl() { upstream_handle.abort(); } +#[test] +fn gateway_settles_stream_attempt_when_client_disconnects_before_first_byte() { + run_lifecycle_test( + "gateway_settles_stream_attempt_when_client_disconnects_before_first_byte", + gateway_settles_stream_attempt_when_client_disconnects_before_first_byte_impl, + ); +} + +async fn gateway_settles_stream_attempt_when_client_disconnects_before_first_byte_impl() { + // The execution runtime accepts the plan and then goes quiet, so the attempt + // is parked between its `pending` rows and the first upstream byte. + let execution_runtime = Router::new().route( + "/v1/execute/stream", + any(|_request: Request| async move { + tokio::time::sleep(std::time::Duration::from_secs(30)).await; + StatusCode::OK + }), + ); + + let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await; + let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( + Some(hash_api_key("sk-client-openai-stream-precommit-disconnect")), + sample_local_openai_auth_snapshot( + "api-key-openai-lifecycle-local-1", + "user-openai-lifecycle-local-1", + ), + )])); + let candidate_selection_repository = + Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![ + sample_local_openai_candidate_row(), + ])); + let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( + vec![sample_local_openai_provider()], + vec![sample_local_openai_endpoint()], + vec![sample_local_openai_key()], + )); + let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); + let gateway = build_router_with_state( + build_state_with_execution_runtime_override(execution_runtime_url) + .with_data_state_for_tests( + GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests( + auth_repository, + candidate_selection_repository, + provider_catalog_repository, + Arc::clone(&request_candidate_repository), + DEVELOPMENT_ENCRYPTION_KEY, + ), + ), + ); + let (gateway_url, gateway_handle) = start_server(gateway).await; + + let request = reqwest::Client::new() + .post(format!("{gateway_url}/v1/chat/completions")) + .header(http::header::CONTENT_TYPE, "application/json") + .header( + http::header::AUTHORIZATION, + "Bearer sk-client-openai-stream-precommit-disconnect", + ) + .header( + TRACE_ID_HEADER, + "trace-openai-chat-stream-precommit-disconnect-123", + ) + .body("{\"model\":\"gpt-5\",\"messages\":[],\"stream\":true}") + .send(); + + // Drop the in-flight request the way a downstream client does when its own + // first-byte timeout fires, before any response header exists. + assert!( + tokio::time::timeout(std::time::Duration::from_millis(750), request) + .await + .is_err(), + "the execution runtime should not have answered before the client gave up" + ); + + let mut stored_candidates = Vec::new(); + for _ in 0..200 { + stored_candidates = request_candidate_repository + .list_by_request_id("trace-openai-chat-stream-precommit-disconnect-123") + .await + .expect("request candidate trace should read"); + if stored_candidates + .iter() + .any(|candidate| candidate.status == RequestCandidateStatus::Cancelled) + { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + + let cancelled = stored_candidates + .iter() + .find(|candidate| candidate.status == RequestCandidateStatus::Cancelled) + .unwrap_or_else(|| { + panic!("dropped stream attempt should settle as cancelled: {stored_candidates:?}") + }); + assert_eq!(cancelled.status_code, Some(499)); + assert_eq!( + cancelled.error_type.as_deref(), + Some("local_stream_attempt_cancelled") + ); + assert!(cancelled.finished_at_unix_ms.is_some()); + + gateway_handle.abort(); + execution_runtime_handle.abort(); +} + #[test] fn gateway_returns_error_body_when_prefetch_detects_embedded_stream_error() { run_lifecycle_test( diff --git a/crates/aether-usage/runtime/src/lib.rs b/crates/aether-usage/runtime/src/lib.rs index 0be0b5b97..e37d38eef 100644 --- a/crates/aether-usage/runtime/src/lib.rs +++ b/crates/aether-usage/runtime/src/lib.rs @@ -61,7 +61,8 @@ pub use write::{ build_sync_terminal_usage_event, build_sync_terminal_usage_outcome, build_sync_terminal_usage_payload_seed, build_sync_terminal_usage_seed, build_terminal_usage_context_seed, build_terminal_usage_event_from_outcome, - build_terminal_usage_event_from_seed, build_usage_event_data_seed, LifecycleUsageSeed, + build_terminal_usage_event_from_seed, build_usage_event_data_seed, + build_usage_event_data_seed_describing_request_bodies, LifecycleUsageSeed, StreamTerminalUsagePayloadSeed, SyncTerminalUsagePayloadSeed, TerminalUsageContextSeed, TerminalUsageOutcome, TerminalUsageSeed, UsageTerminalState, }; diff --git a/crates/aether-usage/runtime/src/write.rs b/crates/aether-usage/runtime/src/write.rs index ba332fdab..cfd61fab2 100644 --- a/crates/aether-usage/runtime/src/write.rs +++ b/crates/aether-usage/runtime/src/write.rs @@ -72,6 +72,20 @@ struct RuntimeRequestCaptureSeed { provider_request: Option, provider_request_body_ref: Option, body_states: UsageBodyStatesSeed, + request_has_inline_body: bool, + provider_request_has_inline_body: bool, +} + +/// Whether a seed keeps the request bodies it describes, or only describes them. +/// +/// A holder that has to outlive the request itself pays for every byte it keeps, +/// and a request body can be megabytes. [`RequestBodyCapture::Describe`] computes +/// every capture state, reference and derived fact from the real plan and report +/// context, and leaves out only the body values themselves. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum RequestBodyCapture { + Keep, + Describe, } #[derive(Debug, Clone, PartialEq)] @@ -804,7 +818,8 @@ pub fn build_terminal_usage_context_seed( report_context: Option<&Value>, ) -> TerminalUsageContextSeed { let context = report_context.and_then(Value::as_object); - let request_capture = build_runtime_request_capture_seed(plan, context); + let request_capture = + build_runtime_request_capture_seed(plan, context, RequestBodyCapture::Keep); let client_contract = context_string(context, "client_contract") .or_else(|| context_string(context, "client_api_format")) .or_else(|| non_empty_str(Some(plan.client_api_format.as_str()))) @@ -1691,16 +1706,34 @@ pub fn build_usage_event_data_seed( plan: &ExecutionPlan, report_context: Option<&Value>, ) -> UsageEventData { - build_usage_event_data_seed_with_detail(plan, report_context) + build_usage_event_data_seed_with_detail(plan, report_context, RequestBodyCapture::Keep) +} + +/// Builds the same seed as [`build_usage_event_data_seed`] without keeping the +/// request bodies. +/// +/// This is for a caller that has to hold a seed for the whole life of an attempt +/// so it can still write a terminal row if the attempt is dropped: a request body +/// can be megabytes, and holding one per in-flight attempt is far more expensive +/// than the row it would eventually capture. Every capture state, body reference +/// and derived request fact is still computed from the real plan and report +/// context, so the resulting terminal write preserves the capture an earlier +/// non-terminal write recorded rather than clearing it. +pub fn build_usage_event_data_seed_describing_request_bodies( + plan: &ExecutionPlan, + report_context: Option<&Value>, +) -> UsageEventData { + build_usage_event_data_seed_with_detail(plan, report_context, RequestBodyCapture::Describe) } fn build_usage_event_data_seed_with_detail( plan: &ExecutionPlan, report_context: Option<&Value>, + capture: RequestBodyCapture, ) -> UsageEventData { let context = report_context.and_then(Value::as_object); let routing = build_runtime_routing_seed(plan, context); - let request_capture = build_runtime_request_capture_seed(plan, context); + let request_capture = build_runtime_request_capture_seed(plan, context, capture); let api_format = context_string(context, "client_api_format") .or_else(|| non_empty_str(Some(plan.client_api_format.as_str()))); let endpoint_api_format = context_string(context, "provider_api_format") @@ -1714,7 +1747,10 @@ fn build_usage_event_data_seed_with_detail( let request_type = Some(infer_request_type_from_contracts( api_format.as_deref(), endpoint_api_format.as_deref(), - request_capture.provider_request.as_ref(), + request_capture + .provider_request + .as_ref() + .or_else(|| provider_request_body_ref_for_inference(plan, context)), )); let api_family = api_format .as_deref() @@ -1737,9 +1773,9 @@ fn build_usage_event_data_seed_with_detail( build_runtime_request_metadata_seed_from_parts( plan, context, - request_capture.request_body.is_some(), + request_capture.request_has_inline_body, request_capture.request_body_ref.as_deref(), - request_capture.provider_request.is_some(), + request_capture.provider_request_has_inline_body, request_capture.provider_request_body_ref.as_deref(), plan.body.body_bytes_b64.as_deref(), ), @@ -1999,17 +2035,29 @@ fn plan_has_inline_json_body_for_usage(plan: &ExecutionPlan) -> bool { fn build_runtime_request_capture_seed( plan: &ExecutionPlan, context: Option<&Map>, + capture: RequestBodyCapture, ) -> RuntimeRequestCaptureSeed { - let request_body = context_body_value(context, "original_request_body"); + // Presence, not the value, is what every capture state and derived fact is + // built from, so both capture modes agree on all of them. + let request_has_inline_body = context_has_inline_body(context, "original_request_body"); + let provider_request_has_inline_body = + context_has_inline_body(context, "provider_request_body") + || plan_has_inline_json_body_for_usage(plan); + let (request_body, provider_request) = match capture { + RequestBodyCapture::Keep => ( + context_body_value(context, "original_request_body"), + context_body_value(context, "provider_request_body") + .or_else(|| plan_json_body_capture_for_usage(plan)), + ), + RequestBodyCapture::Describe => (None, None), + }; let request_body_ref = context_string(context, "request_body_ref"); - let provider_request = context_body_value(context, "provider_request_body") - .or_else(|| plan_json_body_capture_for_usage(plan)); let provider_request_body_ref = context_string(context, "provider_request_body_ref") .or_else(|| non_empty_str(plan.body.body_ref.as_deref())); let body_states = build_runtime_body_states_seed_from_parts( - request_body.is_some(), + request_has_inline_body, request_body_ref.as_deref(), - provider_request.is_some(), + provider_request_has_inline_body, provider_request_body_ref.as_deref(), plan.body.body_bytes_b64.is_some(), ); @@ -2020,9 +2068,26 @@ fn build_runtime_request_capture_seed( provider_request, provider_request_body_ref, body_states, + request_has_inline_body, + provider_request_has_inline_body, } } +/// Borrows the provider request body that request-type inference reads, without +/// cloning it. +fn provider_request_body_ref_for_inference<'a>( + plan: &'a ExecutionPlan, + context: Option<&'a Map>, +) -> Option<&'a Value> { + context_value_ref(context, "provider_request_body") + .filter(|value| !value.is_null()) + .or_else(|| { + plan_has_inline_json_body_for_usage(plan) + .then_some(plan.body.json_body.as_ref()) + .flatten() + }) +} + fn build_runtime_request_metadata_seed( plan: &ExecutionPlan, context: Option<&Map>, @@ -3493,7 +3558,8 @@ mod tests { build_streaming_usage_event_from_owned_seed, build_streaming_usage_record, build_sync_terminal_usage_event, build_sync_terminal_usage_payload_seed, build_sync_terminal_usage_seed, build_terminal_usage_context_seed, - build_terminal_usage_event_from_seed, build_usage_event_data_seed, decode_body_for_storage, + build_terminal_usage_event_from_seed, build_usage_event_data_seed, + build_usage_event_data_seed_describing_request_bodies, decode_body_for_storage, extract_token_counts_from_json, extract_token_counts_from_value, headers_to_json, mask_header_value, mask_sensitive_body_fields, mask_sensitive_headers_in_json_value, parse_sse_body_for_storage, resolve_error_message, trim_owned_non_empty_string, @@ -6842,6 +6908,114 @@ mod tests { assert!(body_size.get("provider_request_body").is_some()); } + #[test] + fn describing_request_bodies_matches_the_capturing_seed_apart_from_the_bodies() { + let plan = ExecutionPlan { + request_id: "req-seed-describe-1".to_string(), + candidate_id: Some("cand-seed-describe-1".to_string()), + provider_name: Some("OpenAI".to_string()), + provider_id: "provider-1".to_string(), + endpoint_id: "endpoint-1".to_string(), + key_id: "key-1".to_string(), + method: "POST".to_string(), + url: "https://example.com/v1/chat/completions".to_string(), + headers: BTreeMap::new(), + content_type: Some("application/json".to_string()), + content_encoding: None, + body: RequestBody::from_json(json!({ + "model": "gpt-5", + "service_tier": "priority", + "reasoning": {"effort": "high"} + })), + stream: false, + client_api_format: "openai:chat".to_string(), + provider_api_format: "openai:chat".to_string(), + model_name: Some("gpt-5".to_string()), + proxy: None, + transport_profile: None, + timeouts: None, + }; + let report_context = json!({ + "client_api_format": "openai:chat", + "provider_api_format": "openai:chat", + "original_request_body": {"model": "gpt-5", "messages": []}, + "original_headers": {"accept": "application/json"} + }); + + let captured = build_usage_event_data_seed(&plan, Some(&report_context)); + let described = + build_usage_event_data_seed_describing_request_bodies(&plan, Some(&report_context)); + + // Only the two heavy values differ. + assert!(captured.request_body.is_some()); + assert!(captured.provider_request_body.is_some()); + assert_eq!(described.request_body, None); + assert_eq!(described.provider_request_body, None); + + // Everything a terminal write reads to decide what to do with the stored + // capture is identical, so the described seed preserves it rather than + // clearing it. + assert_eq!( + described.request_body_state, + Some(UsageBodyCaptureState::Inline) + ); + assert_eq!( + described.provider_request_body_state, + Some(UsageBodyCaptureState::Inline) + ); + assert_eq!(described.request_body_state, captured.request_body_state); + assert_eq!( + described.provider_request_body_state, + captured.provider_request_body_state + ); + assert_eq!(described.request_body_ref, captured.request_body_ref); + assert_eq!( + described.provider_request_body_ref, + captured.provider_request_body_ref + ); + assert_eq!(described.request_type, captured.request_type); + assert_eq!(described.request_metadata, captured.request_metadata); + assert_eq!( + described.provider_request_headers, + captured.provider_request_headers + ); + assert_eq!(described.request_headers, captured.request_headers); + } + + #[test] + fn describing_request_bodies_keeps_the_unavailable_marker_for_raw_bodies() { + let mut plan = ExecutionPlan { + request_id: "req-seed-describe-2".to_string(), + candidate_id: None, + provider_name: Some("OpenAI".to_string()), + provider_id: "provider-1".to_string(), + endpoint_id: "endpoint-1".to_string(), + key_id: "key-1".to_string(), + method: "POST".to_string(), + url: "https://example.com/v1/chat/completions".to_string(), + headers: BTreeMap::new(), + content_type: Some("application/json".to_string()), + content_encoding: None, + body: RequestBody::from_json(json!({"model": "gpt-5"})), + stream: false, + client_api_format: "openai:chat".to_string(), + provider_api_format: "openai:chat".to_string(), + model_name: Some("gpt-5".to_string()), + proxy: None, + transport_profile: None, + timeouts: None, + }; + plan.body.json_body = None; + plan.body.body_bytes_b64 = Some("eyJtb2RlbCI6ICJncHQtNSJ9".to_string()); + + let described = build_usage_event_data_seed_describing_request_bodies(&plan, None); + + assert_eq!( + described.provider_request_body_state, + Some(UsageBodyCaptureState::Unavailable) + ); + } + #[test] fn masks_known_sensitive_header_values() { let token = "Bearer eyJhbGciOiJSUzI1NiJ9.payload-here.signature-tail"; From 206995645b8a5177a285b1025bc7477d0c2e0322 Mon Sep 17 00:00:00 2001 From: ZheFox <77232781+zhefox@users.noreply.github.com> Date: Fri, 4 Sep 2026 18:01:19 +0800 Subject: [PATCH 10/13] fix(gateway): share stream first-byte deadline across retries --- .../src/executor/candidate_loop.rs | 82 ++++++++++++++++++- 1 file changed, 80 insertions(+), 2 deletions(-) diff --git a/apps/aether-gateway/src/executor/candidate_loop.rs b/apps/aether-gateway/src/executor/candidate_loop.rs index 42f262a78..be1f9f876 100644 --- a/apps/aether-gateway/src/executor/candidate_loop.rs +++ b/apps/aether-gateway/src/executor/candidate_loop.rs @@ -412,6 +412,7 @@ where decision, plan_kind, transfer_tracker, + request_first_byte_started_at: Instant::now(), }; match run_ai_attempt_loop(&port, plan_and_reports).await? { AiAttemptLoopOutcome::Responded(response) => { @@ -482,6 +483,7 @@ where decision, plan_kind, transfer_tracker, + request_first_byte_started_at: Instant::now(), }; run_dynamic_attempt_loop( &port, @@ -946,6 +948,10 @@ struct StreamAttemptLoopPort<'a> { decision: &'a GatewayControlDecision, plan_kind: &'a str, transfer_tracker: &'a ProviderTransferTracker, + /// All candidates in one downstream stream request share this origin. + /// Without it every retry receives a fresh full first-byte timeout and a + /// 30-second provider timeout can accumulate into a 60-120 second stall. + request_first_byte_started_at: Instant, } #[async_trait] @@ -1058,6 +1064,7 @@ where self.plan_kind, plan, watchdog_report_context, + self.request_first_byte_started_at, stop_on_transport_errors, move || async move { execute_execution_runtime_stream_with_retry_scope( @@ -1085,7 +1092,7 @@ where http::StatusCode::GATEWAY_TIMEOUT.as_u16(), "local_stream_candidate_watchdog_timeout", stream_candidate_watchdog_timeout_message(), - watchdog_started_at.elapsed().as_millis() as u64, + self.request_first_byte_started_at.elapsed().as_millis() as u64, ) .await?, ) @@ -1368,6 +1375,7 @@ async fn execute_stream_candidate_with_watchdog( plan_kind: &str, plan: &aether_contracts::ExecutionPlan, report_context: Option<&serde_json::Value>, + request_first_byte_started_at: Instant, stop_on_transport_errors: bool, execute: impl FnOnce() -> Fut, ) -> Result @@ -1377,6 +1385,7 @@ where > + Send, { let timeout_duration = resolve_stream_candidate_watchdog_timeout(plan, report_context); + let request_first_byte_deadline = request_first_byte_started_at + timeout_duration; let candidate_started_at = std::time::Instant::now(); let candidate_started_unix_ms = current_unix_ms(); let permit = match acquire_upstream_execution_gate(state, trace_id).await { @@ -1402,7 +1411,14 @@ where let watchdog_progress = StreamCandidateWatchdogProgress::shared(); let execution = watchdog_progress.clone().scope(execute()); tokio::pin!(execution); - let deadline = tokio::time::sleep(timeout_duration); + // This is an absolute request-level deadline, not a new timeout for this + // candidate. Retries therefore consume only the budget left by earlier + // candidates instead of resetting the full provider timeout. + let candidate_budget_ms = request_first_byte_deadline + .saturating_duration_since(Instant::now()) + .as_millis() + .min(u128::from(u64::MAX)) as u64; + let deadline = tokio::time::sleep_until(request_first_byte_deadline); tokio::pin!(deadline); let execution_result = tokio::select! { biased; @@ -1431,6 +1447,10 @@ where .map(|value| value.to_string()) .unwrap_or_else(|| "-".to_string()); let timeout_ms = u64::try_from(timeout_duration.as_millis()).unwrap_or(u64::MAX); + let request_elapsed_ms = request_first_byte_started_at + .elapsed() + .as_millis() + .min(u128::from(u64::MAX)) as u64; record_local_request_candidate_status( state, plan, @@ -1459,6 +1479,8 @@ where model_name, candidate_index = candidate_index.as_str(), timeout_ms, + candidate_budget_ms, + request_elapsed_ms, "gateway local stream candidate watchdog timed out" ); if stop_on_transport_errors { @@ -2389,6 +2411,7 @@ mod tests { "claude_cli_stream", &plan, Some(&report_context), + Instant::now(), false, || { std::future::pending::< @@ -2427,6 +2450,56 @@ mod tests { assert_eq!(record.candidate_index, 2); } + #[tokio::test] + async fn stream_candidate_retry_does_not_reset_an_expired_request_first_byte_budget() { + let writer = Arc::new(TestRequestCandidateWriter::default()); + let plan = test_plan(Some(ExecutionTimeouts { + first_byte_ms: Some(250), + ..ExecutionTimeouts::default() + })); + let report_context = test_report_context(); + // Stand in for earlier candidates having already consumed the request's + // complete first-byte budget. A per-candidate watchdog would wait a new + // 250 ms here; the shared absolute deadline must settle immediately. + let request_first_byte_started_at = Instant::now() - Duration::from_millis(300); + + let result = tokio::time::timeout( + Duration::from_millis(100), + execute_stream_candidate_with_watchdog( + writer.as_ref(), + "trace_watchdog_shared_budget", + "claude_cli_stream", + &plan, + Some(&report_context), + request_first_byte_started_at, + false, + || { + std::future::pending::< + Result>, GatewayError>, + >() + }, + ), + ) + .await + .expect("an expired request-level first-byte budget must not restart per candidate"); + + assert!(matches!( + result, + Ok(StreamCandidateWatchdogOutcome::Executed( + AiAttemptExecutionOutcome::Retry { + scope: AiAttemptRetryScope::Candidate, + fallback_response: None, + } + )) + )); + let records = writer.records.lock().await; + assert_eq!(records.len(), 1); + assert_eq!( + records[0].error_type.as_deref(), + Some("local_stream_candidate_watchdog_timeout") + ); + } + #[tokio::test] async fn stream_candidate_watchdog_can_stop_on_transport_error() { let writer = Arc::new(TestRequestCandidateWriter::default()); @@ -2442,6 +2515,7 @@ mod tests { "claude_cli_stream", &plan, Some(&report_context), + Instant::now(), true, || { std::future::pending::< @@ -2479,6 +2553,7 @@ mod tests { "claude_cli_stream", &plan, Some(&report_context), + Instant::now(), true, || async { mark_stream_candidate_watchdog_terminal_started(); @@ -2511,6 +2586,7 @@ mod tests { "claude_cli_stream", &plan, Some(&report_context), + Instant::now(), true, || async { Err(GatewayError::UpstreamUnavailable { @@ -2550,6 +2626,7 @@ mod tests { "claude_cli_stream", &plan, Some(&report_context), + Instant::now(), false, || async { panic!("execute future should not run while upstream execution gate is saturated") @@ -2597,6 +2674,7 @@ mod tests { "claude_cli_stream", &plan, Some(&report_context), + Instant::now(), false, || async { Err(GatewayError::AdmissionTimeout { From 86f7cc0d588d482a12843a3663a89b3b770e63e6 Mon Sep 17 00:00:00 2001 From: ZheFox <77232781+zhefox@users.noreply.github.com> Date: Fri, 4 Sep 2026 18:01:19 +0800 Subject: [PATCH 11/13] test(gateway): satisfy Rust 1.95 integer lint --- apps/aether-gateway/src/data/state/routing_group_cache.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apps/aether-gateway/src/data/state/routing_group_cache.rs b/apps/aether-gateway/src/data/state/routing_group_cache.rs index 47eda5871..b03ac9011 100644 --- a/apps/aether-gateway/src/data/state/routing_group_cache.rs +++ b/apps/aether-gateway/src/data/state/routing_group_cache.rs @@ -684,7 +684,7 @@ mod tests { .unwrap_or_default(); // One Arc is retained by the map and every active request // owns one through its leader guard or follower state. - if participant_count >= participants + 1 { + if participant_count > participants { break; } tokio::task::yield_now().await; From 344b3031e9f6716ac8ac10b96565b174030d5349 Mon Sep 17 00:00:00 2001 From: ZheFox <77232781+zhefox@users.noreply.github.com> Date: Fri, 4 Sep 2026 18:59:52 +0800 Subject: [PATCH 12/13] test(integration): stabilize truncated h2c stream ordering --- .../integration/src/bin/mock_openai_upstream.rs | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/crates/aether-testing/integration/src/bin/mock_openai_upstream.rs b/crates/aether-testing/integration/src/bin/mock_openai_upstream.rs index 58c297519..8dacebdb2 100644 --- a/crates/aether-testing/integration/src/bin/mock_openai_upstream.rs +++ b/crates/aether-testing/integration/src/bin/mock_openai_upstream.rs @@ -37,6 +37,7 @@ const MAX_MOCK_REQUEST_BODY_BYTES: usize = 1024 * 1024; const BASIS_POINTS: u16 = 10_000; const DEFAULT_TIMEOUT_HOLD_MS: u64 = 60_000; const REQUEST_SEQUENCE_HEADER: &str = "x-mock-request-sequence"; +const TRUNCATED_STREAM_FLUSH_DELAY: Duration = Duration::from_millis(10); const RANDOM_DOMAIN_FAULT: u64 = 0x5c32_22f7_27d4_7a6f; const RANDOM_DOMAIN_FIRST_BYTE: u64 = 0x087d_89d9_3bc3_15db; @@ -646,8 +647,10 @@ fn build_chat_sse_response( if profile.truncate_after_chunks == Some(0) || profile.truncate_after_chunks == Some(index + 1) { - // Force Hyper to flush the successful frame before observing the body error. - tokio::task::yield_now().await; + // Hyper translates body errors into RST_STREAM for HTTP/2. Keep the body + // pending briefly so the response headers and successful DATA frame are + // written before Hyper observes the error. + tokio::time::sleep(TRUNCATED_STREAM_FLUSH_DELAY).await; record_fault(&app, Fault::TruncateStream); yield Err::(truncated_stream_error()); return; @@ -701,8 +704,10 @@ fn build_responses_sse_response( if profile.truncate_after_chunks == Some(0) || profile.truncate_after_chunks == Some(index + 1) { - // Force Hyper to flush the successful frame before observing the body error. - tokio::task::yield_now().await; + // Hyper translates body errors into RST_STREAM for HTTP/2. Keep the body + // pending briefly so the response headers and successful DATA frame are + // written before Hyper observes the error. + tokio::time::sleep(TRUNCATED_STREAM_FLUSH_DELAY).await; record_fault(&app, Fault::TruncateStream); yield Err::(truncated_stream_error()); return; From 57cdef4b8d45bf8ab842d45b9fdbcc61f4c42672 Mon Sep 17 00:00:00 2001 From: ZheFox <77232781+zhefox@users.noreply.github.com> Date: Fri, 4 Sep 2026 19:53:06 +0800 Subject: [PATCH 13/13] fix(gateway): harden cross-format sync finalization Co-authored-by: stabey <36232531+stabey@users.noreply.github.com> --- .../src/execution_runtime/fallback.rs | 87 +++++++- .../src/formats/shared/sync_products.rs | 185 +++++++++++++++++- 2 files changed, 269 insertions(+), 3 deletions(-) diff --git a/apps/aether-gateway/src/execution_runtime/fallback.rs b/apps/aether-gateway/src/execution_runtime/fallback.rs index 3c7287b6d..6954e4eac 100644 --- a/apps/aether-gateway/src/execution_runtime/fallback.rs +++ b/apps/aether-gateway/src/execution_runtime/fallback.rs @@ -156,7 +156,22 @@ pub(crate) fn should_fallback_to_control_sync( return true; }; - body_json.get("error").is_some() + sync_body_has_embedded_error(Some(body_json)) +} + +/// Mirrors the error-like body markers used by the formats layer. Successful OpenAI Responses +/// bodies contain `"error": null`, which must not route them through error finalization. +fn sync_body_has_embedded_error(body_json: Option<&serde_json::Value>) -> bool { + let Some(object) = body_json.and_then(serde_json::Value::as_object) else { + return false; + }; + + object.get("error").is_some_and(|error| !error.is_null()) + || object.get("status").and_then(serde_json::Value::as_str) == Some("failed") + || object + .get("type") + .and_then(serde_json::Value::as_str) + .is_some_and(|value| value == "error") } pub(crate) fn should_finalize_sync_response(report_kind: Option<&str>) -> bool { @@ -168,7 +183,7 @@ pub(crate) fn resolve_core_sync_error_finalize_report_kind( result: &ExecutionResult, body_json: Option<&serde_json::Value>, ) -> Option { - let has_embedded_error = body_json.is_some_and(|value| value.get("error").is_some()); + let has_embedded_error = sync_body_has_embedded_error(body_json); if result.status_code < 400 && !has_embedded_error { return None; } @@ -500,6 +515,74 @@ mod tests { ); } + #[test] + fn successful_responses_body_with_null_error_stays_on_success_path() { + let result = ExecutionResult { + request_id: "req-1".to_string(), + candidate_id: None, + status_code: 200, + headers: Default::default(), + response_observation: None, + body: None, + telemetry: None, + error: None, + }; + let body_json = serde_json::json!({ + "id": "resp_1", + "object": "response", + "status": "completed", + "error": null, + "output": [], + }); + + assert_eq!( + resolve_core_sync_error_finalize_report_kind( + "openai_responses_sync", + &result, + Some(&body_json) + ), + None + ); + assert!(!should_fallback_to_control_sync( + "openai_responses_sync", + &result, + Some(&body_json), + true, + false, + false, + )); + } + + #[test] + fn error_like_success_status_bodies_still_map_to_error_finalize() { + let result = ExecutionResult { + request_id: "req-1".to_string(), + candidate_id: None, + status_code: 200, + headers: Default::default(), + response_observation: None, + body: None, + telemetry: None, + error: None, + }; + + for body_json in [ + serde_json::json!({"status": "failed", "error": null}), + serde_json::json!({"type": "error"}), + serde_json::json!({"error": {"message": "boom"}}), + ] { + assert_eq!( + resolve_core_sync_error_finalize_report_kind( + "openai_responses_sync", + &result, + Some(&body_json) + ), + Some("openai_responses_sync_finalize".to_string()), + "error-like body must not escape through the success path: {body_json}" + ); + } + } + #[test] fn stream_failover_marks_chat_errors() { assert!(should_fallback_to_control_stream( diff --git a/crates/aether-ai/formats/src/formats/shared/sync_products.rs b/crates/aether-ai/formats/src/formats/shared/sync_products.rs index 5ad5c21ea..929b2cfa9 100644 --- a/crates/aether-ai/formats/src/formats/shared/sync_products.rs +++ b/crates/aether-ai/formats/src/formats/shared/sync_products.rs @@ -471,6 +471,20 @@ pub fn maybe_build_standard_sync_finalize_product_from_normalized_payload( }; let body_base64 = body_base64.or(capture_stream_body_base64.as_deref()); + // Cross-format sync attempts can contain raw bytes because the plan requested a stream even + // though the provider returned one complete JSON response. Do not feed that response into an + // SSE aggregator. Capture envelopes and same-format responses retain their existing precedence. + let non_stream_capture_body_json = + if capture_envelope_used || !sync_finalize_needs_conversion(report_context) { + None + } else { + body_base64.and_then(decode_non_stream_sync_capture_body) + }; + let (body_json, body_base64) = match non_stream_capture_body_json.as_ref() { + Some(capture_body_json) => (body_json.or(Some(capture_body_json)), None), + None => (body_json, body_base64), + }; + if let Some(body_json) = maybe_build_standard_same_format_sync_body_from_normalized_payload( report_kind, status_code, @@ -1011,6 +1025,48 @@ fn maybe_build_openai_cross_format_provider_body_from_normalized_payload( })) } +fn sync_finalize_needs_conversion(report_context: Option<&Value>) -> bool { + report_context + .and_then(|report_context| report_context.get("needs_conversion")) + .and_then(Value::as_bool) + .unwrap_or(false) +} + +fn decode_non_stream_sync_capture_body(body_base64: &str) -> Option { + let body_bytes = base64::engine::general_purpose::STANDARD + .decode(body_base64) + .ok()?; + serde_json::from_slice::(&body_bytes) + .ok() + .filter(Value::is_object) + .filter(|body_json| !is_stream_event_object(body_json)) +} + +/// Unframed JSON events are accepted by the stream parsers and must not be mistaken for complete +/// provider response bodies merely because the entire capture parses as one JSON object. +fn is_stream_event_object(value: &Value) -> bool { + let Some(object) = value.as_object() else { + return false; + }; + if object + .get("object") + .and_then(Value::as_str) + .is_some_and(|object| object.ends_with(".chunk")) + { + return true; + } + + object + .get("type") + .and_then(Value::as_str) + .is_some_and(|event_type| { + event_type.contains('.') + || ["response", "message", "item", "delta", "content_block"] + .iter() + .any(|nested| object.contains_key(*nested)) + }) +} + fn is_error_like_sync_body(value: &Value) -> bool { let Some(object) = value.as_object() else { return false; @@ -3979,7 +4035,8 @@ mod tests { aggregate_claude_stream_sync_response, aggregate_gemini_stream_sync_response, aggregate_openai_chat_stream_sync_response, aggregate_openai_responses_stream_sync_response, convert_standard_chat_response, - convert_standard_cli_response, materialize_openai_responses_reasoning_item, + convert_standard_cli_response, decode_non_stream_sync_capture_body, + materialize_openai_responses_reasoning_item, maybe_build_openai_chat_cross_format_sync_product_from_normalized_payload, maybe_build_openai_responses_cross_format_sync_product_from_normalized_payload, maybe_build_openai_responses_same_family_sync_body_from_normalized_payload, @@ -4527,6 +4584,132 @@ mod tests { ); } + #[test] + fn unframed_stream_events_are_not_mistaken_for_provider_bodies() { + for event in [ + json!({"type": "response.completed", "response": {"status": "completed"}}), + json!({"type": "response.output_text.delta", "delta": "hi"}), + json!({"type": "message_start", "message": {"id": "msg_1"}}), + json!({"type": "content_block_delta", "index": 0, "delta": {"text": "hi"}}), + json!({"object": "chat.completion.chunk", "choices": []}), + ] { + let body_base64 = base64::engine::general_purpose::STANDARD + .encode(serde_json::to_vec(&event).expect("serialize event")); + assert!( + decode_non_stream_sync_capture_body(&body_base64).is_none(), + "stream events belong to the aggregators: {event}" + ); + } + } + + #[test] + fn complete_provider_bodies_are_recovered_from_cross_format_captures() { + for body in [ + json!({"id": "resp_1", "object": "response", "status": "completed", "output": []}), + json!({"id": "chatcmpl_1", "object": "chat.completion", "choices": []}), + json!({"id": "msg_1", "type": "message", "role": "assistant", "content": []}), + json!({"candidates": [], "modelVersion": "probe-model"}), + ] { + let body_base64 = base64::engine::general_purpose::STANDARD + .encode(serde_json::to_vec(&body).expect("serialize provider body")); + assert_eq!( + decode_non_stream_sync_capture_body(&body_base64), + Some(body.clone()), + "a complete provider body is not a stream: {body}" + ); + } + } + + #[test] + fn recovers_cross_format_capture_that_is_a_complete_json_body() { + let report_context = json!({ + "provider_api_format": "openai:responses", + "client_api_format": "claude:messages", + "needs_conversion": true, + "upstream_is_stream": true, + }); + let provider_body_json = json!({ + "id": "resp_1", + "object": "response", + "status": "completed", + "error": null, + "model": "probe-model", + "output": [{ + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "hello"}] + }], + "usage": {"input_tokens": 5, "output_tokens": 7, "total_tokens": 12} + }); + let body_base64 = base64::engine::general_purpose::STANDARD + .encode(serde_json::to_vec(&provider_body_json).expect("serialize provider body")); + + let product = maybe_build_standard_sync_finalize_product_from_normalized_payload( + "claude_chat_sync_finalize", + 200, + Some(&report_context), + None, + Some(&body_base64), + ) + .expect("a complete provider body must not fail the stream aggregator") + .expect("product should exist"); + + let StandardSyncFinalizeNormalizedProduct::CrossFormat(product) = product else { + panic!("cross-format attempt should produce a cross-format product"); + }; + assert_eq!(product.provider_body_json, provider_body_json); + assert_eq!(product.client_body_json["type"], "message"); + assert_eq!(product.client_body_json["content"][0]["text"], "hello"); + } + + #[test] + fn keeps_unframed_stream_event_on_the_aggregation_path() { + let report_context = json!({ + "provider_api_format": "openai:responses", + "client_api_format": "claude:messages", + "needs_conversion": true, + "upstream_is_stream": true, + }); + let provider_body_json = json!({ + "id": "resp_1", + "object": "response", + "status": "completed", + "model": "probe-model", + "output": [{ + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "hello"}] + }], + "usage": {"input_tokens": 5, "output_tokens": 7, "total_tokens": 12} + }); + let event = json!({ + "type": "response.completed", + "response": provider_body_json.clone(), + }); + let body_base64 = base64::engine::general_purpose::STANDARD + .encode(serde_json::to_vec(&event).expect("serialize stream event")); + + let product = maybe_build_standard_sync_finalize_product_from_normalized_payload( + "claude_chat_sync_finalize", + 200, + Some(&report_context), + None, + Some(&body_base64), + ) + .expect("unframed stream event should aggregate") + .expect("product should exist"); + + let StandardSyncFinalizeNormalizedProduct::CrossFormat(product) = product else { + panic!("cross-format attempt should produce a cross-format product"); + }; + assert_eq!(product.provider_body_json["id"], provider_body_json["id"]); + assert_eq!(product.provider_body_json["object"], "response"); + assert!(product.provider_body_json.get("response").is_none()); + assert_eq!(product.client_body_json["type"], "message"); + } + #[test] fn builds_standard_same_format_body_from_stream_payload() { let body = concat!(