diff --git a/Cargo.lock b/Cargo.lock index 6808176ba..4a15c563f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -103,6 +103,7 @@ dependencies = [ "aether-pool-core", "aether-scheduler-core", "async-trait", + "base64 0.22.1", "http", "serde", "serde_json", diff --git a/apps/aether-gateway/src/ai_serving/api.rs b/apps/aether-gateway/src/ai_serving/api.rs index 07d622d70..363b708e4 100644 --- a/apps/aether-gateway/src/ai_serving/api.rs +++ b/apps/aether-gateway/src/ai_serving/api.rs @@ -87,28 +87,36 @@ pub(crate) fn resolve_execution_runtime_stream_plan_kind( parts: &http::request::Parts, decision: &GatewayControlDecision, ) -> Option<&'static str> { - aether_ai_formats::api::resolve_execution_runtime_stream_plan_kind( - decision.route_class.as_deref(), - decision.route_family.as_deref(), - decision.route_kind.as_deref(), - decision.request_auth_channel.as_deref(), - &parts.method, - parts.uri.path(), - ) + let plan_kind = + aether_ai_formats::api::resolve_execution_runtime_stream_plan_kind_with_client_surface( + decision.route_class.as_deref(), + decision.route_family.as_deref(), + decision.route_kind.as_deref(), + decision.client_surface, + decision.request_auth_channel.as_deref(), + &parts.method, + parts.uri.path(), + )?; + crate::ai_serving::plan_kind_matches_api_operation(plan_kind, true, decision.api_operation) + .then_some(plan_kind) } pub(crate) fn resolve_execution_runtime_sync_plan_kind( parts: &http::request::Parts, decision: &GatewayControlDecision, ) -> Option<&'static str> { - aether_ai_formats::api::resolve_execution_runtime_sync_plan_kind( - decision.route_class.as_deref(), - decision.route_family.as_deref(), - decision.route_kind.as_deref(), - decision.request_auth_channel.as_deref(), - &parts.method, - parts.uri.path(), - ) + let plan_kind = + aether_ai_formats::api::resolve_execution_runtime_sync_plan_kind_with_client_surface( + decision.route_class.as_deref(), + decision.route_family.as_deref(), + decision.route_kind.as_deref(), + decision.client_surface, + decision.request_auth_channel.as_deref(), + &parts.method, + parts.uri.path(), + )?; + crate::ai_serving::plan_kind_matches_api_operation(plan_kind, false, decision.api_operation) + .then_some(plan_kind) } pub(crate) fn is_matching_stream_request( diff --git a/apps/aether-gateway/src/ai_serving/finalize/tests_stream.rs b/apps/aether-gateway/src/ai_serving/finalize/tests_stream.rs index b36ff32eb..0de41f894 100644 --- a/apps/aether-gateway/src/ai_serving/finalize/tests_stream.rs +++ b/apps/aether-gateway/src/ai_serving/finalize/tests_stream.rs @@ -13,6 +13,7 @@ fn same_format_claude_local_stream_rewriter_sanitizes_read_input_json_delta() { let report_context = json!({ "provider_api_format": "claude:messages", "client_api_format": "claude:messages", + "anthropic_compatibility_profile": "claude_code_legacy", "needs_conversion": false, }); let mut rewriter = 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 d1f0d7807..8b346e5a6 100644 --- a/apps/aether-gateway/src/ai_serving/finalize/tests_sync.rs +++ b/apps/aether-gateway/src/ai_serving/finalize/tests_sync.rs @@ -25,6 +25,9 @@ fn test_decision() -> GatewayControlDecision { route_class: Some("ai_public".to_string()), route_family: Some("openai".to_string()), route_kind: Some("compact".to_string()), + client_surface: None, + api_operation: None, + gateway_credential_carrier: None, request_auth_channel: None, auth_endpoint_signature: Some("openai:responses:compact".to_string()), execution_runtime_candidate: true, @@ -1923,6 +1926,9 @@ fn local_finalize_handles_claude_chat_cross_format_sync_response_from_openai_cha route_class: Some("ai_public".to_string()), route_family: Some("claude".to_string()), route_kind: Some("chat".to_string()), + client_surface: None, + api_operation: None, + gateway_credential_carrier: None, request_auth_channel: None, auth_endpoint_signature: Some("claude:messages".to_string()), execution_runtime_candidate: true, @@ -1991,6 +1997,9 @@ fn local_finalize_handles_gemini_cli_cross_format_sync_response_from_claude_cli( route_class: Some("ai_public".to_string()), route_family: Some("gemini".to_string()), route_kind: Some("cli".to_string()), + client_surface: None, + api_operation: None, + gateway_credential_carrier: None, request_auth_channel: None, auth_endpoint_signature: Some("gemini:generate_content".to_string()), execution_runtime_candidate: true, diff --git a/apps/aether-gateway/src/ai_serving/mod.rs b/apps/aether-gateway/src/ai_serving/mod.rs index 835e27f88..fb696fd3f 100644 --- a/apps/aether-gateway/src/ai_serving/mod.rs +++ b/apps/aether-gateway/src/ai_serving/mod.rs @@ -71,7 +71,7 @@ pub(crate) use self::transport::{ request_pair_allowed_for_transport, request_pair_direct_auth, request_pair_transport_unsupported_reason, CandidateTransportPolicyFacts, }; -pub(crate) use crate::control::GatewayControlDecision; +pub(crate) use crate::control::{GatewayControlDecision, GatewayCredentialCarrier}; pub(crate) use crate::execution_runtime::{ConversionMode, ExecutionStrategy}; pub(crate) use crate::headers::RequestOrigin; pub(crate) use aether_ai_serving::{ @@ -89,6 +89,7 @@ pub(crate) fn build_provider_transport_request_url( upstream_is_stream: bool, request_query: Option<&str>, kiro_api_region: Option<&str>, + api_operation: Option, ) -> Option { self::transport::build_transport_request_url( transport, @@ -98,6 +99,7 @@ pub(crate) fn build_provider_transport_request_url( upstream_is_stream, request_query, kiro_api_region, + api_operation, }, ) } @@ -109,6 +111,7 @@ pub(crate) fn build_provider_transport_request_url_for_request_body( upstream_is_stream: bool, request_query: Option<&str>, kiro_api_region: Option<&str>, + api_operation: Option, provider_request_body: Option<&serde_json::Value>, ) -> Option { self::transport::build_transport_request_url_for_request_body( @@ -119,6 +122,7 @@ pub(crate) fn build_provider_transport_request_url_for_request_body( upstream_is_stream, request_query, kiro_api_region, + api_operation, }, provider_request_body, ) diff --git a/apps/aether-gateway/src/ai_serving/planner/antigravity.rs b/apps/aether-gateway/src/ai_serving/planner/antigravity.rs index 6dc684ba1..d9f2de062 100644 --- a/apps/aether-gateway/src/ai_serving/planner/antigravity.rs +++ b/apps/aether-gateway/src/ai_serving/planner/antigravity.rs @@ -67,6 +67,7 @@ pub(crate) async fn build_antigravity_v1internal_provider_request( input.upstream_is_stream, input.parts.uri.query(), None, + None, Some(&payload.body), ) .ok_or(AntigravityV1InternalRequestError::UpstreamUrlUnavailable)?; diff --git a/apps/aether-gateway/src/ai_serving/planner/candidate_materialization.rs b/apps/aether-gateway/src/ai_serving/planner/candidate_materialization.rs index 6f62f21c6..384b5644c 100644 --- a/apps/aether-gateway/src/ai_serving/planner/candidate_materialization.rs +++ b/apps/aether-gateway/src/ai_serving/planner/candidate_materialization.rs @@ -68,6 +68,8 @@ pub(crate) struct LocalExecutionCandidateAttempt { pub(crate) struct LocalExecutionCandidateAttemptSource<'a> { items: VecDeque>, skipped_provider_ids: BTreeSet, + skipped_endpoint_ids: BTreeSet, + skipped_credential_ids: BTreeSet, } type DecorateSkippedCandidateFn<'a> = Arc< @@ -80,6 +82,10 @@ pub(crate) trait LocalExecutionAttemptSource: Send { async fn drain_execution_attempts(&mut self) -> Result, GatewayError>; + async fn skip_credential(&mut self, key_id: &str) -> Result<(), GatewayError>; + + async fn skip_endpoint(&mut self, endpoint_id: &str) -> Result<(), GatewayError>; + async fn skip_provider(&mut self, provider_id: &str) -> Result<(), GatewayError>; } @@ -111,6 +117,8 @@ impl<'a> LocalExecutionCandidateAttemptSource<'a> { Self { items, skipped_provider_ids: BTreeSet::new(), + skipped_endpoint_ids: BTreeSet::new(), + skipped_credential_ids: BTreeSet::new(), } } @@ -123,7 +131,12 @@ impl<'a> LocalExecutionCandidateAttemptSource<'a> { }; match front { LocalExecutionCandidateAttemptSourceItem::Static { attempts } => { - if dispatch_sequence_provider_is_skipped(attempts, &self.skipped_provider_ids) { + if dispatch_sequence_candidate_is_skipped( + attempts, + &self.skipped_provider_ids, + &self.skipped_endpoint_ids, + &self.skipped_credential_ids, + ) { self.items.pop_front(); continue; } @@ -141,10 +154,20 @@ impl<'a> LocalExecutionCandidateAttemptSource<'a> { pending_attempts, pool_exhaustion_persistence, } => { - if self.skipped_provider_ids.contains(cursor.provider_id()) { + if self.skipped_provider_ids.contains(cursor.provider_id()) + || self.skipped_endpoint_ids.contains(cursor.endpoint_id()) + { self.items.pop_front(); continue; } + if dispatch_sequence_candidate_is_skipped( + pending_attempts, + &self.skipped_provider_ids, + &self.skipped_endpoint_ids, + &self.skipped_credential_ids, + ) { + *pending_attempts = DispatchSequence::new(Vec::new()); + } if let Some(attempt) = next_attempt_from_dispatch_sequence(pending_attempts) { return Ok(Some(attempt)); } @@ -162,6 +185,14 @@ impl<'a> LocalExecutionCandidateAttemptSource<'a> { self.items.pop_front(); continue; }; + if candidate_is_skipped( + &candidate, + &self.skipped_provider_ids, + &self.skipped_endpoint_ids, + &self.skipped_credential_ids, + ) { + continue; + } *pending_attempts = dispatch_sequence_from_attempts( build_unpersisted_local_execution_candidate_attempts( candidate, @@ -174,6 +205,12 @@ impl<'a> LocalExecutionCandidateAttemptSource<'a> { for provider_id in &self.skipped_provider_ids { cursor.skip_provider(provider_id); } + for endpoint_id in &self.skipped_endpoint_ids { + cursor.skip_endpoint(endpoint_id); + } + for key_id in &self.skipped_credential_ids { + cursor.skip_credential(key_id); + } let Some(attempt) = cursor.next_attempt().await? else { self.items.pop_front(); continue; @@ -201,6 +238,32 @@ impl<'a> LocalExecutionCandidateAttemptSource<'a> { } } } + + pub(crate) fn skip_endpoint(&mut self, endpoint_id: &str) { + let endpoint_id = endpoint_id.trim(); + if endpoint_id.is_empty() { + return; + } + self.skipped_endpoint_ids.insert(endpoint_id.to_string()); + for item in &mut self.items { + if let LocalExecutionCandidateAttemptSourceItem::RequestedModelPage { cursor } = item { + cursor.skip_endpoint(endpoint_id); + } + } + } + + pub(crate) fn skip_credential(&mut self, key_id: &str) { + let key_id = key_id.trim(); + if key_id.is_empty() { + return; + } + self.skipped_credential_ids.insert(key_id.to_string()); + for item in &mut self.items { + if let LocalExecutionCandidateAttemptSourceItem::RequestedModelPage { cursor } = item { + cursor.skip_credential(key_id); + } + } + } } impl LocalExecutionCandidateAttempt { @@ -668,6 +731,8 @@ where LocalExecutionCandidateAttemptSource { items, skipped_provider_ids: BTreeSet::new(), + skipped_endpoint_ids: BTreeSet::new(), + skipped_credential_ids: BTreeSet::new(), }, candidate_count, ) @@ -802,6 +867,8 @@ where page_cursor, pending_items: VecDeque::new(), skipped_provider_ids: BTreeSet::new(), + skipped_endpoint_ids: BTreeSet::new(), + skipped_credential_ids: BTreeSet::new(), candidate_count: 0, next_candidate_index: 0, remembered_affinity: false, @@ -825,6 +892,8 @@ where LocalExecutionCandidateAttemptSource { items, skipped_provider_ids: BTreeSet::new(), + skipped_endpoint_ids: BTreeSet::new(), + skipped_credential_ids: BTreeSet::new(), }, candidate_count, ) @@ -851,6 +920,8 @@ struct RequestedModelAttemptPageCursor<'a> { page_cursor: LocalCandidatePreselectionPageCursor<'a>, pending_items: VecDeque>, skipped_provider_ids: BTreeSet, + skipped_endpoint_ids: BTreeSet, + skipped_credential_ids: BTreeSet, candidate_count: usize, next_candidate_index: u32, remembered_affinity: bool, @@ -864,6 +935,14 @@ impl<'a> RequestedModelAttemptPageCursor<'a> { self.skipped_provider_ids.insert(provider_id.to_string()); } + fn skip_endpoint(&mut self, endpoint_id: &str) { + self.skipped_endpoint_ids.insert(endpoint_id.to_string()); + } + + fn skip_credential(&mut self, key_id: &str) { + self.skipped_credential_ids.insert(key_id.to_string()); + } + async fn next_attempt( &mut self, ) -> Result, GatewayError> { @@ -871,8 +950,13 @@ impl<'a> RequestedModelAttemptPageCursor<'a> { return Err(error); } loop { - if let Some(attempt) = - pop_attempt_from_items(&mut self.pending_items, &self.skipped_provider_ids).await + if let Some(attempt) = pop_attempt_from_items( + &mut self.pending_items, + &self.skipped_provider_ids, + &self.skipped_endpoint_ids, + &self.skipped_credential_ids, + ) + .await { return Ok(Some(attempt)); } @@ -1075,12 +1159,19 @@ fn page_is_exact_auth_api_key_concurrency_limited( async fn pop_attempt_from_items( items: &mut VecDeque>, skipped_provider_ids: &BTreeSet, + skipped_endpoint_ids: &BTreeSet, + skipped_credential_ids: &BTreeSet, ) -> Option { loop { let front = items.front_mut()?; match front { LocalExecutionCandidateAttemptSourceItem::Static { attempts } => { - if dispatch_sequence_provider_is_skipped(attempts, skipped_provider_ids) { + if dispatch_sequence_candidate_is_skipped( + attempts, + skipped_provider_ids, + skipped_endpoint_ids, + skipped_credential_ids, + ) { items.pop_front(); continue; } @@ -1098,10 +1189,20 @@ async fn pop_attempt_from_items( pending_attempts, pool_exhaustion_persistence, } => { - if skipped_provider_ids.contains(cursor.provider_id()) { + if skipped_provider_ids.contains(cursor.provider_id()) + || skipped_endpoint_ids.contains(cursor.endpoint_id()) + { items.pop_front(); continue; } + if dispatch_sequence_candidate_is_skipped( + pending_attempts, + skipped_provider_ids, + skipped_endpoint_ids, + skipped_credential_ids, + ) { + *pending_attempts = DispatchSequence::new(Vec::new()); + } if let Some(attempt) = next_attempt_from_dispatch_sequence(pending_attempts) { return Some(attempt); } @@ -1119,6 +1220,14 @@ async fn pop_attempt_from_items( items.pop_front(); continue; }; + if candidate_is_skipped( + &candidate, + skipped_provider_ids, + skipped_endpoint_ids, + skipped_credential_ids, + ) { + continue; + } *pending_attempts = dispatch_sequence_from_attempts( build_unpersisted_local_execution_candidate_attempts( candidate, @@ -1780,15 +1889,33 @@ fn next_attempt_from_dispatch_sequence( Some(attempt) } -fn dispatch_sequence_provider_is_skipped( +fn dispatch_sequence_candidate_is_skipped( sequence: &DispatchSequence, skipped_provider_ids: &BTreeSet, + skipped_endpoint_ids: &BTreeSet, + skipped_credential_ids: &BTreeSet, ) -> bool { sequence.peek_current().is_some_and(|item| { - skipped_provider_ids.contains(&item.candidate.eligible.candidate.provider_id) + candidate_is_skipped( + &item.candidate.eligible, + skipped_provider_ids, + skipped_endpoint_ids, + skipped_credential_ids, + ) }) } +fn candidate_is_skipped( + candidate: &EligibleLocalExecutionCandidate, + skipped_provider_ids: &BTreeSet, + skipped_endpoint_ids: &BTreeSet, + skipped_credential_ids: &BTreeSet, +) -> bool { + skipped_provider_ids.contains(&candidate.candidate.provider_id) + || skipped_endpoint_ids.contains(&candidate.candidate.endpoint_id) + || skipped_credential_ids.contains(&candidate.candidate.key_id) +} + fn dispatch_sequence_exhausted( sequence: &mut DispatchSequence, ) -> bool { @@ -2341,6 +2468,8 @@ mod tests { page_cursor, pending_items: VecDeque::new(), skipped_provider_ids: BTreeSet::new(), + skipped_endpoint_ids: BTreeSet::new(), + skipped_credential_ids: BTreeSet::new(), candidate_count: 0, next_candidate_index: 0, remembered_affinity: false, @@ -2436,6 +2565,8 @@ mod tests { page_cursor, pending_items: VecDeque::new(), skipped_provider_ids: BTreeSet::new(), + skipped_endpoint_ids: BTreeSet::new(), + skipped_credential_ids: BTreeSet::new(), candidate_count: 0, next_candidate_index: 0, remembered_affinity: false, @@ -2634,6 +2765,8 @@ mod tests { ), }]), skipped_provider_ids: BTreeSet::new(), + skipped_endpoint_ids: BTreeSet::new(), + skipped_credential_ids: BTreeSet::new(), }; let first = source @@ -2652,6 +2785,111 @@ mod tests { .is_none()); } + #[tokio::test] + async fn dynamic_attempt_source_skips_credentials_and_endpoints_across_static_candidates() { + let key_a = sample_eligible("key-a", None); + let key_b = sample_eligible("key-b", None); + let mut key_c = sample_eligible("key-c", None); + key_c.candidate.endpoint_id = "endpoint-2".to_string(); + Arc::make_mut(&mut key_c.transport).endpoint.id = "endpoint-2".to_string(); + + let static_item = + |candidate, candidate_index| LocalExecutionCandidateAttemptSourceItem::Static { + attempts: dispatch_sequence_from_attempts( + build_unpersisted_local_execution_candidate_attempts( + candidate, + candidate_index, + ) + .into(), + ), + }; + let mut source = LocalExecutionCandidateAttemptSource { + items: VecDeque::from([ + static_item(key_a, 0), + static_item(key_b, 1), + static_item(key_c, 2), + ]), + skipped_provider_ids: BTreeSet::new(), + skipped_endpoint_ids: BTreeSet::new(), + skipped_credential_ids: BTreeSet::new(), + }; + + source.skip_credential("key-a"); + let key_b_attempt = source + .next_attempt() + .await + .expect("candidate source should succeed") + .expect("a different credential should remain"); + assert_eq!(key_b_attempt.eligible.candidate.key_id, "key-b"); + + source.skip_endpoint("endpoint-1"); + let endpoint_2_attempt = source + .next_attempt() + .await + .expect("candidate source should succeed") + .expect("a different endpoint should remain"); + assert_eq!(endpoint_2_attempt.eligible.candidate.key_id, "key-c"); + assert_eq!( + endpoint_2_attempt.eligible.candidate.endpoint_id, + "endpoint-2" + ); + } + + #[tokio::test] + async fn dynamic_attempt_source_filters_skipped_pool_pending_credential() { + let app = AppState::new().expect("state should build"); + let mut pool_group = sample_eligible("pool-group", None); + pool_group.kind = LocalExecutionCandidateKind::PoolGroup; + pool_group.transport = sample_transport("pool-group", Some(json!({ "pool_advanced": {} }))); + let pool_cursor = PoolKeyCursor::new( + PlannerAppState::new(&app), + pool_group, + None, + Some("gpt-5"), + None, + ); + let pool_key_attempts = dispatch_sequence_from_attempts( + build_unpersisted_local_execution_candidate_attempts( + sample_eligible("pool-key-a", None), + 0, + ) + .into(), + ); + let mut fallback = sample_eligible("fallback-key", None); + fallback.candidate.provider_id = "provider-b".to_string(); + Arc::make_mut(&mut fallback.transport).provider.id = "provider-b".to_string(); + Arc::make_mut(&mut fallback.transport).key.provider_id = "provider-b".to_string(); + let fallback_attempts = dispatch_sequence_from_attempts( + build_unpersisted_local_execution_candidate_attempts(fallback, 1).into(), + ); + let mut source = LocalExecutionCandidateAttemptSource { + items: VecDeque::from([ + LocalExecutionCandidateAttemptSourceItem::Pool { + cursor: pool_cursor, + candidate_index: 0, + pending_attempts: pool_key_attempts, + pool_exhaustion_persistence: None, + }, + LocalExecutionCandidateAttemptSourceItem::Static { + attempts: fallback_attempts, + }, + ]), + skipped_provider_ids: BTreeSet::new(), + skipped_endpoint_ids: BTreeSet::new(), + skipped_credential_ids: BTreeSet::new(), + }; + + source.skip_credential("pool-key-a"); + let attempt = source + .next_attempt() + .await + .expect("candidate source should succeed") + .expect("fallback credential should remain"); + + assert_eq!(attempt.eligible.candidate.provider_id, "provider-b"); + assert_eq!(attempt.eligible.candidate.key_id, "fallback-key"); + } + #[tokio::test] async fn skipped_provider_discards_pool_cursor_and_continues_with_next_provider() { let app = AppState::new().expect("state should build"); @@ -2686,6 +2924,8 @@ mod tests { }, ]), skipped_provider_ids: BTreeSet::new(), + skipped_endpoint_ids: BTreeSet::new(), + skipped_credential_ids: BTreeSet::new(), }; source.skip_provider("provider-1"); @@ -2752,6 +2992,8 @@ mod tests { pool_exhaustion_persistence: Some(pool_exhaustion_persistence), }]), skipped_provider_ids: BTreeSet::new(), + skipped_endpoint_ids: BTreeSet::new(), + skipped_credential_ids: BTreeSet::new(), }; assert!(source diff --git a/apps/aether-gateway/src/ai_serving/planner/common.rs b/apps/aether-gateway/src/ai_serving/planner/common.rs index 33565a583..9bcbdf7b8 100644 --- a/apps/aether-gateway/src/ai_serving/planner/common.rs +++ b/apps/aether-gateway/src/ai_serving/planner/common.rs @@ -10,7 +10,7 @@ use crate::ai_serving::{ }; pub(crate) use crate::ai_serving::{ CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND, - CLAUDE_CLI_SYNC_PLAN_KIND, EXECUTION_RUNTIME_STREAM_ACTION, + CLAUDE_CLI_SYNC_PLAN_KIND, CLAUDE_COUNT_TOKENS_SYNC_PLAN_KIND, EXECUTION_RUNTIME_STREAM_ACTION, EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_ACTION, EXECUTION_RUNTIME_SYNC_DECISION_ACTION, GEMINI_CHAT_STREAM_PLAN_KIND, GEMINI_CHAT_SYNC_PLAN_KIND, GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_SYNC_PLAN_KIND, diff --git a/apps/aether-gateway/src/ai_serving/planner/decision/control_plan.rs b/apps/aether-gateway/src/ai_serving/planner/decision/control_plan.rs index 276df03c0..a101dfdf5 100644 --- a/apps/aether-gateway/src/ai_serving/planner/decision/control_plan.rs +++ b/apps/aether-gateway/src/ai_serving/planner/decision/control_plan.rs @@ -1,17 +1,17 @@ use crate::ai_serving::planner::common::{ CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND, - CLAUDE_CLI_SYNC_PLAN_KIND, GEMINI_CHAT_STREAM_PLAN_KIND, GEMINI_CHAT_SYNC_PLAN_KIND, - GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_SYNC_PLAN_KIND, GEMINI_EMBEDDING_SYNC_PLAN_KIND, - GEMINI_FILES_DELETE_PLAN_KIND, GEMINI_FILES_DOWNLOAD_PLAN_KIND, GEMINI_FILES_GET_PLAN_KIND, - GEMINI_FILES_LIST_PLAN_KIND, GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND, - GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND, - OPENAI_EMBEDDING_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND, - OPENAI_RERANK_SYNC_PLAN_KIND, OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, - OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND, OPENAI_RESPONSES_STREAM_PLAN_KIND, - OPENAI_RESPONSES_SYNC_PLAN_KIND, OPENAI_SEARCH_SYNC_PLAN_KIND, - OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_VIDEO_CONTENT_PLAN_KIND, - OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND, - OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND, + CLAUDE_CLI_SYNC_PLAN_KIND, CLAUDE_COUNT_TOKENS_SYNC_PLAN_KIND, GEMINI_CHAT_STREAM_PLAN_KIND, + GEMINI_CHAT_SYNC_PLAN_KIND, GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_SYNC_PLAN_KIND, + GEMINI_EMBEDDING_SYNC_PLAN_KIND, GEMINI_FILES_DELETE_PLAN_KIND, + GEMINI_FILES_DOWNLOAD_PLAN_KIND, GEMINI_FILES_GET_PLAN_KIND, GEMINI_FILES_LIST_PLAN_KIND, + GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND, GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND, + OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND, OPENAI_EMBEDDING_SYNC_PLAN_KIND, + OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND, OPENAI_RERANK_SYNC_PLAN_KIND, + OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND, + OPENAI_RESPONSES_STREAM_PLAN_KIND, OPENAI_RESPONSES_SYNC_PLAN_KIND, + OPENAI_SEARCH_SYNC_PLAN_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, + OPENAI_VIDEO_CONTENT_PLAN_KIND, OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND, + OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND, OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND, }; use crate::ai_serving::planner::plan_builders::{ build_gemini_stream_plan_from_decision, build_gemini_sync_plan_from_decision, @@ -109,6 +109,7 @@ fn build_sync_plan_payload_from_decision( } CLAUDE_CHAT_SYNC_PLAN_KIND | CLAUDE_CLI_SYNC_PLAN_KIND + | CLAUDE_COUNT_TOKENS_SYNC_PLAN_KIND | OPENAI_EMBEDDING_SYNC_PLAN_KIND | OPENAI_RERANK_SYNC_PLAN_KIND => { build_standard_sync_plan_from_decision(parts, body_json, payload)? diff --git a/apps/aether-gateway/src/ai_serving/planner/decision_input.rs b/apps/aether-gateway/src/ai_serving/planner/decision_input.rs index 272a6a8e9..aff69b55f 100644 --- a/apps/aether-gateway/src/ai_serving/planner/decision_input.rs +++ b/apps/aether-gateway/src/ai_serving/planner/decision_input.rs @@ -15,8 +15,9 @@ use tracing::warn; use crate::ai_serving::planner::common::extract_standard_requested_model; use crate::ai_serving::{ - ExecutionRuntimeAuthContext, GatewayAuthApiKeySnapshot, GatewayProviderTransportSnapshot, - PlannerAppState, CODEX_RESPONSES_LITE_HEADER, + ClientSurface, ExecutionRuntimeAuthContext, GatewayAuthApiKeySnapshot, + GatewayCredentialCarrier, GatewayProviderTransportSnapshot, PlannerAppState, + CODEX_RESPONSES_LITE_HEADER, }; use crate::cache::CacheLoadObserver; use crate::client_session_affinity::client_session_affinity_from_api_request; @@ -52,6 +53,8 @@ pub(crate) struct LocalRequestedModelDecisionInput { pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot, pub(crate) required_capabilities: Option, pub(crate) request_auth_channel: Option, + pub(crate) client_surface: Option, + pub(crate) gateway_credential_carrier: Option, pub(crate) client_session_affinity: Option, pub(crate) routing_policy: Option, pub(crate) routing_trace_seed: Option, @@ -378,6 +381,8 @@ pub(crate) fn build_local_requested_model_decision_input( auth_snapshot: resolved_input.auth_snapshot, required_capabilities: resolved_input.required_capabilities, request_auth_channel: None, + client_surface: None, + gateway_credential_carrier: None, client_session_affinity: None, routing_policy: None, routing_trace_seed: None, @@ -1128,6 +1133,8 @@ mod tests { auth_snapshot: sample_auth_snapshot(), required_capabilities: None, request_auth_channel: None, + client_surface: None, + gateway_credential_carrier: None, client_session_affinity: None, routing_policy: None, routing_trace_seed: None, @@ -1323,6 +1330,8 @@ mod tests { auth_snapshot: sample_auth_snapshot(), required_capabilities: None, request_auth_channel: None, + client_surface: None, + gateway_credential_carrier: None, client_session_affinity: None, routing_policy: None, routing_trace_seed: None, @@ -1390,6 +1399,8 @@ mod tests { auth_snapshot: sample_auth_snapshot(), required_capabilities: None, request_auth_channel: None, + client_surface: None, + gateway_credential_carrier: None, client_session_affinity: None, routing_policy: None, routing_trace_seed: None, @@ -1459,6 +1470,59 @@ mod tests { ); } + #[test] + fn provider_request_routing_policy_cannot_restore_credentials_or_aether_internal_headers() { + for header_name in [ + "authorization", + "proxy-authorization", + "api-key", + "x-api-key", + "x-goog-api-key", + "cookie", + "cookie2", + "set-cookie", + "x-aether-auth-user-id", + "x-aether-control-future", + ] { + let mut input = sample_decision_input(); + set_provider_request_rules( + &mut input, + &["gpt-5"], + json!([{ + "type": "patch_headers", + "patch": [{ + "op": "set", + "name": header_name, + "value": "must-not-reach-upstream" + }] + }]), + ); + let mut decision = sample_decision(); + + let error = + apply_provider_request_routing_policy_to_decision(&input, &mut decision, None) + .expect_err("reserved provider header mutation should fail closed"); + + assert!( + matches!( + &error, + GatewayError::Client { + status: StatusCode::BAD_REQUEST, + .. + } + ), + "unexpected error for {header_name}: {error:?}" + ); + assert!( + !decision + .provider_request_headers + .keys() + .any(|name| name.eq_ignore_ascii_case(header_name)), + "reserved header reached the provider decision: {header_name}" + ); + } + } + #[test] fn codex_prompt_cache_identity_headers_are_terminal_after_routing_mutations() { let mut input = sample_decision_input(); diff --git a/apps/aether-gateway/src/ai_serving/planner/gemini_cli.rs b/apps/aether-gateway/src/ai_serving/planner/gemini_cli.rs index 95de6aba3..2085cb7f6 100644 --- a/apps/aether-gateway/src/ai_serving/planner/gemini_cli.rs +++ b/apps/aether-gateway/src/ai_serving/planner/gemini_cli.rs @@ -59,6 +59,7 @@ pub(crate) async fn build_gemini_cli_v1internal_provider_request( input.upstream_is_stream, input.parts.uri.query(), None, + None, Some(&payload.body), ) .ok_or(GeminiCliV1InternalRequestError::UpstreamUrlUnavailable)?; diff --git a/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/candidates.rs b/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/candidates.rs index fb9d7a102..dd00705f9 100644 --- a/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/candidates.rs +++ b/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/candidates.rs @@ -81,6 +81,8 @@ pub(crate) async fn resolve_local_same_format_provider_decision_input( let mut input = build_local_requested_model_decision_input(resolved_input, requested_model); input.request_auth_channel = decision.request_auth_channel.clone(); + input.client_surface = decision.client_surface; + input.gateway_credential_carrier = decision.gateway_credential_carrier; input.client_session_affinity = client_session_affinity_from_api_request( spec_metadata.api_format, &parts.headers, @@ -128,7 +130,7 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts( .base_model() .unwrap_or(&input.requested_model); let (candidates, preselection_skipped) = planner_state - .list_selectable_candidates_with_skip_reasons( + .list_selectable_candidates_with_skip_reasons_for_request_operation( spec_metadata.api_format, routing_model, spec_metadata.require_streaming, @@ -137,6 +139,7 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts( input.client_session_affinity.as_ref(), current_unix_secs(), false, + spec.operation.map(|operation| operation.as_str()), ) .await?; let outcome = materialize_local_execution_candidates_with_serving( @@ -232,7 +235,7 @@ pub(crate) async fn build_local_same_format_provider_candidate_attempt_source<'a .base_model() .unwrap_or(&input.requested_model); let (candidates, preselection_skipped) = planner_state - .list_selectable_candidates_with_skip_reasons( + .list_selectable_candidates_with_skip_reasons_for_request_operation( spec_metadata.api_format, routing_model, spec_metadata.require_streaming, @@ -241,6 +244,7 @@ pub(crate) async fn build_local_same_format_provider_candidate_attempt_source<'a input.client_session_affinity.as_ref(), current_unix_secs(), false, + spec.operation.map(|operation| operation.as_str()), ) .await?; diff --git a/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/payload.rs b/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/payload.rs index 5fb16dc56..60cffac9f 100644 --- a/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/payload.rs +++ b/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/payload.rs @@ -1,5 +1,8 @@ use serde_json::json; +use aether_ai_serving::{AdaptationMode, AiRequestGzipPolicy, OriginalRequestPayload}; +use aether_contracts::{ExecutionResponseBodyMode, EXECUTION_RESPONSE_BODY_MODE_HEADER}; + use crate::ai_serving::ai_local_execution_contract_for_formats; use crate::ai_serving::build_request_trace_proxy_value; use crate::ai_serving::planner::candidate_materialization::{ @@ -61,6 +64,8 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_ else { return Ok(None); }; + let request_redacted = resolved.request_redacted; + let compatibility_edits_empty = resolved.compatibility_edits.is_empty(); let original_request_body_json = if resolved.request_redacted { Some(&resolved.provider_request_body) } else { @@ -82,6 +87,51 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_ .clone() .or_else(|| resolve_transport_profile(&resolved.transport)); let mut extra_fields = serde_json::Map::new(); + extra_fields.insert( + "provider_type".to_string(), + json!(resolved.transport.provider.provider_type.as_str()), + ); + if let Some(operation) = spec.operation { + extra_fields.insert("api_operation".to_string(), json!(operation.as_str())); + } + if let Some(client_surface) = input.client_surface { + extra_fields.insert("client_surface".to_string(), json!(client_surface.as_str())); + } + if let Some(carrier) = input.gateway_credential_carrier { + extra_fields.insert( + "gateway_credential_carrier".to_string(), + json!(carrier.as_str()), + ); + } + extra_fields.insert( + "upstream_credential_mode".to_string(), + json!(resolved.transport.key.auth_type.trim().to_ascii_lowercase()), + ); + let mut adaptation_mode = if resolved.compatibility_edits.is_empty() { + AdaptationMode::NativeTransparent + } else { + AdaptationMode::SameFormatCompat + }; + if crate::ai_serving::normalize_api_format_alias(&resolved.provider_api_format) + == "claude:messages" + { + let compatibility_profile = + crate::ai_serving::transport::resolve_anthropic_compatibility_profile( + &resolved.transport, + &resolved.provider_api_format, + ); + extra_fields.insert( + "anthropic_compatibility_profile".to_string(), + json!(compatibility_profile.as_str()), + ); + if compatibility_profile.uses_claude_code_compatibility() { + adaptation_mode = AdaptationMode::SameFormatCompat; + } + } + extra_fields.insert( + "adaptation_mode".to_string(), + json!(adaptation_mode.as_str()), + ); if let Some(proxy_value) = build_request_trace_proxy_value(Some(&resolved.transport), proxy.as_ref()) { @@ -227,9 +277,79 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_ &mut decision, Some(transport.as_ref()), )?; + enforce_provider_api_operation_invariants( + spec.operation, + decision.provider_request_body.as_mut(), + &mut decision.provider_request_headers, + ); + decision.provider_request_body_base64 = original_request_body_base64( + parts, + decision.provider_request_body.as_ref(), + adaptation_mode, + request_redacted, + compatibility_edits_empty, + decision.content_encoding.as_deref(), + decision.request_gzip.as_ref(), + ); + decision + .provider_request_headers + .retain(|name, _| !name.eq_ignore_ascii_case(EXECUTION_RESPONSE_BODY_MODE_HEADER)); + if !spec_metadata.require_streaming && decision.provider_request_body_base64.is_some() { + decision.provider_request_headers.insert( + EXECUTION_RESPONSE_BODY_MODE_HEADER.to_string(), + ExecutionResponseBodyMode::PreserveBytes + .as_str() + .to_string(), + ); + } Ok(Some(decision)) } +fn enforce_provider_api_operation_invariants( + operation: Option, + provider_request_body: Option<&mut serde_json::Value>, + provider_request_headers: &mut std::collections::BTreeMap, +) { + if operation != Some(crate::ai_serving::ApiOperation::ClaudeCountTokens) { + return; + } + + if let Some(provider_request_body) = provider_request_body { + crate::ai_serving::transport::enforce_same_format_provider_api_operation_body_policy( + provider_request_body, + operation, + ); + } + for header_name in ["accept", "content-type"] { + provider_request_headers.retain(|name, _| !name.eq_ignore_ascii_case(header_name)); + provider_request_headers.insert(header_name.to_string(), "application/json".to_string()); + } +} + +fn original_request_body_base64( + parts: &http::request::Parts, + provider_request_body: Option<&serde_json::Value>, + adaptation_mode: AdaptationMode, + request_redacted: bool, + compatibility_edits_empty: bool, + content_encoding: Option<&str>, + request_gzip: Option<&AiRequestGzipPolicy>, +) -> Option { + if adaptation_mode != AdaptationMode::NativeTransparent + || request_redacted + || !compatibility_edits_empty + || content_encoding.is_some_and(|value| !value.trim().is_empty()) + || request_gzip.is_some_and(|policy| policy.enabled != Some(false)) + { + return None; + } + + parts + .extensions + .get::()? + .body_bytes_base64_if_unchanged(provider_request_body?) +} + pub(super) async fn mark_skipped_local_same_format_provider_candidate( state: &AppState, input: &LocalSameFormatProviderDecisionInput, @@ -313,3 +433,177 @@ pub(super) async fn mark_skipped_local_same_format_provider_candidate_with_failu ) .await; } + +#[cfg(test)] +mod tests { + use std::collections::BTreeMap; + + use base64::Engine as _; + + use super::{ + enforce_provider_api_operation_invariants, original_request_body_base64, AdaptationMode, + AiRequestGzipPolicy, OriginalRequestPayload, + }; + use crate::ai_serving::ApiOperation; + + fn request_parts_with_original_payload( + body_json: serde_json::Value, + body_bytes: &[u8], + ) -> http::request::Parts { + let (mut parts, ()) = http::Request::new(()).into_parts(); + parts + .extensions + .insert(OriginalRequestPayload::from_parsed_json( + body_json, body_bytes, + )); + parts + } + + #[test] + fn count_tokens_invariants_win_after_provider_routing_mutations() { + let mut body = serde_json::json!({ + "model": "claude-sonnet-4", + "messages": [], + "stream": true + }); + let mut headers = BTreeMap::from([ + ("Accept".to_string(), "text/event-stream".to_string()), + ("Content-Type".to_string(), "text/plain".to_string()), + ("x-provider-route".to_string(), "kept".to_string()), + ]); + + enforce_provider_api_operation_invariants( + Some(ApiOperation::ClaudeCountTokens), + Some(&mut body), + &mut headers, + ); + + assert!(body.get("stream").is_none()); + assert_eq!( + headers.get("accept").map(String::as_str), + Some("application/json") + ); + assert_eq!( + headers.get("content-type").map(String::as_str), + Some("application/json") + ); + assert_eq!( + headers.get("x-provider-route").map(String::as_str), + Some("kept") + ); + assert_eq!( + headers + .keys() + .filter(|name| name.eq_ignore_ascii_case("accept")) + .count(), + 1 + ); + assert_eq!( + headers + .keys() + .filter(|name| name.eq_ignore_ascii_case("content-type")) + .count(), + 1 + ); + } + + #[test] + fn unchanged_same_format_body_preserves_original_json_bytes() { + let raw = br#"{ "unknown": {"enabled":true}, "messages": [], "model": "claude-sonnet-4" }"#; + let body_json: serde_json::Value = serde_json::from_slice(raw).expect("body should parse"); + let parts = request_parts_with_original_payload(body_json.clone(), raw); + + let encoded = original_request_body_base64( + &parts, + Some(&body_json), + AdaptationMode::NativeTransparent, + false, + true, + None, + None, + ) + .expect("unchanged request should retain exact bytes"); + + assert_eq!( + base64::engine::general_purpose::STANDARD + .decode(encoded) + .expect("body should decode"), + raw + ); + } + + #[test] + fn request_edits_or_encoding_disable_original_json_bytes() { + let raw = br#"{"model":"claude-sonnet-4","messages":[]}"#; + let body_json: serde_json::Value = serde_json::from_slice(raw).expect("body should parse"); + let parts = request_parts_with_original_payload(body_json.clone(), raw); + let changed_body = serde_json::json!({ + "model": "claude-sonnet-4-5", + "messages": [] + }); + + assert!(original_request_body_base64( + &parts, + Some(&changed_body), + AdaptationMode::NativeTransparent, + false, + true, + None, + None, + ) + .is_none()); + assert!(original_request_body_base64( + &parts, + Some(&body_json), + AdaptationMode::NativeTransparent, + true, + true, + None, + None, + ) + .is_none()); + assert!(original_request_body_base64( + &parts, + Some(&body_json), + AdaptationMode::NativeTransparent, + false, + false, + None, + None, + ) + .is_none()); + assert!(original_request_body_base64( + &parts, + Some(&body_json), + AdaptationMode::SameFormatCompat, + false, + true, + None, + None, + ) + .is_none()); + assert!(original_request_body_base64( + &parts, + Some(&body_json), + AdaptationMode::NativeTransparent, + false, + true, + Some("gzip"), + None, + ) + .is_none()); + assert!(original_request_body_base64( + &parts, + Some(&body_json), + AdaptationMode::NativeTransparent, + false, + true, + None, + Some(&AiRequestGzipPolicy { + enabled: Some(true), + min_bytes: Some(1), + }), + ) + .is_none()); + } +} diff --git a/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request.rs b/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request.rs index 848e8dd52..d2406c7bb 100644 --- a/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request.rs +++ b/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request.rs @@ -75,9 +75,10 @@ pub(crate) fn resolve_same_format_provider_transport_unsupported_reason_for_trac decision_kind: "trace_candidate_metadata", report_kind: Some("trace_candidate_metadata"), }, + None, ); if !behavior.is_antigravity - && !behavior.is_claude_code + && !behavior.is_claude_code_transport && !behavior.is_gemini_cli && !behavior.is_vertex && !behavior.is_kiro @@ -127,6 +128,23 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts( spec: LocalSameFormatProviderSpec, ) -> Result, GatewayError> { let candidate = &attempt.eligible.candidate; + if let Some(skip_reason) = same_format_provider_operation_skip_reason( + &attempt.eligible.transport, + attempt.eligible.provider_api_format.as_str(), + spec.operation, + ) { + mark_skipped_local_same_format_provider_candidate( + state, + input, + trace_id, + candidate, + attempt.candidate_index, + &attempt.candidate_id, + skip_reason, + ) + .await; + return Ok(None); + } let Some(prepared) = prepare_local_same_format_provider_candidate( state, trace_id, @@ -364,7 +382,7 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts( } else { None }; - let provider_request_body = if let Some(antigravity_auth) = antigravity_auth.as_ref() { + let mut provider_request_body = if let Some(antigravity_auth) = antigravity_auth.as_ref() { match build_antigravity_safe_v1internal_request( antigravity_auth, trace_id, @@ -424,6 +442,16 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts( } else { base_provider_request_body }; + if crate::ai_serving::transport::enforce_same_format_provider_api_operation_body_policy( + &mut provider_request_body, + spec.operation, + ) { + compatibility_edits.push(SameFormatProviderCompatibilityEdit { + field: "stream".to_string(), + action: SameFormatProviderCompatibilityEditAction::RuntimeRewrite, + detail: "removed stream field for non-streaming API operation".to_string(), + }); + } let is_grok = prepared .transport @@ -490,10 +518,10 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts( original_request_body: body_json, header_rules: transport.endpoint.header_rules.as_ref(), behavior: prepared.behavior, + api_operation: spec.operation, auth_header: prepared.auth_header.as_deref(), auth_value: prepared.auth_value.as_deref(), extra_headers: &extra_headers, - key_fingerprint: transport.key.fingerprint.as_ref(), kiro_auth_config: prepared.kiro_auth.as_ref().map(|auth| &auth.auth_config), kiro_machine_id: prepared .kiro_auth @@ -564,3 +592,98 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts( request_redacted: redaction.redacted, })) } + +fn same_format_provider_operation_skip_reason( + transport: &GatewayProviderTransportSnapshot, + provider_api_format: &str, + operation: Option, +) -> Option<&'static str> { + (!crate::ai_serving::transport::transport_supports_api_operation( + transport, + provider_api_format, + operation, + )) + .then_some("transport_operation_unsupported") +} + +#[cfg(test)] +mod tests { + use super::same_format_provider_operation_skip_reason; + use crate::ai_serving::transport::snapshot::{ + GatewayProviderTransportEndpoint, GatewayProviderTransportKey, + GatewayProviderTransportProvider, + }; + use crate::ai_serving::{ApiOperation, GatewayProviderTransportSnapshot}; + + fn private_adapter_transport(provider_type: &str) -> GatewayProviderTransportSnapshot { + GatewayProviderTransportSnapshot { + provider: GatewayProviderTransportProvider { + id: "provider-1".to_string(), + name: provider_type.to_string(), + provider_type: provider_type.to_string(), + website: None, + is_active: true, + keep_priority_on_conversion: false, + enable_format_conversion: true, + concurrent_limit: None, + max_retries: None, + proxy: None, + request_timeout_secs: None, + stream_first_byte_timeout_secs: None, + config: None, + }, + endpoint: GatewayProviderTransportEndpoint { + id: "endpoint-1".to_string(), + provider_id: "provider-1".to_string(), + api_format: "claude:messages".to_string(), + api_family: Some("claude".to_string()), + endpoint_kind: Some("chat".to_string()), + is_active: true, + base_url: "https://private.example".to_string(), + header_rules: None, + body_rules: None, + max_retries: None, + custom_path: None, + config: None, + format_acceptance_config: None, + proxy: None, + }, + key: GatewayProviderTransportKey { + id: "key-1".to_string(), + provider_id: "provider-1".to_string(), + name: "key".to_string(), + auth_type: "oauth".to_string(), + is_active: true, + api_formats: None, + auth_type_by_format: None, + allow_auth_channel_mismatch_formats: None, + allowed_models: None, + capabilities: None, + rate_multipliers: None, + global_priority_by_format: None, + expires_at_unix_secs: None, + proxy: None, + fingerprint: None, + upstream_metadata: None, + decrypted_api_key: String::new(), + decrypted_auth_config: None, + }, + } + } + + #[test] + fn private_adapter_count_tokens_is_rejected_by_pre_auth_operation_gate() { + for provider_type in ["kiro", "grok"] { + let transport = private_adapter_transport(provider_type); + assert_eq!( + same_format_provider_operation_skip_reason( + &transport, + "claude:messages", + Some(ApiOperation::ClaudeCountTokens), + ), + Some("transport_operation_unsupported"), + "provider_type={provider_type}" + ); + } + } +} diff --git a/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request/policy.rs b/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request/policy.rs index ba709cf42..ef61e8d28 100644 --- a/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request/policy.rs +++ b/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request/policy.rs @@ -1,6 +1,6 @@ use crate::ai_serving::planner::spec_metadata::LocalExecutionSurfaceSpecMetadata; use crate::ai_serving::transport::{ - classify_same_format_provider_request_behavior as classify_same_format_provider_request_behavior_impl, + classify_same_format_provider_request_behavior_for_operation as classify_same_format_provider_request_behavior_impl, resolve_same_format_provider_direct_auth as resolve_same_format_provider_direct_auth_impl, same_format_provider_transport_supported as same_format_provider_transport_supported_impl, same_format_provider_transport_unsupported_reason as same_format_provider_transport_unsupported_reason_impl, @@ -15,6 +15,7 @@ pub(super) fn classify_same_format_provider_request_behavior( transport: &GatewayProviderTransportSnapshot, provider_api_format: &str, spec_metadata: LocalExecutionSurfaceSpecMetadata, + api_operation: Option, ) -> SameFormatProviderRequestBehavior { classify_same_format_provider_request_behavior_impl( transport, @@ -25,6 +26,7 @@ pub(super) fn classify_same_format_provider_request_behavior( .report_kind .expect("same-format provider specs should declare report kind"), }, + api_operation, ) } diff --git a/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request/prepare.rs b/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request/prepare.rs index 5f993ca65..8ca86d48b 100644 --- a/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request/prepare.rs +++ b/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request/prepare.rs @@ -59,6 +59,7 @@ pub(super) async fn prepare_local_same_format_provider_candidate( &transport, provider_api_format, spec_metadata, + spec.operation, ); if !same_format_provider_transport_supported( diff --git a/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/plans.rs b/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/plans.rs index 169a147f6..fe4070b74 100644 --- a/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/plans.rs +++ b/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/plans.rs @@ -214,6 +214,16 @@ impl LocalExecutionAttemptSource for LocalSameFormatProviderSyncA Ok(drained) } + async fn skip_credential(&mut self, key_id: &str) -> Result<(), GatewayError> { + self.candidates.skip_credential(key_id); + Ok(()) + } + + async fn skip_endpoint(&mut self, endpoint_id: &str) -> Result<(), GatewayError> { + self.candidates.skip_endpoint(endpoint_id); + Ok(()) + } + async fn skip_provider(&mut self, provider_id: &str) -> Result<(), GatewayError> { self.candidates.skip_provider(provider_id); Ok(()) @@ -249,6 +259,16 @@ impl LocalExecutionAttemptSource Ok(drained) } + async fn skip_credential(&mut self, key_id: &str) -> Result<(), GatewayError> { + self.candidates.skip_credential(key_id); + Ok(()) + } + + async fn skip_endpoint(&mut self, endpoint_id: &str) -> Result<(), GatewayError> { + self.candidates.skip_endpoint(endpoint_id); + Ok(()) + } + async fn skip_provider(&mut self, provider_id: &str) -> Result<(), GatewayError> { self.candidates.skip_provider(provider_id); Ok(()) diff --git a/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/request/url.rs b/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/request/url.rs index b9f064f24..3d42f140a 100644 --- a/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/request/url.rs +++ b/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/request/url.rs @@ -24,6 +24,7 @@ pub(crate) fn build_same_format_upstream_url( upstream_is_stream, request_query: parts.uri.query(), kiro_api_region: kiro_auth.map(|auth| auth.auth_config.effective_api_region()), + api_operation: spec.operation, provider_request_body, }, ) diff --git a/apps/aether-gateway/src/ai_serving/planner/route.rs b/apps/aether-gateway/src/ai_serving/planner/route.rs index 362a6b623..11c2658f9 100644 --- a/apps/aether-gateway/src/ai_serving/planner/route.rs +++ b/apps/aether-gateway/src/ai_serving/planner/route.rs @@ -1,8 +1,8 @@ use crate::ai_serving::GatewayControlDecision; use crate::ai_serving::{ is_matching_stream_http_request as is_matching_stream_http_request_impl, - resolve_execution_runtime_stream_plan_kind as resolve_execution_runtime_stream_plan_kind_impl, - resolve_execution_runtime_sync_plan_kind as resolve_execution_runtime_sync_plan_kind_impl, + resolve_execution_runtime_stream_plan_kind_with_client_surface as resolve_execution_runtime_stream_plan_kind_impl, + resolve_execution_runtime_sync_plan_kind_with_client_surface as resolve_execution_runtime_sync_plan_kind_impl, supports_stream_execution_decision_kind as supports_stream_execution_decision_kind_impl, supports_sync_execution_decision_kind as supports_sync_execution_decision_kind_impl, }; @@ -11,28 +11,34 @@ pub(crate) fn resolve_execution_runtime_stream_plan_kind( parts: &http::request::Parts, decision: &GatewayControlDecision, ) -> Option<&'static str> { - resolve_execution_runtime_stream_plan_kind_impl( + let plan_kind = resolve_execution_runtime_stream_plan_kind_impl( decision.route_class.as_deref(), decision.route_family.as_deref(), decision.route_kind.as_deref(), + decision.client_surface, decision.request_auth_channel.as_deref(), &parts.method, parts.uri.path(), - ) + )?; + crate::ai_serving::plan_kind_matches_api_operation(plan_kind, true, decision.api_operation) + .then_some(plan_kind) } pub(crate) fn resolve_execution_runtime_sync_plan_kind( parts: &http::request::Parts, decision: &GatewayControlDecision, ) -> Option<&'static str> { - resolve_execution_runtime_sync_plan_kind_impl( + let plan_kind = resolve_execution_runtime_sync_plan_kind_impl( decision.route_class.as_deref(), decision.route_family.as_deref(), decision.route_kind.as_deref(), + decision.client_surface, decision.request_auth_channel.as_deref(), &parts.method, parts.uri.path(), - ) + )?; + crate::ai_serving::plan_kind_matches_api_operation(plan_kind, false, decision.api_operation) + .then_some(plan_kind) } pub(crate) fn is_matching_stream_request( @@ -62,7 +68,7 @@ mod tests { resolve_execution_runtime_sync_plan_kind, supports_stream_execution_decision_kind, supports_sync_execution_decision_kind, }; - use crate::ai_serving::GatewayControlDecision; + use crate::ai_serving::{ApiOperation, ClientSurface, GatewayControlDecision}; fn sample_decision(route_family: &str, route_kind: &str) -> GatewayControlDecision { GatewayControlDecision { @@ -71,6 +77,9 @@ mod tests { route_class: Some("ai_public".to_string()), route_family: Some(route_family.to_string()), route_kind: Some(route_kind.to_string()), + client_surface: None, + api_operation: None, + gateway_credential_carrier: None, request_auth_channel: None, auth_context: None, admin_principal: None, @@ -121,7 +130,9 @@ mod tests { let (claude_parts, _) = claude_request.into_parts(); let claude_api_key = sample_decision_with_auth_channel("claude", "messages", "api_key"); - let claude_bearer = sample_decision_with_auth_channel("claude", "messages", "bearer_like"); + let mut claude_bearer = + sample_decision_with_auth_channel("claude", "messages", "bearer_like"); + claude_bearer.client_surface = Some(ClientSurface::ClaudeCode); assert_eq!( resolve_execution_runtime_sync_plan_kind(&claude_parts, &claude_api_key), Some("claude_chat_sync") @@ -131,6 +142,13 @@ mod tests { Some("claude_cli_stream") ); + let claude_sdk_bearer = + sample_decision_with_auth_channel("claude", "messages", "bearer_like"); + assert_eq!( + resolve_execution_runtime_sync_plan_kind(&claude_parts, &claude_sdk_bearer), + Some("claude_chat_sync") + ); + let gemini_request = Request::builder() .method(Method::POST) .uri("/v1beta/models/gemini-2.5-pro:generateContent") @@ -152,6 +170,36 @@ mod tests { ); } + #[test] + fn resolves_claude_count_tokens_as_native_sync_operation() { + let request = Request::builder() + .method(Method::POST) + .uri("/v1/messages/count_tokens") + .body(()) + .expect("request should build"); + let (parts, _) = request.into_parts(); + let mut decision = sample_decision("claude", "count_tokens"); + decision.api_operation = Some(ApiOperation::ClaudeCountTokens); + + assert_eq!( + resolve_execution_runtime_sync_plan_kind(&parts, &decision), + Some("claude_count_tokens_sync") + ); + assert!(supports_sync_execution_decision_kind( + "claude_count_tokens_sync" + )); + + decision.api_operation = Some(ApiOperation::ClaudeMessagesCreate); + assert_eq!( + resolve_execution_runtime_sync_plan_kind(&parts, &decision), + None + ); + assert_eq!( + resolve_execution_runtime_stream_plan_kind(&parts, &decision), + None + ); + } + #[test] fn stream_matching_uses_surface_route_logic() { let request = Request::builder() diff --git a/apps/aether-gateway/src/ai_serving/planner/specialized/files.rs b/apps/aether-gateway/src/ai_serving/planner/specialized/files.rs index edba6e57c..2d8ec9a62 100644 --- a/apps/aether-gateway/src/ai_serving/planner/specialized/files.rs +++ b/apps/aether-gateway/src/ai_serving/planner/specialized/files.rs @@ -194,6 +194,16 @@ impl LocalExecutionAttemptSource for LocalGeminiFilesSyncAttemptS Ok(drained) } + async fn skip_credential(&mut self, key_id: &str) -> Result<(), GatewayError> { + self.candidates.skip_credential(key_id); + Ok(()) + } + + async fn skip_endpoint(&mut self, endpoint_id: &str) -> Result<(), GatewayError> { + self.candidates.skip_endpoint(endpoint_id); + Ok(()) + } + async fn skip_provider(&mut self, provider_id: &str) -> Result<(), GatewayError> { self.candidates.skip_provider(provider_id); Ok(()) @@ -222,6 +232,16 @@ impl LocalExecutionAttemptSource for LocalGeminiFilesStreamAtte Ok(drained) } + async fn skip_credential(&mut self, key_id: &str) -> Result<(), GatewayError> { + self.candidates.skip_credential(key_id); + Ok(()) + } + + async fn skip_endpoint(&mut self, endpoint_id: &str) -> Result<(), GatewayError> { + self.candidates.skip_endpoint(endpoint_id); + Ok(()) + } + async fn skip_provider(&mut self, provider_id: &str) -> Result<(), GatewayError> { self.candidates.skip_provider(provider_id); Ok(()) diff --git a/apps/aether-gateway/src/ai_serving/planner/specialized/image.rs b/apps/aether-gateway/src/ai_serving/planner/specialized/image.rs index 8f407a32d..b02027337 100644 --- a/apps/aether-gateway/src/ai_serving/planner/specialized/image.rs +++ b/apps/aether-gateway/src/ai_serving/planner/specialized/image.rs @@ -272,6 +272,16 @@ impl LocalExecutionAttemptSource for LocalOpenAiImageSyncAttemptS Ok(drained) } + async fn skip_credential(&mut self, key_id: &str) -> Result<(), GatewayError> { + self.candidates.skip_credential(key_id); + Ok(()) + } + + async fn skip_endpoint(&mut self, endpoint_id: &str) -> Result<(), GatewayError> { + self.candidates.skip_endpoint(endpoint_id); + Ok(()) + } + async fn skip_provider(&mut self, provider_id: &str) -> Result<(), GatewayError> { self.candidates.skip_provider(provider_id); Ok(()) @@ -300,6 +310,16 @@ impl LocalExecutionAttemptSource for LocalOpenAiImageStreamAtte Ok(drained) } + async fn skip_credential(&mut self, key_id: &str) -> Result<(), GatewayError> { + self.candidates.skip_credential(key_id); + Ok(()) + } + + async fn skip_endpoint(&mut self, endpoint_id: &str) -> Result<(), GatewayError> { + self.candidates.skip_endpoint(endpoint_id); + Ok(()) + } + async fn skip_provider(&mut self, provider_id: &str) -> Result<(), GatewayError> { self.candidates.skip_provider(provider_id); Ok(()) diff --git a/apps/aether-gateway/src/ai_serving/planner/specialized/video.rs b/apps/aether-gateway/src/ai_serving/planner/specialized/video.rs index 4548fccb9..172bd0a69 100644 --- a/apps/aether-gateway/src/ai_serving/planner/specialized/video.rs +++ b/apps/aether-gateway/src/ai_serving/planner/specialized/video.rs @@ -124,6 +124,16 @@ impl LocalExecutionAttemptSource for LocalVideoCreateSyncAttemptS Ok(drained) } + async fn skip_credential(&mut self, key_id: &str) -> Result<(), GatewayError> { + self.candidates.skip_credential(key_id); + Ok(()) + } + + async fn skip_endpoint(&mut self, endpoint_id: &str) -> Result<(), GatewayError> { + self.candidates.skip_endpoint(endpoint_id); + Ok(()) + } + async fn skip_provider(&mut self, provider_id: &str) -> Result<(), GatewayError> { self.candidates.skip_provider(provider_id); Ok(()) diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/family/build.rs b/apps/aether-gateway/src/ai_serving/planner/standard/family/build.rs index a4e700bff..d2d599b49 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/family/build.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/family/build.rs @@ -202,6 +202,16 @@ impl LocalExecutionAttemptSource for LocalStandardSyncAttemptSour Ok(drained) } + async fn skip_credential(&mut self, key_id: &str) -> Result<(), GatewayError> { + self.candidates.skip_credential(key_id); + Ok(()) + } + + async fn skip_endpoint(&mut self, endpoint_id: &str) -> Result<(), GatewayError> { + self.candidates.skip_endpoint(endpoint_id); + Ok(()) + } + async fn skip_provider(&mut self, provider_id: &str) -> Result<(), GatewayError> { self.candidates.skip_provider(provider_id); Ok(()) @@ -235,6 +245,16 @@ impl LocalExecutionAttemptSource for LocalStandardStreamAttempt Ok(drained) } + async fn skip_credential(&mut self, key_id: &str) -> Result<(), GatewayError> { + self.candidates.skip_credential(key_id); + Ok(()) + } + + async fn skip_endpoint(&mut self, endpoint_id: &str) -> Result<(), GatewayError> { + self.candidates.skip_endpoint(endpoint_id); + Ok(()) + } + async fn skip_provider(&mut self, provider_id: &str) -> Result<(), GatewayError> { self.candidates.skip_provider(provider_id); Ok(()) diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/family/payload.rs b/apps/aether-gateway/src/ai_serving/planner/standard/family/payload.rs index 151cde139..6e5ec6ef5 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/family/payload.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/family/payload.rs @@ -373,6 +373,8 @@ mod tests { auth_snapshot: sample_auth_snapshot(), required_capabilities: None, request_auth_channel: None, + client_surface: None, + gateway_credential_carrier: None, client_session_affinity: None, routing_policy: None, routing_trace_seed: None, diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/gemini/plan_builders.rs b/apps/aether-gateway/src/ai_serving/planner/standard/gemini/plan_builders.rs index 1bf47d1e6..e84dd04cf 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/gemini/plan_builders.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/gemini/plan_builders.rs @@ -1,12 +1,10 @@ use std::collections::BTreeMap; -use aether_contracts::RequestBody; - use super::{ augment_sync_report_context, build_ai_execution_plan_from_decision, - generic_decision_missing_exact_provider_request, take_ai_decision_plan_core, - take_ai_upstream_auth_pair, take_non_empty_string, AiExecutionPlanFromDecisionParts, - AiStreamAttempt, AiSyncAttempt, + generic_decision_missing_exact_provider_request, resolve_ai_passthrough_sync_request_body, + take_ai_decision_plan_core, take_ai_upstream_auth_pair, take_non_empty_string, + AiExecutionPlanFromDecisionParts, AiStreamAttempt, AiSyncAttempt, }; use crate::ai_serving::transport::{ build_standard_plan_fallback_headers, StandardPlanFallbackAcceptPolicy, @@ -61,6 +59,10 @@ pub(crate) fn build_gemini_sync_plan_from_decision( &provider_request_headers, &provider_request_body_value, )?; + let request_body = resolve_ai_passthrough_sync_request_body( + Some(provider_request_body_value), + payload.provider_request_body_base64.take(), + ); let stream = payload.upstream_is_stream; let plan = build_ai_execution_plan_from_decision( &mut payload, @@ -70,7 +72,7 @@ pub(crate) fn build_gemini_sync_plan_from_decision( url, headers: std::mem::take(&mut provider_request_headers), content_type, - body: RequestBody::from_json(provider_request_body_value), + body: request_body, stream, }, ); @@ -129,6 +131,10 @@ pub(crate) fn build_gemini_stream_plan_from_decision( &provider_request_headers, &provider_request_body_value, )?; + let request_body = resolve_ai_passthrough_sync_request_body( + Some(provider_request_body_value), + payload.provider_request_body_base64.take(), + ); let plan = build_ai_execution_plan_from_decision( &mut payload, AiExecutionPlanFromDecisionParts { @@ -137,7 +143,7 @@ pub(crate) fn build_gemini_stream_plan_from_decision( url, headers: std::mem::take(&mut provider_request_headers), content_type, - body: RequestBody::from_json(provider_request_body_value), + body: request_body, stream: true, }, ); diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/mod.rs b/apps/aether-gateway/src/ai_serving/planner/standard/mod.rs index 03d9699a6..c8ab13128 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/mod.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/mod.rs @@ -84,6 +84,7 @@ pub(crate) fn build_standard_upstream_url( upstream_is_stream, parts.uri.query(), None, + None, provider_request_body, ) } diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs index b8e142673..ace3fbc23 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs @@ -2199,6 +2199,8 @@ mod tests { auth_snapshot: sample_auth_snapshot(), required_capabilities: None, request_auth_channel: None, + client_surface: None, + gateway_credential_carrier: None, client_session_affinity: None, routing_policy: None, routing_trace_seed: None, diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/stream.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/stream.rs index 76ac901af..b25a1be84 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/stream.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/stream.rs @@ -139,6 +139,20 @@ impl LocalExecutionAttemptSource for LocalOpenAiChatStreamAttem Ok(drained) } + async fn skip_credential(&mut self, key_id: &str) -> Result<(), GatewayError> { + self.prefetched_attempts + .retain(|attempt| attempt.eligible.candidate.key_id != key_id); + self.candidates.skip_credential(key_id); + Ok(()) + } + + async fn skip_endpoint(&mut self, endpoint_id: &str) -> Result<(), GatewayError> { + self.prefetched_attempts + .retain(|attempt| attempt.eligible.candidate.endpoint_id != endpoint_id); + self.candidates.skip_endpoint(endpoint_id); + Ok(()) + } + async fn skip_provider(&mut self, provider_id: &str) -> Result<(), GatewayError> { self.prefetched_attempts .retain(|attempt| attempt.eligible.candidate.provider_id != provider_id); diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/sync.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/sync.rs index 7772c0645..812def7c6 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/sync.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/sync.rs @@ -117,6 +117,16 @@ impl LocalExecutionAttemptSource for LocalOpenAiChatSyncAttemptSo Ok(drained) } + async fn skip_credential(&mut self, key_id: &str) -> Result<(), GatewayError> { + self.candidates.skip_credential(key_id); + Ok(()) + } + + async fn skip_endpoint(&mut self, endpoint_id: &str) -> Result<(), GatewayError> { + self.candidates.skip_endpoint(endpoint_id); + Ok(()) + } + async fn skip_provider(&mut self, provider_id: &str) -> Result<(), GatewayError> { self.candidates.skip_provider(provider_id); Ok(()) diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/plans.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/plans.rs index be88eb8d9..649f173b7 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/plans.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/plans.rs @@ -186,6 +186,16 @@ impl LocalExecutionAttemptSource for LocalOpenAiResponsesSyncAtte Ok(drained) } + async fn skip_credential(&mut self, key_id: &str) -> Result<(), GatewayError> { + self.candidates.skip_credential(key_id); + Ok(()) + } + + async fn skip_endpoint(&mut self, endpoint_id: &str) -> Result<(), GatewayError> { + self.candidates.skip_endpoint(endpoint_id); + Ok(()) + } + async fn skip_provider(&mut self, provider_id: &str) -> Result<(), GatewayError> { self.candidates.skip_provider(provider_id); Ok(()) @@ -219,6 +229,16 @@ impl LocalExecutionAttemptSource for LocalOpenAiResponsesStream Ok(drained) } + async fn skip_credential(&mut self, key_id: &str) -> Result<(), GatewayError> { + self.candidates.skip_credential(key_id); + Ok(()) + } + + async fn skip_endpoint(&mut self, endpoint_id: &str) -> Result<(), GatewayError> { + self.candidates.skip_endpoint(endpoint_id); + Ok(()) + } + async fn skip_provider(&mut self, provider_id: &str) -> Result<(), GatewayError> { self.candidates.skip_provider(provider_id); Ok(()) diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/plan_builders.rs b/apps/aether-gateway/src/ai_serving/planner/standard/plan_builders.rs index 99078f2e4..53127882a 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/plan_builders.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/plan_builders.rs @@ -1,9 +1,8 @@ use std::collections::BTreeMap; -use aether_contracts::RequestBody; - use super::{ - augment_sync_report_context, build_ai_execution_plan_from_decision, take_ai_decision_plan_core, + augment_sync_report_context, build_ai_execution_plan_from_decision, + resolve_ai_passthrough_sync_request_body, take_ai_decision_plan_core, take_ai_upstream_auth_pair, take_non_empty_string, AiExecutionPlanFromDecisionParts, AiStreamAttempt, AiSyncAttempt, }; @@ -63,6 +62,10 @@ pub(crate) fn build_standard_sync_plan_from_decision( &provider_request_headers, &provider_request_body_value, )?; + let request_body = resolve_ai_passthrough_sync_request_body( + Some(provider_request_body_value), + payload.provider_request_body_base64.take(), + ); let stream = payload.upstream_is_stream; let plan = build_ai_execution_plan_from_decision( &mut payload, @@ -72,7 +75,7 @@ pub(crate) fn build_standard_sync_plan_from_decision( url, headers: std::mem::take(&mut provider_request_headers), content_type, - body: RequestBody::from_json(provider_request_body_value), + body: request_body, stream, }, ); @@ -146,6 +149,10 @@ pub(crate) fn build_standard_stream_plan_from_decision( &provider_request_headers, &provider_request_body_value, )?; + let request_body = resolve_ai_passthrough_sync_request_body( + Some(provider_request_body_value), + payload.provider_request_body_base64.take(), + ); let stream = payload.upstream_is_stream; let plan = build_ai_execution_plan_from_decision( &mut payload, @@ -155,7 +162,7 @@ pub(crate) fn build_standard_stream_plan_from_decision( url, headers: std::mem::take(&mut provider_request_headers), content_type, - body: RequestBody::from_json(provider_request_body_value), + body: request_body, stream, }, ); @@ -166,3 +173,88 @@ pub(crate) fn build_standard_stream_plan_from_decision( report_context, })) } + +#[cfg(test)] +mod tests { + use aether_contracts::{ExecutionResponseBodyMode, EXECUTION_RESPONSE_BODY_MODE_HEADER}; + use serde_json::json; + + use super::{ + build_standard_stream_plan_from_decision, build_standard_sync_plan_from_decision, + AiExecutionDecision, + }; + + fn decision_with_raw_body(upstream_is_stream: bool) -> AiExecutionDecision { + serde_json::from_value(json!({ + "action": if upstream_is_stream { "stream" } else { "sync" }, + "request_id": "req-raw", + "provider_id": "provider-raw", + "endpoint_id": "endpoint-raw", + "key_id": "key-raw", + "upstream_url": "https://api.anthropic.test/v1/messages", + "provider_api_format": "claude:messages", + "client_api_format": "claude:messages", + "provider_request_headers": { + "content-type": "application/json", + (EXECUTION_RESPONSE_BODY_MODE_HEADER): ExecutionResponseBodyMode::PreserveBytes.as_str() + }, + "provider_request_body": { + "model": "claude-sonnet-4", + "messages": [] + }, + "provider_request_body_base64": "eyAibW9kZWwiOiAiY2xhdWRlLXNvbm5ldC00IiwgIm1lc3NhZ2VzIjogW10gfQ==", + "content_type": "application/json", + "upstream_is_stream": upstream_is_stream + })) + .expect("decision should deserialize") + } + + fn request_parts() -> http::request::Parts { + http::Request::builder() + .uri("http://localhost/v1/messages") + .body(()) + .expect("request should build") + .into_parts() + .0 + } + + #[test] + fn standard_sync_plan_prefers_exact_request_body_bytes() { + let built = build_standard_sync_plan_from_decision( + &request_parts(), + &json!({}), + decision_with_raw_body(false), + ) + .expect("plan should build") + .expect("plan should exist"); + + assert!(built.plan.body.json_body.is_none()); + assert_eq!( + built.plan.body.body_bytes_b64.as_deref(), + Some("eyAibW9kZWwiOiAiY2xhdWRlLXNvbm5ldC00IiwgIm1lc3NhZ2VzIjogW10gfQ==") + ); + assert_eq!( + built + .plan + .headers + .get(EXECUTION_RESPONSE_BODY_MODE_HEADER) + .map(String::as_str), + Some(ExecutionResponseBodyMode::PreserveBytes.as_str()) + ); + } + + #[test] + fn standard_stream_plan_prefers_exact_request_body_bytes() { + let built = build_standard_stream_plan_from_decision( + &request_parts(), + &json!({}), + decision_with_raw_body(true), + false, + ) + .expect("plan should build") + .expect("plan should exist"); + + assert!(built.plan.body.json_body.is_none()); + assert!(built.plan.body.body_bytes_b64.is_some()); + } +} diff --git a/apps/aether-gateway/src/ai_serving/pure/mod.rs b/apps/aether-gateway/src/ai_serving/pure/mod.rs index a258dd07e..d710ff9d7 100644 --- a/apps/aether-gateway/src/ai_serving/pure/mod.rs +++ b/apps/aether-gateway/src/ai_serving/pure/mod.rs @@ -92,9 +92,12 @@ pub(crate) use aether_ai_formats::api::{ request_conversion_requires_enable_flag, request_path_implies_stream_request, resolve_claude_stream_spec, resolve_claude_sync_spec, resolve_codex_responses_model_capabilities, resolve_execution_runtime_stream_plan_kind, - resolve_execution_runtime_sync_plan_kind, resolve_finalize_stream_rewrite_mode, - resolve_gemini_files_stream_spec, resolve_gemini_files_sync_spec, resolve_gemini_stream_spec, - resolve_gemini_sync_spec, resolve_local_image_stream_spec, resolve_local_image_sync_spec, + resolve_execution_runtime_stream_plan_kind_with_client_surface, + resolve_execution_runtime_sync_plan_kind, + resolve_execution_runtime_sync_plan_kind_with_client_surface, + resolve_finalize_stream_rewrite_mode, resolve_gemini_files_stream_spec, + resolve_gemini_files_sync_spec, resolve_gemini_stream_spec, resolve_gemini_sync_spec, + resolve_local_image_stream_spec, resolve_local_image_sync_spec, resolve_local_same_format_stream_spec, resolve_local_same_format_sync_spec, resolve_local_video_sync_spec, resolve_openai_chat_max_tokens, resolve_openai_embedding_sync_spec, resolve_openai_responses_stream_spec, @@ -129,14 +132,15 @@ pub(crate) use aether_ai_formats::api::{ CLAUDE_CHAT_SYNC_SUCCESS_REPORT_KIND, CLAUDE_CLI_STREAM_PLAN_KIND, CLAUDE_CLI_STREAM_SUCCESS_REPORT_KIND, CLAUDE_CLI_SYNC_ERROR_REPORT_KIND, CLAUDE_CLI_SYNC_FINALIZE_REPORT_KIND, CLAUDE_CLI_SYNC_PLAN_KIND, - CLAUDE_CLI_SYNC_SUCCESS_REPORT_KIND, CODEX_OPENAI_IMAGE_DEFAULT_MODEL, - CODEX_OPENAI_IMAGE_DEFAULT_OUTPUT_FORMAT, CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_MODEL, - CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_PROMPT, CODEX_OPENAI_IMAGE_INTERNAL_MODEL, - EXECUTION_RUNTIME_STREAM_ACTION, EXECUTION_RUNTIME_STREAM_DECISION_ACTION, - EXECUTION_RUNTIME_SYNC_ACTION, EXECUTION_RUNTIME_SYNC_DECISION_ACTION, - GEMINI_CHAT_STREAM_PLAN_KIND, GEMINI_CHAT_STREAM_SUCCESS_REPORT_KIND, - GEMINI_CHAT_SYNC_ERROR_REPORT_KIND, GEMINI_CHAT_SYNC_FINALIZE_REPORT_KIND, - GEMINI_CHAT_SYNC_PLAN_KIND, GEMINI_CHAT_SYNC_SUCCESS_REPORT_KIND, GEMINI_CLI_STREAM_PLAN_KIND, + CLAUDE_CLI_SYNC_SUCCESS_REPORT_KIND, CLAUDE_COUNT_TOKENS_SYNC_PLAN_KIND, + CODEX_OPENAI_IMAGE_DEFAULT_MODEL, CODEX_OPENAI_IMAGE_DEFAULT_OUTPUT_FORMAT, + CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_MODEL, CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_PROMPT, + CODEX_OPENAI_IMAGE_INTERNAL_MODEL, EXECUTION_RUNTIME_STREAM_ACTION, + EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_ACTION, + EXECUTION_RUNTIME_SYNC_DECISION_ACTION, GEMINI_CHAT_STREAM_PLAN_KIND, + GEMINI_CHAT_STREAM_SUCCESS_REPORT_KIND, GEMINI_CHAT_SYNC_ERROR_REPORT_KIND, + GEMINI_CHAT_SYNC_FINALIZE_REPORT_KIND, GEMINI_CHAT_SYNC_PLAN_KIND, + GEMINI_CHAT_SYNC_SUCCESS_REPORT_KIND, GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_STREAM_SUCCESS_REPORT_KIND, GEMINI_CLI_SYNC_ERROR_REPORT_KIND, GEMINI_CLI_SYNC_FINALIZE_REPORT_KIND, GEMINI_CLI_SYNC_PLAN_KIND, GEMINI_CLI_SYNC_SUCCESS_REPORT_KIND, GEMINI_CLI_V1INTERNAL_ENVELOPE_NAME, @@ -167,5 +171,28 @@ pub(crate) use aether_ai_formats::api::{ pub(crate) use aether_ai_formats::{ api_format_defaults_to_client_error_failover, api_format_defaults_to_non_stream, api_format_permission_covers, intersect_api_format_allowed_lists, is_embedding_api_format, - is_rerank_api_format, openai_responses_request_operation, + is_rerank_api_format, openai_responses_request_operation, ApiOperation, ClientSurface, }; + +pub(crate) fn plan_kind_matches_api_operation( + plan_kind: &str, + require_streaming: bool, + expected_operation: Option, +) -> bool { + let Some(expected_operation) = expected_operation else { + return true; + }; + if expected_operation == ApiOperation::OpenAiResponsesCompact { + return if require_streaming { + plan_kind == OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND + } else { + plan_kind == OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND + }; + } + let resolved_operation = if require_streaming { + resolve_local_same_format_stream_spec(plan_kind).and_then(|spec| spec.operation) + } else { + resolve_local_same_format_sync_spec(plan_kind).and_then(|spec| spec.operation) + }; + resolved_operation == Some(expected_operation) +} diff --git a/apps/aether-gateway/src/ai_serving/transport.rs b/apps/aether-gateway/src/ai_serving/transport.rs index 6ea43c541..a7eb7d5cb 100644 --- a/apps/aether-gateway/src/ai_serving/transport.rs +++ b/apps/aether-gateway/src/ai_serving/transport.rs @@ -82,10 +82,11 @@ pub(crate) use aether_provider_transport::{ build_windsurf_cascade_headers, build_windsurf_cascade_request_body, build_windsurf_cascade_upstream_url, candidate_common_transport_skip_reason, candidate_transport_pair_skip_reason, classify_same_format_provider_request_behavior, - ensure_upstream_auth_header, gemini_files_transport_unsupported_reason, - header_rules_are_locally_supported, header_rules_have_enabled_rules, - is_gemini_cli_provider_transport, is_windsurf_provider_transport, - local_gemini_transport_unsupported_reason_with_network, + classify_same_format_provider_request_behavior_for_operation, + enforce_same_format_provider_api_operation_body_policy, ensure_upstream_auth_header, + gemini_files_transport_unsupported_reason, header_rules_are_locally_supported, + header_rules_have_enabled_rules, is_gemini_cli_provider_transport, + is_windsurf_provider_transport, local_gemini_transport_unsupported_reason_with_network, local_openai_chat_transport_unsupported_reason, local_standard_transport_unsupported_reason_with_network, local_windsurf_request_transport_unsupported_reason_with_network, @@ -93,22 +94,22 @@ pub(crate) use aether_provider_transport::{ request_conversion_enabled_for_transport, request_conversion_transport_supported, request_conversion_transport_unsupported_reason, request_pair_allowed_for_transport, request_pair_direct_auth, request_pair_transport_unsupported_reason, - resolve_gemini_cli_project_id, resolve_gemini_files_auth, resolve_grok_session_auth, - resolve_local_gemini_cli_request_auth, resolve_openai_image_auth, - resolve_same_format_provider_direct_auth, resolve_transport_execution_timeouts, - resolve_transport_profile, resolve_transport_proxy_snapshot, - resolve_transport_proxy_snapshot_with_tunnel_affinity, resolve_video_create_auth, - same_format_provider_transport_supported, same_format_provider_transport_unsupported_reason, - should_skip_upstream_passthrough_header, should_try_same_format_provider_oauth_auth, - supports_local_gemini_transport_with_network, + resolve_anthropic_compatibility_profile, resolve_gemini_cli_project_id, + resolve_gemini_files_auth, resolve_grok_session_auth, resolve_local_gemini_cli_request_auth, + resolve_openai_image_auth, resolve_same_format_provider_direct_auth, + resolve_transport_execution_timeouts, resolve_transport_profile, + resolve_transport_proxy_snapshot, resolve_transport_proxy_snapshot_with_tunnel_affinity, + resolve_video_create_auth, same_format_provider_transport_supported, + same_format_provider_transport_unsupported_reason, should_skip_upstream_passthrough_header, + should_try_same_format_provider_oauth_auth, supports_local_gemini_transport_with_network, supports_local_generic_oauth_request_auth_resolution, supports_local_oauth_request_auth_resolution, transport_proxy_is_locally_supported, - video_create_transport_unsupported_reason, CandidateTransportPolicyFacts, - GatewayProviderTransportSnapshot, GeminiCliRequestAuth, GeminiCliRequestAuthSupport, - GeminiCliRequestAuthUnsupportedReason, GeminiCliRequestEnvelopeSupport, - GeminiFilesHeadersInput, GeminiFilesRequestBodyError, GeminiFilesRequestBodyParts, - GrokHeaderInput, LocalResolvedOAuthRequestAuth, ProviderOpenAiImageHeadersInput, - ProviderVideoCreateFamily, ProviderVideoCreateHeadersInput, + transport_supports_api_operation, video_create_transport_unsupported_reason, + AnthropicCompatibilityProfile, CandidateTransportPolicyFacts, GatewayProviderTransportSnapshot, + GeminiCliRequestAuth, GeminiCliRequestAuthSupport, GeminiCliRequestAuthUnsupportedReason, + GeminiCliRequestEnvelopeSupport, GeminiFilesHeadersInput, GeminiFilesRequestBodyError, + GeminiFilesRequestBodyParts, GrokHeaderInput, LocalResolvedOAuthRequestAuth, + ProviderOpenAiImageHeadersInput, ProviderVideoCreateFamily, ProviderVideoCreateHeadersInput, SameFormatProviderCompatibilityEdit, SameFormatProviderCompatibilityEditAction, SameFormatProviderFamily, SameFormatProviderHeadersInput, SameFormatProviderRequestBehavior, SameFormatProviderRequestBehaviorParams, SameFormatProviderRequestBodyInput, diff --git a/apps/aether-gateway/src/api/ai/registry.rs b/apps/aether-gateway/src/api/ai/registry.rs index 7ae1df4e6..172fd9ed1 100644 --- a/apps/aether-gateway/src/api/ai/registry.rs +++ b/apps/aether-gateway/src/api/ai/registry.rs @@ -1,8 +1,13 @@ +use axum::body::Body; +use axum::extract::Request; +use axum::http::{header, HeaderValue, Response, StatusCode}; use axum::routing::{any, post}; use axum::Router; use super::{aliyun, claude, doubao, gemini, jina, openai}; -use crate::{handlers::proxy::proxy_request, state::AppState}; +use crate::api::response::build_local_http_error_response_with_request_path; +use crate::headers::extract_or_generate_trace_id; +use crate::{handlers::proxy::proxy_request, state::AppState, GatewayError}; // Router registration patterns live here so AI public ingress has a single mount registry. // They intentionally stay separate from manifest-facing route inventories in constants.rs, @@ -11,8 +16,6 @@ const AI_POST_ROUTE_PATTERNS: &[&str] = &[ "/v1/chat/completions", "/v1/embeddings", "/v1/rerank", - "/v1/messages", - "/v1/messages/count_tokens", "/v1/responses", "/v1/responses/compact", "/v1/alpha/search", @@ -32,6 +35,8 @@ const AI_POST_ROUTE_PATTERNS: &[&str] = &[ "/v1internal:streamGenerateContent", ]; +const CLAUDE_POST_ROUTE_PATTERNS: &[&str] = &["/v1/messages", "/v1/messages/count_tokens"]; + const AI_ANY_ROUTE_PATTERNS: &[&str] = &[ "/v1/models/{*gemini_path}", "/v1beta/models/{*gemini_path}", @@ -48,12 +53,33 @@ pub(crate) fn mount_ai_routes(mut router: Router) -> Router for path in AI_POST_ROUTE_PATTERNS { router = router.route(path, post(proxy_request)); } + for path in CLAUDE_POST_ROUTE_PATTERNS { + router = router.route( + path, + post(proxy_request).fallback(claude_method_not_allowed), + ); + } for path in AI_ANY_ROUTE_PATTERNS { router = router.route(path, any(proxy_request)); } router } +async fn claude_method_not_allowed(request: Request) -> Result, GatewayError> { + let trace_id = extract_or_generate_trace_id(request.headers()); + let mut response = build_local_http_error_response_with_request_path( + &trace_id, + None, + Some(request.uri().path()), + StatusCode::METHOD_NOT_ALLOWED, + "Method not allowed", + )?; + response + .headers_mut() + .insert(header::ALLOW, HeaderValue::from_static("POST")); + Ok(response) +} + pub(crate) fn public_api_format_local_path(api_format: &str) -> &'static str { let normalized = api_format.trim().to_ascii_lowercase(); openai::local_path(&normalized) diff --git a/apps/aether-gateway/src/api/response.rs b/apps/aether-gateway/src/api/response.rs index 4535193e5..403037c7a 100644 --- a/apps/aether-gateway/src/api/response.rs +++ b/apps/aether-gateway/src/api/response.rs @@ -6,6 +6,7 @@ use axum::http::Response; use axum::http::StatusCode; use serde_json::json; +use crate::ai_serving::{build_core_error_body_for_client_format, LocalCoreSyncErrorKind}; use crate::constants::*; use crate::control::GatewayControlDecision; use crate::control::GatewayLocalAuthRejection; @@ -191,7 +192,7 @@ pub(crate) fn build_local_balance_denied_response( Some(remaining) => format!("余额不足（剩余: ${remaining:.2}）"), None => "余额不足".to_string(), }; - let payload = json!({ + let fallback_payload = json!({ "error": { "type": "balance_exceeded", "message": message, @@ -201,6 +202,13 @@ pub(crate) fn build_local_balance_denied_response( } } }); + let payload = build_local_error_payload( + control_decision, + None, + &message, + LocalCoreSyncErrorKind::RateLimit, + fallback_payload, + ); let body = serde_json::to_vec(&payload).map_err(|err| GatewayError::Internal(err.to_string()))?; let headers = BTreeMap::from([("content-type".to_string(), "application/json".to_string())]); @@ -218,12 +226,20 @@ pub(crate) fn build_local_user_rpm_limited_response( control_decision: Option<&GatewayControlDecision>, rejection: &FrontdoorUserRpmRejection, ) -> Result, GatewayError> { - let payload = json!({ + let message = "请求过于频繁，请稍后重试"; + let fallback_payload = json!({ "error": { "type": "rate_limit_exceeded", - "message": "请求过于频繁，请稍后重试", + "message": message, } }); + let payload = build_local_error_payload( + control_decision, + None, + message, + LocalCoreSyncErrorKind::RateLimit, + fallback_payload, + ); let body = serde_json::to_vec(&payload).map_err(|err| GatewayError::Internal(err.to_string()))?; let headers = BTreeMap::from([ @@ -248,12 +264,35 @@ pub(crate) fn build_local_http_error_response( status_code: StatusCode, message: &str, ) -> Result, GatewayError> { - let payload = json!({ + build_local_http_error_response_with_request_path( + trace_id, + control_decision, + None, + status_code, + message, + ) +} + +pub(crate) fn build_local_http_error_response_with_request_path( + trace_id: &str, + control_decision: Option<&GatewayControlDecision>, + request_path: Option<&str>, + status_code: StatusCode, + message: &str, +) -> Result, GatewayError> { + let fallback_payload = json!({ "error": { "type": "http_error", "message": message, } }); + let payload = build_local_error_payload( + control_decision, + request_path, + message, + local_error_kind_for_status(status_code), + fallback_payload, + ); let body = serde_json::to_vec(&payload).map_err(|err| GatewayError::Internal(err.to_string()))?; let headers = BTreeMap::from([("content-type".to_string(), "application/json".to_string())]); @@ -329,19 +368,28 @@ pub(crate) fn build_local_auth_rejection_response( pub(crate) fn build_local_overloaded_response( trace_id: &str, control_decision: Option<&GatewayControlDecision>, + request_path: Option<&str>, gate: &str, limit: usize, ) -> Result, GatewayError> { - let payload = json!({ + let message = "服务繁忙，请稍后重试"; + let fallback_payload = json!({ "error": { "type": "overloaded", - "message": "服务繁忙，请稍后重试", + "message": message, "details": { "gate": gate, "limit": limit, } } }); + let payload = build_local_error_payload( + control_decision, + request_path, + message, + LocalCoreSyncErrorKind::Overloaded, + fallback_payload, + ); let body = serde_json::to_vec(&payload).map_err(|err| GatewayError::Internal(err.to_string()))?; let headers = BTreeMap::from([("content-type".to_string(), "application/json".to_string())]); @@ -354,10 +402,65 @@ pub(crate) fn build_local_overloaded_response( ) } +fn build_local_error_payload( + control_decision: Option<&GatewayControlDecision>, + request_path: Option<&str>, + message: &str, + kind: LocalCoreSyncErrorKind, + fallback_payload: serde_json::Value, +) -> serde_json::Value { + if !local_error_uses_claude_format(control_decision, request_path) { + return fallback_payload; + } + + build_core_error_body_for_client_format("claude:messages", message, None, kind) + .unwrap_or(fallback_payload) +} + +fn local_error_uses_claude_format( + control_decision: Option<&GatewayControlDecision>, + request_path: Option<&str>, +) -> bool { + control_decision.is_some_and(|decision| { + decision.route_family.as_deref() == Some("claude") + || decision + .auth_endpoint_signature + .as_deref() + .is_some_and(|format| { + crate::ai_serving::normalize_api_format_alias(format) + .eq_ignore_ascii_case("claude:messages") + }) + }) || request_path.is_some_and(|path| { + matches!( + path.trim_end_matches('/'), + "/v1/messages" | "/v1/messages/count_tokens" + ) + }) +} + +fn local_error_kind_for_status(status: StatusCode) -> LocalCoreSyncErrorKind { + match status.as_u16() { + 400 | 405 | 422 => LocalCoreSyncErrorKind::InvalidRequest, + 401 => LocalCoreSyncErrorKind::Authentication, + 403 => LocalCoreSyncErrorKind::PermissionDenied, + 404 => LocalCoreSyncErrorKind::NotFound, + 413 => LocalCoreSyncErrorKind::RequestTooLarge, + 429 => LocalCoreSyncErrorKind::RateLimit, + 503 | 529 => LocalCoreSyncErrorKind::Overloaded, + _ => LocalCoreSyncErrorKind::ServerError, + } +} + #[cfg(test)] mod tests { - use super::build_client_response_from_parts; - use axum::body::Body; + use super::{ + build_client_response_from_parts, build_local_auth_rejection_response, + build_local_http_error_response_with_request_path, build_local_overloaded_response, + build_local_user_rpm_limited_response, + }; + use crate::control::{GatewayControlDecision, GatewayLocalAuthRejection}; + use crate::rate_limit::FrontdoorUserRpmRejection; + use axum::body::{to_bytes, Body}; use std::collections::BTreeMap; #[test] @@ -386,4 +489,96 @@ mod tests { Some("no") ); } + + fn claude_decision() -> GatewayControlDecision { + GatewayControlDecision::synthetic( + "/v1/messages", + Some("ai_public".to_string()), + Some("claude".to_string()), + Some("messages".to_string()), + Some("claude:messages".to_string()), + ) + } + + async fn response_json(response: http::Response) -> serde_json::Value { + let body = to_bytes(response.into_body(), usize::MAX) + .await + .expect("response body should read"); + serde_json::from_slice(&body).expect("response body should be JSON") + } + + #[tokio::test] + async fn claude_local_errors_use_anthropic_envelopes() { + let decision = claude_decision(); + let invalid_key = build_local_auth_rejection_response( + "trace-auth", + Some(&decision), + &GatewayLocalAuthRejection::InvalidApiKey, + ) + .expect("invalid-key response should build"); + let invalid_key = response_json(invalid_key).await; + assert_eq!(invalid_key["type"], "error"); + assert_eq!(invalid_key["error"]["type"], "authentication_error"); + + let rpm = build_local_user_rpm_limited_response( + "trace-rpm", + Some(&decision), + &FrontdoorUserRpmRejection { + scope: "api_key", + limit: 1, + retry_after: 60, + }, + ) + .expect("RPM response should build"); + let rpm = response_json(rpm).await; + assert_eq!(rpm["type"], "error"); + assert_eq!(rpm["error"]["type"], "rate_limit_error"); + + let overloaded = build_local_overloaded_response( + "trace-overload", + None, + Some("/v1/messages/count_tokens"), + "requests", + 10, + ) + .expect("overload response should build"); + let overloaded = response_json(overloaded).await; + assert_eq!(overloaded["type"], "error"); + assert_eq!(overloaded["error"]["type"], "overloaded_error"); + } + + #[tokio::test] + async fn claude_path_shapes_pre_control_http_errors_and_413() { + for path in ["/v1/messages", "/v1/messages/count_tokens"] { + let forbidden = build_local_http_error_response_with_request_path( + "trace-pre-control", + None, + Some(path), + http::StatusCode::FORBIDDEN, + "blocked", + ) + .expect("forbidden response should build"); + let forbidden = response_json(forbidden).await; + assert_eq!(forbidden["type"], "error", "path: {path}"); + assert_eq!( + forbidden["error"]["type"], "permission_error", + "path: {path}" + ); + + let too_large = build_local_http_error_response_with_request_path( + "trace-too-large", + None, + Some(path), + http::StatusCode::PAYLOAD_TOO_LARGE, + "too large", + ) + .expect("payload-too-large response should build"); + let too_large = response_json(too_large).await; + assert_eq!(too_large["type"], "error", "path: {path}"); + assert_eq!( + too_large["error"]["type"], "request_too_large", + "path: {path}" + ); + } + } } diff --git a/apps/aether-gateway/src/control/auth/credentials.rs b/apps/aether-gateway/src/control/auth/credentials.rs index 9105baf44..23f93dc58 100644 --- a/apps/aether-gateway/src/control/auth/credentials.rs +++ b/apps/aether-gateway/src/control/auth/credentials.rs @@ -71,6 +71,20 @@ pub(super) fn extract_request_credentials( } } +pub(in crate::control) fn resolve_gateway_credential_carrier( + headers: &http::HeaderMap, + uri: &Uri, + auth_endpoint_signature: &str, +) -> Option { + extract_request_credentials(headers, uri, auth_endpoint_signature) + .primary + .map(|credential| match credential { + GatewayPrimaryCredential::ProviderApiKey { carrier, .. } + | GatewayPrimaryCredential::BearerToken { carrier, .. } + | GatewayPrimaryCredential::CookieHeader { carrier, .. } => carrier, + }) +} + fn has_trusted_gateway_marker(headers: &http::HeaderMap) -> bool { header_value_str(headers, crate::constants::GATEWAY_HEADER) .unwrap_or_default() diff --git a/apps/aether-gateway/src/control/auth/mod.rs b/apps/aether-gateway/src/control/auth/mod.rs index 412832acf..1aa699566 100644 --- a/apps/aether-gateway/src/control/auth/mod.rs +++ b/apps/aether-gateway/src/control/auth/mod.rs @@ -5,6 +5,7 @@ mod resolution; mod types; pub(crate) use credentials::extract_requested_model; +pub(super) use credentials::resolve_gateway_credential_carrier; pub(crate) use gate::{ execution_plan_balance_capacity_rejection, request_model_local_rejection, should_buffer_request_for_local_auth, trusted_auth_local_rejection, GatewayLocalAuthRejection, @@ -14,3 +15,4 @@ pub(crate) use resolution::{ GatewayAdminPrincipalContext, GatewayControlAuthContext, }; pub(super) use resolution::{resolve_control_decision_auth, ControlDecisionAuthResolution}; +pub(crate) use types::GatewayCredentialCarrier; diff --git a/apps/aether-gateway/src/control/auth/types.rs b/apps/aether-gateway/src/control/auth/types.rs index 930e7be41..5773edfaa 100644 --- a/apps/aether-gateway/src/control/auth/types.rs +++ b/apps/aether-gateway/src/control/auth/types.rs @@ -1,5 +1,5 @@ #[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(super) enum GatewayCredentialCarrier { +pub(crate) enum GatewayCredentialCarrier { AuthorizationBearer, XApiKey, ApiKey, @@ -8,6 +8,26 @@ pub(super) enum GatewayCredentialCarrier { CookieHeader, } +impl GatewayCredentialCarrier { + pub(crate) const fn as_str(self) -> &'static str { + match self { + Self::AuthorizationBearer => "authorization_bearer", + Self::XApiKey => "x_api_key", + Self::ApiKey => "api_key", + Self::XGoogApiKey => "x_goog_api_key", + Self::QueryKey => "query_key", + Self::CookieHeader => "cookie_header", + } + } + + pub(crate) const fn request_auth_channel(self) -> &'static str { + match self { + Self::AuthorizationBearer | Self::CookieHeader => "bearer_like", + Self::XApiKey | Self::ApiKey | Self::XGoogApiKey | Self::QueryKey => "api_key", + } + } +} + #[derive(Debug, Clone, PartialEq)] pub(super) struct GatewayTrustedAuthHeaders { pub(super) user_id: String, diff --git a/apps/aether-gateway/src/control/mod.rs b/apps/aether-gateway/src/control/mod.rs index edf284777..f62651d4a 100644 --- a/apps/aether-gateway/src/control/mod.rs +++ b/apps/aether-gateway/src/control/mod.rs @@ -12,7 +12,7 @@ pub(crate) use auth::{ refresh_execution_runtime_auth_context, request_model_local_rejection, resolve_execution_runtime_auth_context, should_buffer_request_for_local_auth, trusted_auth_local_rejection, GatewayAdminPrincipalContext, GatewayControlAuthContext, - GatewayLocalAuthRejection, + GatewayCredentialCarrier, GatewayLocalAuthRejection, }; pub(crate) use execute::{allows_control_execute_emergency, maybe_execute_via_control}; pub(crate) use management_token_permissions::{ diff --git a/apps/aether-gateway/src/control/route/ai.rs b/apps/aether-gateway/src/control/route/ai.rs index 65a4781ef..6d36f9a3b 100644 --- a/apps/aether-gateway/src/control/route/ai.rs +++ b/apps/aether-gateway/src/control/route/ai.rs @@ -1,7 +1,8 @@ use super::{ - classified, classified_with_request_auth_channel, is_claude_cli_request, is_gemini_cli_request, - is_gemini_models_route, is_gemini_operation_route, ClassifiedRoute, + classified, classified_with_request_auth_channel, detect_claude_client_surface, + is_gemini_cli_request, is_gemini_models_route, is_gemini_operation_route, ClassifiedRoute, }; +use crate::ai_serving::ApiOperation; pub(super) fn classify_ai_public_route( method: &http::Method, @@ -76,27 +77,33 @@ pub(super) fn classify_ai_public_route( true, )) } else if method == http::Method::POST && normalized_path == "/v1/messages/count_tokens" { - Some(classified( - "ai_public", - "claude", - "count_tokens", - "claude:messages", - false, - )) + let request_auth_channel = claude_request_auth_channel(headers); + Some( + classified_with_request_auth_channel( + "ai_public", + "claude", + "count_tokens", + request_auth_channel, + "claude:messages", + true, + ) + .with_client_surface(detect_claude_client_surface(headers)) + .with_api_operation(ApiOperation::ClaudeCountTokens), + ) } else if method == http::Method::POST && normalized_path == "/v1/messages" { - let request_auth_channel = if is_claude_cli_request(headers) { - "bearer_like" - } else { - "api_key" - }; - Some(classified_with_request_auth_channel( - "ai_public", - "claude", - "messages", - request_auth_channel, - "claude:messages", - true, - )) + let request_auth_channel = claude_request_auth_channel(headers); + Some( + classified_with_request_auth_channel( + "ai_public", + "claude", + "messages", + request_auth_channel, + "claude:messages", + true, + ) + .with_client_surface(detect_claude_client_surface(headers)) + .with_api_operation(ApiOperation::ClaudeMessagesCreate), + ) } else if normalized_path.starts_with("/v1/videos") { Some(classified( "ai_public", @@ -178,6 +185,20 @@ pub(super) fn classify_ai_public_route( } } +fn claude_request_auth_channel(headers: &http::HeaderMap) -> &'static str { + if crate::headers::header_value_str(headers, "x-api-key").is_some() + || crate::headers::header_value_str(headers, "api-key").is_some() + { + "api_key" + } else if crate::headers::header_value_str(headers, http::header::AUTHORIZATION.as_str()) + .is_some_and(|value| value.trim().to_ascii_lowercase().starts_with("bearer ")) + { + "bearer_like" + } else { + "api_key" + } +} + fn is_gemini_operation_method(method: &http::Method, normalized_path: &str) -> bool { method == http::Method::GET || (method == http::Method::POST && normalized_path.ends_with(":cancel")) diff --git a/apps/aether-gateway/src/control/route/mod.rs b/apps/aether-gateway/src/control/route/mod.rs index bd1c513a5..c18db0f34 100644 --- a/apps/aether-gateway/src/control/route/mod.rs +++ b/apps/aether-gateway/src/control/route/mod.rs @@ -1,5 +1,6 @@ use axum::http::Uri; +use crate::ai_serving::{ApiOperation, ClientSurface}; use crate::headers::header_value_str; use crate::{AppState, GatewayError}; @@ -9,7 +10,10 @@ mod internal; mod oauth; mod public_support; -use super::auth::{resolve_control_decision_auth, ControlDecisionAuthResolution}; +use super::auth::{ + resolve_control_decision_auth, resolve_gateway_credential_carrier, + ControlDecisionAuthResolution, GatewayCredentialCarrier, +}; use super::{GatewayAdminPrincipalContext, GatewayControlAuthContext, GatewayLocalAuthRejection}; #[derive(Debug, Clone)] @@ -19,6 +23,9 @@ pub(crate) struct GatewayControlDecision { pub(crate) route_class: Option, pub(crate) route_family: Option, pub(crate) route_kind: Option, + pub(crate) client_surface: Option, + pub(crate) api_operation: Option, + pub(crate) gateway_credential_carrier: Option, pub(crate) request_auth_channel: Option, pub(crate) auth_endpoint_signature: Option, pub(crate) execution_runtime_candidate: bool, @@ -42,6 +49,9 @@ impl GatewayControlDecision { route_class, route_family, route_kind, + client_surface: None, + api_operation: None, + gateway_credential_carrier: None, request_auth_channel: None, auth_endpoint_signature, execution_runtime_candidate: false, @@ -80,6 +90,8 @@ pub(super) struct ClassifiedRoute { route_family: &'static str, route_kind: &'static str, request_auth_channel: Option<&'static str>, + client_surface: Option, + api_operation: Option, auth_endpoint_signature: String, execution_runtime_candidate: bool, } @@ -96,6 +108,8 @@ pub(super) fn classified( route_family, route_kind, request_auth_channel: None, + client_surface: None, + api_operation: None, auth_endpoint_signature: auth_endpoint_signature.into(), execution_runtime_candidate, } @@ -114,11 +128,25 @@ pub(super) fn classified_with_request_auth_channel( route_family, route_kind, request_auth_channel: Some(request_auth_channel), + client_surface: None, + api_operation: None, auth_endpoint_signature: auth_endpoint_signature.into(), execution_runtime_candidate, } } +impl ClassifiedRoute { + pub(super) fn with_client_surface(mut self, client_surface: ClientSurface) -> Self { + self.client_surface = Some(client_surface); + self + } + + pub(super) fn with_api_operation(mut self, api_operation: ApiOperation) -> Self { + self.api_operation = Some(api_operation); + self + } +} + impl ClassifiedRoute { fn into_decision(self, public_path: String) -> GatewayControlDecision { GatewayControlDecision { @@ -127,6 +155,9 @@ impl ClassifiedRoute { route_class: Some(self.route_class.to_string()), route_family: Some(self.route_family.to_string()), route_kind: Some(self.route_kind.to_string()), + client_surface: self.client_surface, + api_operation: self.api_operation, + gateway_credential_carrier: None, request_auth_channel: self.request_auth_channel.map(str::to_string), auth_endpoint_signature: Some(self.auth_endpoint_signature), execution_runtime_candidate: self.execution_runtime_candidate, @@ -183,7 +214,17 @@ pub(crate) fn classify_control_route( .or_else(|| internal::classify_internal_route(method, &normalized_path)) .or_else(|| ai::classify_ai_public_route(method, &normalized_path, headers))?; - Some(classified.into_decision(normalized_path)) + let mut decision = classified.into_decision(normalized_path); + if let Some(signature) = decision.auth_endpoint_signature.as_deref() { + decision.gateway_credential_carrier = + resolve_gateway_credential_carrier(headers, uri, signature); + } + if decision.route_family.as_deref() == Some("claude") { + if let Some(carrier) = decision.gateway_credential_carrier { + decision.request_auth_channel = Some(carrier.request_auth_channel().to_string()); + } + } + Some(decision) } pub(super) fn detect_public_models_auth_signature(uri: &Uri, headers: &http::HeaderMap) -> String { @@ -220,11 +261,26 @@ pub(super) fn detect_public_models_auth_signature(uri: &Uri, headers: &http::Hea "openai:chat".to_string() } -pub(super) fn is_claude_cli_request(headers: &http::HeaderMap) -> bool { - let auth_header = header_value_str(headers, http::header::AUTHORIZATION.as_str()) +pub(super) fn detect_claude_client_surface(headers: &http::HeaderMap) -> ClientSurface { + let user_agent = header_value_str(headers, http::header::USER_AGENT.as_str()) .unwrap_or_default() .to_ascii_lowercase(); - auth_header.starts_with("bearer ") + let x_app_is_cli = header_value_str(headers, "x-app") + .is_some_and(|value| value.trim().eq_ignore_ascii_case("cli")); + if user_agent.contains("claude-code") + || user_agent.contains("claude-cli") + || user_agent.contains("claude code") + || x_app_is_cli + || header_value_str(headers, "x-claude-code-session-id").is_some() + { + ClientSurface::ClaudeCode + } else if user_agent.contains("anthropic/") + || header_value_str(headers, "x-stainless-lang").is_some() + { + ClientSurface::AnthropicSdk + } else { + ClientSurface::GenericCompatible + } } pub(super) fn is_gemini_cli_request(headers: &http::HeaderMap) -> bool { diff --git a/apps/aether-gateway/src/control/tests/ai.rs b/apps/aether-gateway/src/control/tests/ai.rs index c645cb54c..fbfde74bf 100644 --- a/apps/aether-gateway/src/control/tests/ai.rs +++ b/apps/aether-gateway/src/control/tests/ai.rs @@ -1,9 +1,11 @@ +use aether_ai_formats::{ApiOperation, ClientSurface}; use http::Uri; +use super::super::auth::GatewayCredentialCarrier; use super::{classify_control_route, headers}; #[test] -fn classifies_claude_count_tokens_as_non_execution_runtime_public_route() { +fn classifies_claude_count_tokens_as_execution_runtime_operation() { let headers = headers(&[("x-api-key", "sk-test")]); let uri: Uri = "/v1/messages/count_tokens" .parse() @@ -17,7 +19,11 @@ fn classifies_claude_count_tokens_as_non_execution_runtime_public_route() { decision.auth_endpoint_signature.as_deref(), Some("claude:messages") ); - assert!(!decision.is_execution_runtime_candidate()); + assert!(decision.is_execution_runtime_candidate()); + assert_eq!( + decision.api_operation, + Some(ApiOperation::ClaudeCountTokens) + ); } #[test] @@ -141,7 +147,7 @@ fn classifies_models_list_as_claude_when_headers_match() { } #[test] -fn classifies_claude_messages_cli_when_bearer_without_api_key() { +fn bearer_auth_does_not_imply_claude_code_client_surface() { let headers = headers(&[("authorization", "Bearer token-123")]); let uri: Uri = "/v1/messages".parse().expect("uri should parse"); let decision = @@ -149,6 +155,10 @@ fn classifies_claude_messages_cli_when_bearer_without_api_key() { assert_eq!(decision.route_family.as_deref(), Some("claude")); assert_eq!(decision.route_kind.as_deref(), Some("messages")); + assert_eq!( + decision.client_surface, + Some(ClientSurface::GenericCompatible) + ); assert_eq!( decision.request_auth_channel.as_deref(), Some("bearer_like") @@ -161,7 +171,7 @@ fn classifies_claude_messages_cli_when_bearer_without_api_key() { } #[test] -fn classifies_claude_messages_cli_when_bearer_is_present_even_with_api_key() { +fn claude_api_key_carrier_keeps_precedence_over_bearer() { let headers = headers(&[ ("authorization", "Bearer token-123"), ("x-api-key", "sk-client"), @@ -172,9 +182,10 @@ fn classifies_claude_messages_cli_when_bearer_is_present_even_with_api_key() { assert_eq!(decision.route_family.as_deref(), Some("claude")); assert_eq!(decision.route_kind.as_deref(), Some("messages")); + assert_eq!(decision.request_auth_channel.as_deref(), Some("api_key")); assert_eq!( - decision.request_auth_channel.as_deref(), - Some("bearer_like") + decision.gateway_credential_carrier, + Some(GatewayCredentialCarrier::XApiKey) ); assert_eq!( decision.auth_endpoint_signature.as_deref(), @@ -183,6 +194,63 @@ fn classifies_claude_messages_cli_when_bearer_is_present_even_with_api_key() { assert!(decision.is_execution_runtime_candidate()); } +#[test] +fn detects_claude_code_independently_from_bearer_auth() { + let headers = headers(&[ + ("authorization", "Bearer token-123"), + ("user-agent", "Claude-Code/2.1.0"), + ]); + let uri: Uri = "/v1/messages".parse().expect("uri should parse"); + let decision = + classify_control_route(&http::Method::POST, &uri, &headers).expect("route should classify"); + + assert_eq!(decision.client_surface, Some(ClientSurface::ClaudeCode)); + assert_eq!( + decision.gateway_credential_carrier, + Some(GatewayCredentialCarrier::AuthorizationBearer) + ); + assert_eq!( + decision.api_operation, + Some(ApiOperation::ClaudeMessagesCreate) + ); +} + +#[test] +fn detects_current_claude_cli_user_agent() { + let headers = headers(&[ + ("x-api-key", "sk-client"), + ("user-agent", "claude-cli/2.1.161 (external, cli)"), + ]); + let uri: Uri = "/v1/messages".parse().expect("uri should parse"); + let decision = + classify_control_route(&http::Method::POST, &uri, &headers).expect("route should classify"); + + assert_eq!(decision.client_surface, Some(ClientSurface::ClaudeCode)); + assert_eq!(decision.request_auth_channel.as_deref(), Some("api_key")); + assert_eq!( + decision.gateway_credential_carrier, + Some(GatewayCredentialCarrier::XApiKey) + ); +} + +#[test] +fn detects_claude_code_from_explicit_x_app_signal() { + let headers = headers(&[ + ("x-api-key", "sk-client"), + ("user-agent", "rewritten-by-proxy"), + ("x-app", "cli"), + ]); + let uri: Uri = "/v1/messages".parse().expect("uri should parse"); + let decision = + classify_control_route(&http::Method::POST, &uri, &headers).expect("route should classify"); + + assert_eq!(decision.client_surface, Some(ClientSurface::ClaudeCode)); + assert_eq!( + decision.api_operation, + Some(ApiOperation::ClaudeMessagesCreate) + ); +} + #[test] fn classifies_claude_messages_when_api_key_without_bearer() { let headers = headers(&[("x-api-key", "sk-client")]); diff --git a/apps/aether-gateway/src/dispatch/pool_scheduler.rs b/apps/aether-gateway/src/dispatch/pool_scheduler.rs index cd814d37c..f4eecd995 100644 --- a/apps/aether-gateway/src/dispatch/pool_scheduler.rs +++ b/apps/aether-gateway/src/dispatch/pool_scheduler.rs @@ -375,6 +375,10 @@ impl<'a> PoolKeyCursor<'a> { self.group.candidate.provider_id.as_str() } + pub(crate) fn endpoint_id(&self) -> &str { + self.group.candidate.endpoint_id.as_str() + } + pub(crate) fn new( state: PlannerAppState<'a>, group: EligibleLocalExecutionCandidate, diff --git a/apps/aether-gateway/src/execution_runtime/grok.rs b/apps/aether-gateway/src/execution_runtime/grok.rs index 5691bdb8b..5ec901d43 100644 --- a/apps/aether-gateway/src/execution_runtime/grok.rs +++ b/apps/aether-gateway/src/execution_runtime/grok.rs @@ -223,8 +223,8 @@ async fn execute_grok_app_chat( let elapsed_ms = started_at.elapsed().as_millis() as u64; if !(200..300).contains(&status_code) { - let decoded = decode_response_body_bytes(&headers, &raw_body).unwrap_or(raw_body); - let text = String::from_utf8_lossy(&decoded).to_string(); + let decoded = decode_response_body_bytes(&headers, &raw_body)?; + let text = String::from_utf8_lossy(decoded.as_ref()).to_string(); return Ok(GrokCollected { status_code, headers, @@ -274,8 +274,8 @@ async fn execute_grok_app_chat_stream( &mut adapter, ) .await?; - let decoded = decode_response_body_bytes(&headers, &raw_body).unwrap_or(raw_body); - let text = String::from_utf8_lossy(&decoded).to_string(); + let decoded = decode_response_body_bytes(&headers, &raw_body)?; + let text = String::from_utf8_lossy(decoded.as_ref()).to_string(); let elapsed_ms = started_at.elapsed().as_millis() as u64; let collected = GrokCollected { status_code, diff --git a/apps/aether-gateway/src/execution_runtime/mod.rs b/apps/aether-gateway/src/execution_runtime/mod.rs index 0bc96f8bf..4f074b315 100644 --- a/apps/aether-gateway/src/execution_runtime/mod.rs +++ b/apps/aether-gateway/src/execution_runtime/mod.rs @@ -42,7 +42,67 @@ pub(crate) use self::response_header_rules::{ pub(crate) use crate::orchestration::{ append_local_failover_policy_to_value, LocalFailoverAnalysis, LocalFailoverDecision, }; +pub(crate) use aether_ai_serving::AdaptationMode; pub(crate) use aether_ai_serving::{ConversionMode, ExecutionStrategy}; + +pub(crate) fn ai_attempt_retry_scope_from_failure_disposition( + disposition: crate::orchestration::FailureDisposition, +) -> aether_ai_serving::AiAttemptRetryScope { + use crate::orchestration::{FailureRetryAction, FailureScope}; + use aether_ai_serving::AiAttemptRetryScope; + + match disposition.failure_scope { + FailureScope::Credential | FailureScope::CredentialModel => AiAttemptRetryScope::Credential, + FailureScope::Endpoint => AiAttemptRetryScope::Endpoint, + FailureScope::Provider => AiAttemptRetryScope::Provider, + FailureScope::None => match disposition.retry_action { + FailureRetryAction::NextCredential => AiAttemptRetryScope::Credential, + FailureRetryAction::NextEndpoint => AiAttemptRetryScope::Endpoint, + FailureRetryAction::Stop + | FailureRetryAction::SameCredential + | FailureRetryAction::NextCandidate => AiAttemptRetryScope::Candidate, + }, + } +} + +#[cfg(test)] +mod retry_scope_tests { + use aether_ai_serving::AiAttemptRetryScope; + + use super::ai_attempt_retry_scope_from_failure_disposition; + use crate::orchestration::{classify_failure_disposition, LocalFailoverClassification}; + + #[test] + fn anthropic_failure_scope_survives_runtime_mapping() { + let retry_scope = |status_code| { + ai_attempt_retry_scope_from_failure_disposition(classify_failure_disposition( + "claude:messages", + LocalFailoverClassification::RetryUpstreamFailure, + status_code, + )) + }; + + assert_eq!(retry_scope(429), AiAttemptRetryScope::Credential); + assert_eq!(retry_scope(500), AiAttemptRetryScope::Endpoint); + assert_eq!(retry_scope(529), AiAttemptRetryScope::Provider); + assert_eq!(retry_scope(400), AiAttemptRetryScope::Candidate); + } + + #[test] + fn non_anthropic_retry_keeps_existing_candidate_order() { + let disposition = classify_failure_disposition( + "openai:chat", + LocalFailoverClassification::RetryUpstreamFailure, + 429, + ); + + assert_eq!( + ai_attempt_retry_scope_from_failure_disposition(disposition), + AiAttemptRetryScope::Candidate + ); + assert!(!disposition.preserve_upstream_error); + } +} pub use server::{ build_execution_runtime_router, build_execution_runtime_router_with_request_concurrency_limit, build_execution_runtime_router_with_request_gates, serve_execution_runtime_tcp, @@ -57,12 +117,14 @@ pub async fn prewarm_direct_h2c_sender_cache_from_env_for_startup( .map_err(|err| err.to_string()) } -pub(crate) use stream::execute_execution_runtime_stream; +pub(crate) use stream::{ + execute_execution_runtime_stream, execute_execution_runtime_stream_with_retry_scope, +}; pub(crate) use stream_pump::build_direct_execution_frame_stream; pub(crate) use sync::{ - execute_execution_runtime_sync, maybe_build_local_sync_finalize_response, - maybe_build_local_video_error_response, maybe_build_local_video_success_outcome, - resolve_local_sync_error_background_report_kind, + execute_execution_runtime_sync, execute_execution_runtime_sync_with_retry_scope, + maybe_build_local_sync_finalize_response, maybe_build_local_video_error_response, + maybe_build_local_video_success_outcome, resolve_local_sync_error_background_report_kind, resolve_local_sync_success_background_report_kind, LocalVideoSyncSuccessBuild, LocalVideoSyncSuccessOutcome, }; @@ -220,6 +282,16 @@ pub(crate) fn append_execution_contract_fields( "provider_contract".to_string(), Value::String(provider_contract.to_string()), ); + let default_adaptation_mode = if execution_strategy == ExecutionStrategy::LocalCrossFormat + || conversion_mode != ConversionMode::None + { + AdaptationMode::CrossFormat + } else { + AdaptationMode::NativeTransparent + }; + object + .entry("adaptation_mode".to_string()) + .or_insert_with(|| Value::String(default_adaptation_mode.as_str().to_string())); } pub(crate) fn append_execution_contract_fields_to_value( @@ -263,6 +335,7 @@ mod tests { assert_eq!(value["conversion_mode"], "bidirectional"); assert_eq!(value["client_contract"], "openai:chat"); assert_eq!(value["provider_contract"], "gemini:generate_content"); + assert_eq!(value["adaptation_mode"], "cross_format"); assert_eq!(value["provider_api_format"], "gemini:generate_content"); } } diff --git a/apps/aether-gateway/src/execution_runtime/oauth_retry.rs b/apps/aether-gateway/src/execution_runtime/oauth_retry.rs index ca40e38f3..2334ab3d0 100644 --- a/apps/aether-gateway/src/execution_runtime/oauth_retry.rs +++ b/apps/aether-gateway/src/execution_runtime/oauth_retry.rs @@ -1,6 +1,10 @@ use aether_contracts::ExecutionPlan; use tracing::warn; +use crate::orchestration::{ + oauth_status_may_be_invalid as status_may_be_oauth_invalid, + oauth_status_proves_access_token_invalid as status_proves_access_token_invalid, +}; use crate::state::AgentIdentityAuthConfigFence; use crate::{provider_transport::LocalOAuthRefreshError, AppState}; @@ -70,17 +74,17 @@ pub(crate) async fn refresh_oauth_plan_auth_for_retry( // A bearer-token response cannot authorize refreshing an Agent Identity // installed under the same key id while the request was in flight. return false; - } else if transport - .provider - .provider_type - .trim() - .eq_ignore_ascii_case("codex") - && transport.key.auth_type.trim().eq_ignore_ascii_case("oauth") - && !request_authorization.is_some_and(|authorization| { - bearer_authorization_matches_transport(authorization, &transport) - }) - { - return false; + } else if aether_provider_transport::supports_local_generic_oauth_request_auth_resolution( + &transport, + ) { + if let Some(current_authorization) = generic_oauth_transport_authorization(&transport) { + if !request_authorization.is_some_and(|authorization| { + authorizations_use_same_access_token(authorization, ¤t_authorization) + }) { + replace_execution_plan_authorization(plan, current_authorization); + return true; + } + } } if transport.key.decrypted_auth_config.is_none() @@ -164,63 +168,31 @@ fn execution_plan_authorization(plan: &ExecutionPlan) -> Option<&str> { .map(|(_, value)| value.as_str()) } -fn bearer_authorization_matches_transport( - authorization: &str, +fn generic_oauth_transport_authorization( transport: &aether_provider_transport::GatewayProviderTransportSnapshot, -) -> bool { - let current_token = transport.key.decrypted_api_key.trim(); - !current_token.is_empty() - && authorization - .trim() - .strip_prefix("Bearer ") - .map(str::trim) - .is_some_and(|token| token == current_token) +) -> Option { + aether_provider_transport::resolve_local_generic_oauth_transport_authorization(transport) } -fn status_may_be_oauth_invalid(status_code: u16, response_text: Option<&str>) -> bool { - if status_code == 401 { - return true; +fn authorizations_use_same_access_token(left: &str, right: &str) -> bool { + match (bearer_access_token(left), bearer_access_token(right)) { + (Some(left), Some(right)) => left == right, + _ => left.trim() == right.trim(), } - if status_code != 403 { - return false; - } - - let Some(response_text) = response_text else { - return true; - }; - let response_text = response_text.to_ascii_lowercase(); - ["oauth", "token", "auth", "credential", "expired"] - .iter() - .any(|needle| response_text.contains(needle)) } -fn status_proves_access_token_invalid(status_code: u16, response_text: Option<&str>) -> bool { - if status_code == 401 { - return true; - } - if status_code != 403 { - return false; - } +fn bearer_access_token(authorization: &str) -> Option<&str> { + let mut parts = authorization.split_ascii_whitespace(); + let scheme = parts.next()?; + let token = parts.next()?; + (scheme.eq_ignore_ascii_case("bearer") && parts.next().is_none()).then_some(token) +} - let Some(response_text) = response_text else { - return false; - }; - let response_text = response_text.to_ascii_lowercase(); - [ - "oauth_token_invalid", - "invalid_token", - "invalid access token", - "access token invalid", - "access token expired", - "expired access token", - "authentication token has been invalidated", - "token has been invalidated", - "personal access token owner is inactive", - "biscuit_baker_service_auth_credential_error_status", - "security token included in the request is expired", - ] - .iter() - .any(|needle| response_text.contains(needle)) +fn replace_execution_plan_authorization(plan: &mut ExecutionPlan, authorization: String) { + plan.headers + .retain(|name, _| !name.eq_ignore_ascii_case("authorization")); + plan.headers + .insert("authorization".to_string(), authorization); } #[cfg(test)] @@ -230,14 +202,15 @@ mod tests { status_proves_access_token_invalid, }; use std::collections::BTreeMap; + use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use aether_contracts::{ExecutionPlan, RequestBody}; use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY}; use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository; use aether_data_contracts::repository::provider_catalog::{ - ProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey, - StoredProviderCatalogProvider, + ProviderCatalogReadRepository, ProviderCatalogWriteRepository, + StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider, }; use axum::routing::post; use axum::{extract::Request, Json, Router}; @@ -252,8 +225,65 @@ mod tests { 403, Some("The security token included in the request is expired") )); - assert!(status_may_be_oauth_invalid(403, None)); + assert!(status_may_be_oauth_invalid( + 403, + Some("oauth_token_invalid") + )); + assert!(!status_may_be_oauth_invalid(403, None)); + assert!(!status_may_be_oauth_invalid( + 403, + Some( + r#"{"type":"error","error":{"type":"permission_error","message":"this token is not authorized for the workspace"}}"# + ) + )); + assert!(!status_may_be_oauth_invalid( + 403, + Some( + r#"{"error":{"type":"permission_error","message":"the authentication token has been invalidated for this workspace"}}"# + ) + )); + assert!(status_may_be_oauth_invalid( + 403, + Some( + r#"{"type":"error","error":{"type":"authentication_error","message":"credential expired"}}"# + ) + )); + assert!(status_may_be_oauth_invalid( + 403, + Some(r#"{"error":{"type":"oauth_token_invalid","message":"sign in again"}}"#) + )); + assert!(status_may_be_oauth_invalid( + 403, + Some( + r#"{"error":{"code":"biscuit_baker_service_auth_credential_error_status","message":"Personal access token owner is inactive."}}"# + ) + )); + assert!(!status_may_be_oauth_invalid( + 403, + Some( + r#"{"error":{"type":"invalid_request_error","message":"Your authentication token has been invalidated. Please try signing in again."}}"# + ) + )); + assert!(!status_may_be_oauth_invalid( + 403, + Some( + r#"{"error":{"type":"invalid_request_error","message":"invalid request: token budget is invalid"}}"# + ) + )); assert!(!status_may_be_oauth_invalid(403, Some("quota exceeded"))); + assert!(!status_may_be_oauth_invalid( + 403, + Some("invalid request: max token budget is invalid") + )); + assert!(!status_may_be_oauth_invalid( + 403, + Some("invalid_token_budget") + )); + assert!(!status_may_be_oauth_invalid(403, Some("not authorized"))); + assert!(!status_may_be_oauth_invalid( + 403, + Some("authorization denied") + )); assert!(!status_may_be_oauth_invalid(429, Some("token bucket"))); } @@ -282,7 +312,7 @@ mod tests { } #[tokio::test] - async fn auto_removes_request_proven_oauth_failure_after_terminal_refresh_failure() { + async fn retains_codex_key_after_request_proven_terminal_refresh_failure() { let token_hits = Arc::new(Mutex::new(0usize)); let token_hits_clone = Arc::clone(&token_hits); let token_server = Router::new().route( @@ -437,7 +467,233 @@ mod tests { .list_keys_by_ids(&["key-codex-oauth-retry".to_string()]) .await .expect("keys should read"); - assert!(keys.is_empty()); + assert_eq!(keys.len(), 1); + assert!(keys[0].oauth_invalid_at_unix_secs.is_some()); + assert!(keys[0] + .oauth_invalid_reason + .as_deref() + .is_some_and(|reason| reason.contains("[REFRESH_FAILED]") + && reason.contains("Token 续期失败 (401)"))); + + token_handle.abort(); + } + + #[tokio::test] + async fn stale_claude_code_request_reuses_rotated_access_token_without_second_refresh() { + let refresh_hits = Arc::new(AtomicUsize::new(0)); + let refresh_hits_for_server = Arc::clone(&refresh_hits); + let token_server = Router::new().route( + "/oauth/token", + post(move |_request: Request| { + let hits = Arc::clone(&refresh_hits_for_server); + async move { + hits.fetch_add(1, Ordering::SeqCst); + Json(json!({ + "access_token": "fresh-claude-access-token", + "refresh_token": "fresh-claude-refresh-token", + "expires_in": 3600, + "token_type": "Bearer" + })) + } + }), + ); + + let provider = StoredProviderCatalogProvider::new( + "provider-claude-code".to_string(), + "Claude Code".to_string(), + Some("https://api.anthropic.com".to_string()), + "claude_code".to_string(), + ) + .expect("provider should build"); + let endpoint = StoredProviderCatalogEndpoint::new( + "endpoint-claude-code".to_string(), + "provider-claude-code".to_string(), + "claude:messages".to_string(), + None, + None, + true, + ) + .expect("endpoint should build") + .with_transport_fields( + "https://api.anthropic.com".to_string(), + None, + None, + None, + None, + None, + None, + None, + ) + .expect("endpoint transport should build"); + let encrypted_api_key = encrypt_python_fernet_plaintext( + DEVELOPMENT_ENCRYPTION_KEY, + "stale-claude-access-token", + ) + .expect("api key ciphertext should build"); + let encrypted_auth_config = encrypt_python_fernet_plaintext( + DEVELOPMENT_ENCRYPTION_KEY, + r#"{"provider_type":"claude_code","access_token":"stale-claude-access-token","refresh_token":"stale-claude-refresh-token","expires_at":4102444800}"#, + ) + .expect("auth config ciphertext should build"); + let mut key = StoredProviderCatalogKey::new( + "key-claude-code".to_string(), + "provider-claude-code".to_string(), + "Claude OAuth".to_string(), + "oauth".to_string(), + None, + true, + ) + .expect("key should build") + .with_transport_fields( + Some(json!(["claude:messages"])), + encrypted_api_key, + Some(encrypted_auth_config), + None, + None, + None, + None, + None, + None, + ) + .expect("key transport should build"); + key.expires_at_unix_secs = Some(4_102_444_800); + + let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( + vec![provider], + vec![endpoint], + vec![key], + )); + let (token_url, token_handle) = start_test_server(token_server).await; + let oauth_refresh = + crate::provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![ + Arc::new( + crate::provider_transport::oauth_refresh::GenericOAuthRefreshAdapter::default() + .with_token_url_for_tests( + "claude_code", + format!("{token_url}/oauth/token"), + ), + ), + ]); + let state = crate::AppState::new() + .expect("state should build") + .with_data_state_for_tests( + crate::data::GatewayDataState::with_provider_catalog_repository_for_tests( + provider_catalog_repository.clone(), + ) + .with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY), + ) + .with_oauth_refresh_coordinator_for_tests(oauth_refresh); + let stale_transport = state + .read_provider_transport_snapshot( + "provider-claude-code", + "endpoint-claude-code", + "key-claude-code", + ) + .await + .expect("stale transport should load") + .expect("stale transport should exist"); + let stale_plan = ExecutionPlan { + request_id: "req-claude-oauth-fence".to_string(), + candidate_id: None, + provider_name: Some("claude_code".to_string()), + provider_id: "provider-claude-code".to_string(), + endpoint_id: "endpoint-claude-code".to_string(), + key_id: "key-claude-code".to_string(), + method: "POST".to_string(), + url: "https://api.anthropic.com/v1/messages".to_string(), + headers: BTreeMap::from([( + "authorization".to_string(), + "Bearer stale-claude-access-token".to_string(), + )]), + content_type: Some("application/json".to_string()), + content_encoding: None, + body: RequestBody::from_json(json!({"model": "claude-sonnet-4-5"})), + stream: false, + 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, + }; + + let mut first_plan = stale_plan.clone(); + assert!( + refresh_oauth_plan_auth_for_retry( + &state, + &mut first_plan, + 401, + Some(r#"{"error":"invalid_token"}"#), + "trace-claude-oauth-fence-first", + ) + .await + ); + assert_eq!( + first_plan.headers.get("authorization").map(String::as_str), + Some("Bearer fresh-claude-access-token") + ); + assert_eq!(refresh_hits.load(Ordering::SeqCst), 1); + + let stale_force_result = state + .force_local_oauth_refresh_entry(&stale_transport) + .await + .expect("stale force should reuse the persisted winner") + .expect("stale force should return the winner entry"); + assert_eq!( + stale_force_result.auth_header_value, + "Bearer fresh-claude-access-token" + ); + assert_eq!(refresh_hits.load(Ordering::SeqCst), 1); + + let mut stale_in_flight_plan = stale_plan; + assert!( + refresh_oauth_plan_auth_for_retry( + &state, + &mut stale_in_flight_plan, + 401, + Some(r#"{"error":"invalid_token"}"#), + "trace-claude-oauth-fence-stale", + ) + .await + ); + assert_eq!( + stale_in_flight_plan + .headers + .get("authorization") + .map(String::as_str), + Some("Bearer fresh-claude-access-token") + ); + assert_eq!(refresh_hits.load(Ordering::SeqCst), 1); + + let mut admin_replacement = provider_catalog_repository + .list_keys_by_ids(&["key-claude-code".to_string()]) + .await + .expect("Claude key should load") + .pop() + .expect("Claude key should exist"); + admin_replacement.encrypted_api_key = Some( + encrypt_python_fernet_plaintext( + DEVELOPMENT_ENCRYPTION_KEY, + "admin-claude-access-token", + ) + .expect("admin access token should encrypt"), + ); + admin_replacement.expires_at_unix_secs = Some(4_102_444_800); + provider_catalog_repository + .update_key(&admin_replacement) + .await + .expect("admin replacement should persist"); + + let admin_result = state + .force_local_oauth_refresh_entry(&stale_transport) + .await + .expect("stale force should reuse the admin replacement") + .expect("admin replacement should resolve"); + assert_eq!( + admin_result.auth_header_value, + "Bearer admin-claude-access-token" + ); + assert_eq!(refresh_hits.load(Ordering::SeqCst), 1); token_handle.abort(); } diff --git a/apps/aether-gateway/src/execution_runtime/server.rs b/apps/aether-gateway/src/execution_runtime/server.rs index f2a8bf8c0..ddc00cf04 100644 --- a/apps/aether-gateway/src/execution_runtime/server.rs +++ b/apps/aether-gateway/src/execution_runtime/server.rs @@ -375,6 +375,8 @@ impl IntoResponse for ExecutionRuntimeAppError { | ExecutionRuntimeTransportError::BrowserClientBuild(_) | ExecutionRuntimeTransportError::BrowserBody(_) | ExecutionRuntimeTransportError::UpstreamRequest(_) + | ExecutionRuntimeTransportError::UpstreamResponseTooLarge { .. } + | ExecutionRuntimeTransportError::UpstreamResponseDecode { .. } | ExecutionRuntimeTransportError::RelayError(_) | ExecutionRuntimeTransportError::InvalidJson(_), ) => StatusCode::BAD_GATEWAY, diff --git a/apps/aether-gateway/src/execution_runtime/stream/commit_policy.rs b/apps/aether-gateway/src/execution_runtime/stream/commit_policy.rs new file mode 100644 index 000000000..4fa24d156 --- /dev/null +++ b/apps/aether-gateway/src/execution_runtime/stream/commit_policy.rs @@ -0,0 +1,527 @@ +use std::time::Duration; + +use serde_json::Value; + +use crate::execution_runtime::MAX_STREAM_PREFETCH_BYTES; + +const ANTHROPIC_PRECOMMIT_MAX_WAIT: Duration = Duration::from_millis(750); + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum StreamCommitPolicy { + OnResponseHeaders, + OnFirstClassifiedBody, + OnFirstAnthropicSemanticEvent { + max_bytes: usize, + max_wait: Duration, + }, +} + +impl StreamCommitPolicy { + #[allow(clippy::too_many_arguments)] + pub(super) fn for_response( + has_direct_finalize: bool, + content_type: Option<&str>, + provider_api_format: &str, + client_api_format: &str, + has_private_stream_normalizer: bool, + has_local_stream_rewriter: bool, + force_prefetch: bool, + ) -> Self { + if !has_direct_finalize { + return Self::OnFirstClassifiedBody; + } + + if force_prefetch { + return Self::OnFirstClassifiedBody; + } + + let content_type = content_type + .map(str::trim) + .filter(|value| !value.is_empty()) + .unwrap_or_default() + .to_ascii_lowercase(); + if content_type.contains("text/event-stream") { + if provider_api_format.eq_ignore_ascii_case("claude:messages") + && provider_api_format.eq_ignore_ascii_case(client_api_format) + && !has_private_stream_normalizer + && !has_local_stream_rewriter + { + return Self::OnFirstAnthropicSemanticEvent { + max_bytes: MAX_STREAM_PREFETCH_BYTES, + max_wait: ANTHROPIC_PRECOMMIT_MAX_WAIT, + }; + } + return Self::OnResponseHeaders; + } + + if has_private_stream_normalizer || has_local_stream_rewriter { + return Self::OnFirstClassifiedBody; + } + + if !provider_api_format.eq_ignore_ascii_case(client_api_format) { + return Self::OnFirstClassifiedBody; + } + + if content_type.is_empty() { + return Self::OnResponseHeaders; + } + + if content_type.contains("json") || content_type.ends_with("+json") { + Self::OnFirstClassifiedBody + } else { + Self::OnResponseHeaders + } + } + + pub(super) const fn commits_on_response_headers(self) -> bool { + matches!(self, Self::OnResponseHeaders) + } + + pub(super) const fn requires_bounded_frame_wait(self) -> bool { + matches!(self, Self::OnFirstAnthropicSemanticEvent { .. }) + } + + pub(super) const fn max_precommit_wait(self) -> Option { + match self { + Self::OnFirstAnthropicSemanticEvent { max_wait, .. } => Some(max_wait), + Self::OnResponseHeaders | Self::OnFirstClassifiedBody => None, + } + } + + pub(super) const fn is_native_anthropic(self) -> bool { + matches!(self, Self::OnFirstAnthropicSemanticEvent { .. }) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum StreamCommitState { + Uncommitted, + Committed, + Terminal, +} + +#[derive(Debug, PartialEq)] +pub(super) enum StreamPrecommitObservation { + Pending, + Commit, + UpstreamError { status_code: u16, body_json: Value }, +} + +#[derive(Debug)] +pub(super) struct StreamCommitGate { + policy: StreamCommitPolicy, + state: StreamCommitState, + observed_bytes: usize, + anthropic: AnthropicSsePrecommitInspector, +} + +impl StreamCommitGate { + pub(super) fn new(policy: StreamCommitPolicy) -> Self { + let state = if policy.commits_on_response_headers() { + StreamCommitState::Committed + } else { + StreamCommitState::Uncommitted + }; + Self { + policy, + state, + observed_bytes: 0, + anthropic: AnthropicSsePrecommitInspector::default(), + } + } + + pub(super) const fn state(&self) -> StreamCommitState { + self.state + } + + pub(super) const fn is_uncommitted(&self) -> bool { + matches!(self.state, StreamCommitState::Uncommitted) + } + + pub(super) fn observe_provider_bytes(&mut self, chunk: &[u8]) -> StreamPrecommitObservation { + if self.state != StreamCommitState::Uncommitted { + return StreamPrecommitObservation::Commit; + } + + let StreamCommitPolicy::OnFirstAnthropicSemanticEvent { max_bytes, .. } = self.policy + else { + return StreamPrecommitObservation::Pending; + }; + + self.observed_bytes = self.observed_bytes.saturating_add(chunk.len()); + match self.anthropic.observe(chunk, max_bytes) { + AnthropicSseObservation::Pending => {} + AnthropicSseObservation::SemanticEvent => { + self.state = StreamCommitState::Committed; + return StreamPrecommitObservation::Commit; + } + AnthropicSseObservation::Error(body_json) => { + self.state = StreamCommitState::Terminal; + return StreamPrecommitObservation::UpstreamError { + status_code: anthropic_error_status_code(&body_json), + body_json, + }; + } + } + + if self.observed_bytes >= max_bytes { + self.commit(); + StreamPrecommitObservation::Commit + } else { + StreamPrecommitObservation::Pending + } + } + + pub(super) fn commit(&mut self) { + if self.state == StreamCommitState::Uncommitted { + self.state = StreamCommitState::Committed; + } + } +} + +#[derive(Debug)] +enum AnthropicSseObservation { + Pending, + SemanticEvent, + Error(Value), +} + +#[derive(Debug, Default)] +struct AnthropicSsePrecommitInspector { + buffered: Vec, +} + +impl AnthropicSsePrecommitInspector { + fn observe(&mut self, chunk: &[u8], max_bytes: usize) -> AnthropicSseObservation { + let remaining = max_bytes.saturating_sub(self.buffered.len()); + let truncated = chunk.len() > remaining; + self.buffered + .extend_from_slice(&chunk[..chunk.len().min(remaining)]); + + while let Some((record_end, separator_len)) = find_sse_record_boundary(&self.buffered) { + let record = self.buffered[..record_end].to_vec(); + self.buffered.drain(..record_end + separator_len); + match classify_anthropic_sse_record(&record) { + AnthropicSseObservation::Pending => {} + decision => return decision, + } + } + + if truncated { + AnthropicSseObservation::SemanticEvent + } else { + AnthropicSseObservation::Pending + } + } +} + +pub(super) fn find_sse_record_boundary(buffer: &[u8]) -> Option<(usize, usize)> { + let mut cursor = 0; + while cursor < buffer.len() { + let (line_end, line_ending_len) = next_sse_line_ending(buffer, cursor)?; + let next_line_start = line_end + line_ending_len; + let Some((next_line_end, next_line_ending_len)) = + next_sse_line_ending(buffer, next_line_start) + else { + return None; + }; + if next_line_end == next_line_start { + return Some(( + line_end, + line_ending_len.saturating_add(next_line_ending_len), + )); + } + cursor = next_line_start; + } + None +} + +fn next_sse_line_ending(buffer: &[u8], start: usize) -> Option<(usize, usize)> { + let relative = buffer + .get(start..)? + .iter() + .position(|byte| matches!(byte, b'\r' | b'\n'))?; + let index = start + relative; + let ending_len = if buffer[index] == b'\r' && buffer.get(index + 1) == Some(&b'\n') { + 2 + } else { + 1 + }; + Some((index, ending_len)) +} + +fn classify_anthropic_sse_record(record: &[u8]) -> AnthropicSseObservation { + let Ok(record) = std::str::from_utf8(record) else { + return AnthropicSseObservation::Pending; + }; + let normalized_record = record.replace("\r\n", "\n").replace('\r', "\n"); + let mut event_type = None; + let mut data = String::new(); + for line in normalized_record.lines() { + if line.starts_with(':') { + continue; + } + if let Some(value) = line.strip_prefix("event:") { + let value = value.trim(); + if !value.is_empty() { + event_type = Some(value); + } + continue; + } + if let Some(value) = line.strip_prefix("data:") { + if !data.is_empty() { + data.push('\n'); + } + data.push_str(value.trim_start()); + } + } + if data.trim().is_empty() { + return AnthropicSseObservation::Pending; + } + + let Ok(body_json) = serde_json::from_str::(data.trim()) else { + return AnthropicSseObservation::Pending; + }; + let payload_type = body_json.get("type").and_then(Value::as_str).map(str::trim); + if event_type == Some("error") || payload_type == Some("error") { + return AnthropicSseObservation::Error(body_json); + } + + let semantic_type = match (event_type, payload_type) { + (Some(event_type), Some(payload_type)) if event_type == payload_type => Some(event_type), + (None, Some(payload_type)) => Some(payload_type), + _ => None, + }; + if semantic_type.is_some_and(is_anthropic_semantic_event_type) { + AnthropicSseObservation::SemanticEvent + } else { + AnthropicSseObservation::Pending + } +} + +fn is_anthropic_semantic_event_type(event_type: &str) -> bool { + matches!( + event_type, + "message_start" + | "content_block_start" + | "content_block_delta" + | "content_block_stop" + | "message_delta" + | "message_stop" + ) +} + +pub(super) fn anthropic_error_status_code(body_json: &Value) -> u16 { + let error_type = body_json + .get("error") + .and_then(|error| error.get("type")) + .or_else(|| body_json.get("type")) + .and_then(Value::as_str) + .map(str::trim) + .unwrap_or_default(); + match error_type { + "invalid_request_error" => 400, + "authentication_error" => 401, + "permission_error" => 403, + "not_found_error" => 404, + "request_too_large" => 413, + "rate_limit_error" => 429, + "overloaded_error" => 529, + _ => 500, + } +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use super::{ + anthropic_error_status_code, StreamCommitGate, StreamCommitPolicy, StreamCommitState, + StreamPrecommitObservation, + }; + + fn native_anthropic_policy() -> StreamCommitPolicy { + StreamCommitPolicy::OnFirstAnthropicSemanticEvent { + max_bytes: 16_384, + max_wait: Duration::from_millis(750), + } + } + + #[test] + fn policy_selects_bounded_anthropic_gate_only_for_native_same_format_sse() { + let native = StreamCommitPolicy::for_response( + true, + Some("text/event-stream; charset=utf-8"), + "claude:messages", + "claude:messages", + false, + false, + false, + ); + assert!(native.is_native_anthropic()); + assert_eq!( + native.max_precommit_wait(), + Some(Duration::from_millis(750)) + ); + assert!(StreamCommitPolicy::for_response( + true, + Some("text/event-stream"), + "openai:chat", + "claude:messages", + false, + false, + false, + ) + .commits_on_response_headers()); + assert!(StreamCommitPolicy::for_response( + true, + Some("text/event-stream"), + "claude:messages", + "claude:messages", + false, + true, + false, + ) + .commits_on_response_headers()); + } + + #[test] + fn gate_detects_anthropic_error_across_every_chunk_boundary() { + let event = b"event: error\r\ndata: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",\"message\":\"busy\"}}\r\n\r\n"; + for split in 1..event.len() { + let mut gate = StreamCommitGate::new(native_anthropic_policy()); + let first_observation = gate.observe_provider_bytes(&event[..split]); + if matches!( + first_observation, + StreamPrecommitObservation::UpstreamError { + status_code: 529, + .. + } + ) { + assert_eq!(event[split - 1], b'\r'); + } else { + assert_eq!(first_observation, StreamPrecommitObservation::Pending); + assert!(matches!( + gate.observe_provider_bytes(&event[split..]), + StreamPrecommitObservation::UpstreamError { + status_code: 529, + .. + } + )); + } + assert_eq!(gate.state(), StreamCommitState::Terminal); + } + } + + #[test] + fn gate_detects_cr_only_and_mixed_line_ending_errors() { + for event in [ + "event: error\rdata: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\"}}\r\r", + "event: error\r\ndata: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\"}}\n\r", + ] { + for split in 1..event.len() { + let mut gate = StreamCommitGate::new(native_anthropic_policy()); + assert_eq!( + gate.observe_provider_bytes(&event.as_bytes()[..split]), + StreamPrecommitObservation::Pending, + "gate committed before complete mixed-line event at split {split}", + ); + assert!(matches!( + gate.observe_provider_bytes(&event.as_bytes()[split..]), + StreamPrecommitObservation::UpstreamError { + status_code: 529, + .. + } + )); + } + } + } + + #[test] + fn unknown_and_ping_events_do_not_commit_before_anthropic_error() { + let mut gate = StreamCommitGate::new(native_anthropic_policy()); + assert_eq!( + gate.observe_provider_bytes( + b"event: future_event\ndata: {\"type\":\"future_event\",\"value\":1}\n\n" + ), + StreamPrecommitObservation::Pending + ); + assert_eq!( + gate.observe_provider_bytes(b"event: ping\ndata: {\"type\":\"ping\"}\n\n"), + StreamPrecommitObservation::Pending + ); + assert!(matches!( + gate.observe_provider_bytes( + b"event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"rate_limit_error\"}}\n\n" + ), + StreamPrecommitObservation::UpstreamError { + status_code: 429, + .. + } + )); + } + + #[test] + fn first_semantic_event_commits_before_later_error_in_same_chunk() { + let mut gate = StreamCommitGate::new(native_anthropic_policy()); + let observation = gate.observe_provider_bytes( + concat!( + "event: message_start\n", + "data: {\"type\":\"message_start\",\"message\":{}}\n\n", + "event: error\n", + "data: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\"}}\n\n", + ) + .as_bytes(), + ); + + assert_eq!(observation, StreamPrecommitObservation::Commit); + assert_eq!(gate.state(), StreamCommitState::Committed); + } + + #[test] + fn transport_fragment_count_does_not_commit_an_incomplete_anthropic_error() { + let policy = StreamCommitPolicy::OnFirstAnthropicSemanticEvent { + max_bytes: 1024, + max_wait: Duration::from_millis(750), + }; + let mut gate = StreamCommitGate::new(policy); + let event = b"event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\"}}\n\n"; + for byte in &event[..event.len() - 1] { + assert_eq!( + gate.observe_provider_bytes(std::slice::from_ref(byte)), + StreamPrecommitObservation::Pending, + ); + } + assert!(matches!( + gate.observe_provider_bytes(&event[event.len() - 1..]), + StreamPrecommitObservation::UpstreamError { + status_code: 529, + .. + } + )); + } + + #[test] + fn anthropic_error_status_mapping_matches_messages_api_taxonomy() { + for (error_type, status_code) in [ + ("invalid_request_error", 400), + ("authentication_error", 401), + ("permission_error", 403), + ("not_found_error", 404), + ("request_too_large", 413), + ("rate_limit_error", 429), + ("overloaded_error", 529), + ("api_error", 500), + ] { + let body = serde_json::json!({ + "type": "error", + "error": { "type": error_type, "message": "upstream failure" } + }); + assert_eq!( + anthropic_error_status_code(&body), + status_code, + "unexpected status for {error_type}" + ); + } + } +} diff --git a/apps/aether-gateway/src/execution_runtime/stream/execution.rs b/apps/aether-gateway/src/execution_runtime/stream/execution.rs index 99ea7f037..75a45f8ba 100644 --- a/apps/aether-gateway/src/execution_runtime/stream/execution.rs +++ b/apps/aether-gateway/src/execution_runtime/stream/execution.rs @@ -3,11 +3,13 @@ use std::future::Future; use std::io::Error as IoError; use std::pin::Pin; use std::sync::{ - atomic::{AtomicBool, AtomicU64, Ordering}, + atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}, Arc, }; +use std::task::{Context, Poll}; use std::time::{Duration, Instant}; +use aether_ai_serving::{AiAttemptExecutionOutcome, AiAttemptRetryScope}; use aether_contracts::{ ExecutionPlan, ExecutionStreamTerminalSummary, ExecutionTelemetry, StandardizedUsage, StreamFrame, StreamFramePayload, @@ -30,15 +32,19 @@ use axum::body::{Body, Bytes}; use axum::http::Response; use base64::Engine as _; use futures_util::stream::{self as futures_stream, BoxStream}; -use futures_util::{StreamExt, TryStreamExt}; +use futures_util::{Stream, StreamExt, TryStreamExt}; use http_body_util::BodyExt; use serde_json::{json, Value}; +use tokio::io::{AsyncRead, ReadBuf}; use tokio::sync::mpsc; use tokio::time::MissedTickBehavior; use tokio_util::codec::{FramedRead, LinesCodec}; -use tokio_util::io::StreamReader; use tracing::{debug, info, warn}; +use super::commit_policy::{ + anthropic_error_status_code, find_sse_record_boundary, StreamCommitGate, StreamCommitPolicy, + StreamPrecommitObservation, +}; use super::error::{ build_synthetic_non_success_stream_error_body, collect_error_body, decode_stream_error_body, inspect_prefetched_stream_body, read_next_frame, @@ -92,8 +98,9 @@ use crate::execution_runtime::transport::{ }; use crate::execution_runtime::windsurf::maybe_execute_windsurf_stream; use crate::execution_runtime::{ - apply_endpoint_response_header_rules, attach_provider_response_headers_to_report_context, - local_failover_response_text, resolve_core_stream_direct_finalize_report_kind, + ai_attempt_retry_scope_from_failure_disposition, apply_endpoint_response_header_rules, + attach_provider_response_headers_to_report_context, local_failover_response_text, + resolve_core_stream_direct_finalize_report_kind, resolve_core_stream_error_finalize_report_kind, resolve_local_candidate_failover_analysis_stream, should_fallback_to_control_stream, should_retry_next_local_candidate_stream, LocalFailoverDecision, @@ -103,12 +110,13 @@ use crate::execution_runtime::{ }; use crate::log_ids::short_request_id; use crate::orchestration::{ - apply_local_execution_effect, build_local_error_flow_metadata, cyber_continue_failover_enabled, - trace_upstream_response_body, with_error_flow_report_context, - with_upstream_response_report_context, LocalAdaptiveRateLimitEffect, - LocalAdaptiveSuccessEffect, LocalAttemptFailureEffect, LocalExecutionEffect, - LocalExecutionEffectContext, LocalHealthFailureEffect, LocalHealthSuccessEffect, - LocalOAuthInvalidationEffect, LocalPoolErrorEffect, + apply_local_execution_effect, build_local_error_flow_metadata, classify_failure_disposition, + cyber_continue_failover_enabled, trace_upstream_response_body, with_error_flow_report_context, + with_upstream_response_report_context, FailureDisposition, FailureTokenAction, + LocalAdaptiveRateLimitEffect, LocalAdaptiveSuccessEffect, LocalAttemptFailureEffect, + LocalExecutionEffect, LocalExecutionEffectContext, LocalFailoverAnalysis, + LocalHealthFailureEffect, LocalHealthSuccessEffect, LocalOAuthInvalidationEffect, + LocalPoolErrorEffect, }; use crate::provider_pool_demand::{ acquire_provider_pool_in_flight_guard, ProviderPoolInFlightGuard, @@ -138,10 +146,17 @@ 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; const SSE_TERMINAL_DETECTOR_MAX_LINE_BYTES: usize = 1024 * 1024; +const SSE_TERMINAL_DETECTOR_MAX_RECORD_BYTES: usize = SSE_TERMINAL_DETECTOR_MAX_LINE_BYTES; const PROVIDER_STREAM_ERROR_INSPECTION_MAX_BYTES: usize = SSE_TERMINAL_DETECTOR_MAX_LINE_BYTES; const STREAM_IDLE_LOG_INTERVAL: Duration = Duration::from_secs(60); const STREAM_IDLE_LOG_INTERVAL_MS: u64 = 60_000; const REWRITTEN_STREAM_PREFETCH_TIMEOUT: Duration = Duration::from_millis(750); +const OAUTH_ERROR_PREFETCH_MAX_WAIT: Duration = Duration::from_millis(750); +const ANTHROPIC_POST_STOP_DRAIN_MAX_WAIT: Duration = Duration::from_millis(250); +const ANTHROPIC_POST_STOP_DRAIN_MAX_FRAMES: usize = 8; +const ANTHROPIC_POST_STOP_DRAIN_MAX_BYTES: usize = 64 * 1024; +const POST_STOP_FRAME_READ_BUDGET_INACTIVE: usize = usize::MAX; +const POST_STOP_MAX_EMPTY_CHUNKS_PER_POLL: usize = 32; const DEFAULT_DIRECT_PASSTHROUGH_CHANNEL_CAPACITY: usize = 16; const MAX_DIRECT_PASSTHROUGH_CHANNEL_CAPACITY: usize = 1024; const DIRECT_PASSTHROUGH_CHANNEL_CAPACITY_ENV: &str = @@ -915,17 +930,37 @@ fn observe_stream_usage_bytes( buffered: &mut Vec, chunk: &[u8], ) { - if chunk.is_empty() { + if chunk.is_empty() + || observer + .latest_summary() + .and_then(|summary| summary.parser_error.as_deref()) + .is_some() + { return; } - buffered.extend_from_slice(chunk); - while let Some(line_end) = buffered.iter().position(|byte| *byte == b'\n') { - let line = buffered.drain(..=line_end).collect::>(); - if let Err(err) = observer.push_line(report_context, line) { - observer.disable_with_error(err.to_string()); + let mut remaining = chunk; + while !remaining.is_empty() { + let line_part_len = remaining + .iter() + .position(|byte| *byte == b'\n') + .map_or(remaining.len(), |index| index + 1); + if buffered.len().saturating_add(line_part_len) > SSE_TERMINAL_DETECTOR_MAX_LINE_BYTES { + observer.disable_with_error(format!( + "stream usage event exceeded {SSE_TERMINAL_DETECTOR_MAX_LINE_BYTES} bytes" + )); buffered.clear(); - break; + return; + } + buffered.extend_from_slice(&remaining[..line_part_len]); + remaining = &remaining[line_part_len..]; + if buffered.last() == Some(&b'\n') { + let line = std::mem::take(buffered); + if let Err(err) = observer.push_line(report_context, line) { + observer.disable_with_error(err.to_string()); + buffered.clear(); + return; + } } } } @@ -1122,18 +1157,53 @@ async fn execute_in_process_stream_with_oauth_retry( ) -> Result { let mut execution = execute_in_process_stream(state, plan, trace_id).await?; apply_stream_summary_report_context(&mut execution, report_context); - let response_text = if execution.status_code == 401 + let uses_oauth_credential = stream_plan_uses_oauth_credential(state, plan).await; + let embedded_oauth_credential = execution.status_code == 200 + && plan + .provider_api_format + .eq_ignore_ascii_case("claude:messages") + && uses_oauth_credential; + let prefetched_failure = if embedded_oauth_credential { + prefetch_direct_anthropic_stream_failure(&mut execution, plan, report_context).await + } else { + None + }; + let analyzed_prefetched_failure = match prefetched_failure { + Some(failure) => { + Some(analyze_prefetched_stream_failure(state, plan, report_context, failure).await) + } + None => None, + }; + let response_text = if let Some(failure) = analyzed_prefetched_failure.as_ref() { + Some(failure.response_text.clone()) + } else if execution.status_code == 403 && uses_oauth_credential { + prefetch_direct_stream_error_body(&mut execution).await + } else if execution.status_code == 401 && stream_plan_uses_codex_agent_identity(state, plan).await { prefetch_direct_stream_error_body(&mut execution).await } else { None }; - if execution.status_code >= 400 + let retry_status_code = analyzed_prefetched_failure + .as_ref() + .map(|failure| failure.status_code) + .unwrap_or(execution.status_code); + let retry_requested = + analyzed_prefetched_failure + .as_ref() + .map_or(execution.status_code >= 400, |failure| { + matches!( + failure.disposition.token_action, + FailureTokenAction::ForceRefresh + ) + }); + if retry_requested + && uses_oauth_credential && refresh_oauth_plan_auth_for_retry( state, plan, - execution.status_code, + retry_status_code, response_text.as_deref(), trace_id, ) @@ -1146,6 +1216,196 @@ async fn execute_in_process_stream_with_oauth_retry( Ok(execution) } +#[derive(Debug)] +struct PrefetchedStreamFailure { + status_code: u16, + response_text: String, +} + +#[derive(Debug)] +struct AnalyzedPrefetchedStreamFailure { + status_code: u16, + response_text: String, + #[allow(dead_code)] + analysis: LocalFailoverAnalysis, + disposition: FailureDisposition, +} + +async fn analyze_prefetched_stream_failure( + state: &AppState, + plan: &ExecutionPlan, + report_context: Option<&Value>, + failure: PrefetchedStreamFailure, +) -> AnalyzedPrefetchedStreamFailure { + let analysis = resolve_local_candidate_failover_analysis_stream( + state, + plan, + report_context, + failure.status_code, + Some(failure.response_text.as_str()), + ) + .await; + let disposition = classify_failure_disposition( + plan.provider_api_format.as_str(), + analysis.classification, + failure.status_code, + ); + AnalyzedPrefetchedStreamFailure { + status_code: failure.status_code, + response_text: failure.response_text, + analysis, + disposition, + } +} + +async fn stream_plan_uses_oauth_credential(state: &AppState, plan: &ExecutionPlan) -> bool { + state + .read_provider_transport_snapshot(&plan.provider_id, &plan.endpoint_id, &plan.key_id) + .await + .ok() + .flatten() + .as_ref() + .is_some_and(|transport| { + aether_provider_transport::auth::resolve_local_auth_type_for_transport_format(transport) + .eq_ignore_ascii_case("oauth") + }) +} + +async fn prefetch_direct_anthropic_stream_failure( + execution: &mut DirectUpstreamStreamExecution, + plan: &ExecutionPlan, + report_context: Option<&Value>, +) -> Option { + if execution.status_code != 200 { + return None; + } + let normalized_stream_report_context = + normalize_provider_private_report_context(report_context); + let policy = StreamCommitPolicy::for_response( + true, + execution.headers.get("content-type").map(String::as_str), + plan.provider_api_format.as_str(), + plan.client_api_format.as_str(), + maybe_build_provider_private_stream_normalizer(report_context).is_some(), + maybe_build_stream_response_rewriter(normalized_stream_report_context.as_ref()).is_some(), + false, + ); + if !policy.is_native_anthropic() { + return None; + } + + let mut gate = StreamCommitGate::new(policy); + let precommit_started_at = Instant::now(); + let max_wait = policy.max_precommit_wait()?; + let mut observed_first_body = execution + .prefetched_body + .iter() + .any(|item| item.as_ref().is_ok_and(|chunk| !chunk.is_empty())); + while gate.is_uncommitted() { + let wait = select_direct_anthropic_prefetch_wait( + precommit_started_at, + max_wait, + execution.started_at, + execution.stream_first_byte_timeout, + observed_first_body, + Instant::now(), + ); + if wait.remaining.is_zero() { + if wait.commit_on_timeout { + gate.commit(); + } + break; + } + let next_chunk = match tokio::time::timeout( + wait.remaining, + next_direct_upstream_response_chunk(&mut execution.response), + ) + .await + { + Ok(result) => result, + Err(_) => { + if wait.commit_on_timeout { + gate.commit(); + } + break; + } + }; + let chunk = match next_chunk { + Ok(Some(chunk)) => chunk, + Ok(None) => break, + Err(error) => { + execution.prefetched_body.push_back(Err(error)); + break; + } + }; + if chunk.is_empty() { + continue; + } + observed_first_body = true; + execution.prefetched_body.push_back(Ok(chunk.clone())); + match gate.observe_provider_bytes(&chunk) { + StreamPrecommitObservation::Pending => {} + StreamPrecommitObservation::Commit => break, + StreamPrecommitObservation::UpstreamError { + status_code, + body_json, + } => { + let response_text = serde_json::to_string(&body_json) + .unwrap_or_else(|_| String::from_utf8_lossy(&chunk).into_owned()); + return Some(PrefetchedStreamFailure { + status_code, + response_text, + }); + } + } + } + execution.stream_precommit_committed = !gate.is_uncommitted(); + None +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct DirectAnthropicPrefetchWait { + remaining: Duration, + commit_on_timeout: bool, +} + +fn select_direct_anthropic_prefetch_wait( + precommit_started_at: Instant, + max_precommit_wait: Duration, + upstream_started_at: Instant, + first_byte_timeout: Option, + observed_first_body: bool, + now: Instant, +) -> DirectAnthropicPrefetchWait { + let precommit_remaining = + max_precommit_wait.saturating_sub(now.saturating_duration_since(precommit_started_at)); + if observed_first_body { + return DirectAnthropicPrefetchWait { + remaining: precommit_remaining, + commit_on_timeout: true, + }; + } + let Some(first_byte_timeout) = first_byte_timeout else { + return DirectAnthropicPrefetchWait { + remaining: precommit_remaining, + commit_on_timeout: true, + }; + }; + let first_byte_remaining = + first_byte_timeout.saturating_sub(now.saturating_duration_since(upstream_started_at)); + if first_byte_remaining <= precommit_remaining { + DirectAnthropicPrefetchWait { + remaining: first_byte_remaining, + commit_on_timeout: false, + } + } else { + DirectAnthropicPrefetchWait { + remaining: precommit_remaining, + commit_on_timeout: true, + } + } +} + async fn stream_plan_uses_codex_agent_identity(state: &AppState, plan: &ExecutionPlan) -> bool { state .read_provider_transport_snapshot(&plan.provider_id, &plan.endpoint_id, &plan.key_id) @@ -1184,22 +1444,38 @@ async fn next_direct_upstream_response_chunk( async fn prefetch_direct_stream_error_body( execution: &mut DirectUpstreamStreamExecution, ) -> Option { + let prefetch_started_at = Instant::now(); let mut inspected = Vec::with_capacity(MAX_ERROR_BODY_BYTES); let mut fully_buffered = false; while inspected.len() < MAX_ERROR_BODY_BYTES { + let remaining = OAUTH_ERROR_PREFETCH_MAX_WAIT.saturating_sub(prefetch_started_at.elapsed()); + if remaining.is_zero() { + break; + } let next_chunk = if execution.prefetched_body.is_empty() { - match await_direct_passthrough_first_item( + match tokio::time::timeout( + remaining, + await_direct_passthrough_first_item( + next_direct_upstream_response_chunk(&mut execution.response), + execution.started_at, + execution.stream_first_byte_timeout, + ), + ) + .await + { + Ok(Ok(item)) => item, + Ok(Err(_)) | Err(_) => break, + } + } else { + match tokio::time::timeout( + remaining, next_direct_upstream_response_chunk(&mut execution.response), - execution.started_at, - execution.stream_first_byte_timeout, ) .await { Ok(item) => item, Err(_) => break, } - } else { - next_direct_upstream_response_chunk(&mut execution.response).await }; let chunk = match next_chunk { Ok(Some(chunk)) => chunk, @@ -1361,6 +1637,7 @@ async fn forward_direct_passthrough_client_chunk( downstream_dropped: &mut bool, client_visible_stream_completed: &mut bool, client_stream_completion_tracker: &mut ClientVisibleStreamCompletionTracker, + observe_stream_completion: bool, client_stream_bytes: &mut u64, buffered_body: &mut Vec, client_body_truncated: &mut bool, @@ -1404,8 +1681,10 @@ async fn forward_direct_passthrough_client_chunk( observe_gateway_stage_ms("direct_passthrough_first_client_send_wait", send_wait_ms); } - *client_visible_stream_completed |= - client_stream_completion_tracker.observe_chunk(chunk.as_ref()); + if observe_stream_completion { + *client_visible_stream_completed |= + client_stream_completion_tracker.observe_chunk(chunk.as_ref()); + } *client_stream_bytes = client_stream_bytes.saturating_add(chunk_len); *last_client_chunk_elapsed_ms = stream_started_at .elapsed() @@ -1445,6 +1724,7 @@ struct DirectPassthroughFinalizerCore { provider_body_truncated: bool, client_body_truncated: bool, client_stream_completion_tracker: ClientVisibleStreamCompletionTracker, + requires_anthropic_message_stop: bool, client_visible_stream_completed: bool, usage_stream_telemetry: Option, telemetry: Option, @@ -1495,6 +1775,43 @@ impl DirectPassthroughFinalizer { self.core_mut().terminal_failure = Some(failure); } + fn prepare_upstream_chunk(&mut self, mut chunk: Bytes) -> Option { + let core = self.core_mut(); + if !core.requires_anthropic_message_stop { + return Some(chunk); + } + if core.client_visible_stream_completed { + return None; + } + if let Some(terminal_end) = core + .client_stream_completion_tracker + .observe_anthropic_message_stop_terminal_end(chunk.as_ref()) + { + chunk.truncate(terminal_end); + core.client_visible_stream_completed = true; + } + Some(chunk) + } + + fn fail_if_anthropic_message_stop_missing(&mut self) { + let core = self.core_mut(); + if core.requires_anthropic_message_stop + && !core.client_visible_stream_completed + && core.terminal_failure.is_none() + { + core.terminal_failure = Some(build_anthropic_premature_eof_failure( + "upstream Anthropic stream ended before message_stop", + )); + } + } + + fn completed_native_anthropic_stream(&self) -> bool { + let core = self.core(); + core.requires_anthropic_message_stop + && core.client_visible_stream_completed + && core.terminal_failure.is_none() + } + fn log_terminal_error_event_encode_failed(&self, err: impl std::fmt::Debug) { let core = self.core(); warn!( @@ -1570,8 +1887,11 @@ impl DirectPassthroughFinalizer { .provider_error_inspection .observe(core.stream_usage_report_context.as_ref(), chunk.as_ref()) { - let error_status_code = - resolve_local_sync_error_status_code(core.status_code, &error_body_json); + let error_status_code = resolve_provider_stream_error_status_code( + core.plan.provider_api_format.as_str(), + core.status_code, + &error_body_json, + ); core.terminal_failure = Some(build_stream_failure_from_provider_error_body( error_status_code, &error_body_json, @@ -1590,9 +1910,11 @@ impl DirectPassthroughFinalizer { DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES, &mut core.client_body_truncated, ); - core.client_visible_stream_completed |= core - .client_stream_completion_tracker - .observe_chunk(chunk.as_ref()); + if !core.requires_anthropic_message_stop { + core.client_visible_stream_completed |= core + .client_stream_completion_tracker + .observe_chunk(chunk.as_ref()); + } core.client_stream_bytes = core .client_stream_bytes .saturating_add(u64::try_from(chunk.len()).unwrap_or(u64::MAX)); @@ -1777,6 +2099,7 @@ impl DirectPassthroughFinalizerCore { provider_body_truncated, client_body_truncated, client_stream_completion_tracker: _, + requires_anthropic_message_stop: _, client_visible_stream_completed, usage_stream_telemetry, telemetry, @@ -2127,6 +2450,16 @@ impl DirectPassthroughInlineBodyState { if self.finalized { return None; } + if self + .finalizer + .as_ref() + .is_some_and(DirectPassthroughFinalizer::completed_native_anthropic_stream) + { + self.upstream.take(); + self.finalized = true; + drop(self.finalizer.take()); + return None; + } if !self.observed_first_body_poll { self.observed_first_body_poll = true; if let Some(finalizer) = self.finalizer.as_mut() { @@ -2162,17 +2495,40 @@ impl DirectPassthroughInlineBodyState { } let observed_at = Instant::now(); - if let Some(finalizer) = self.finalizer.as_mut() { + let Some(chunk) = self + .finalizer + .as_mut() + .and_then(|finalizer| finalizer.prepare_upstream_chunk(chunk)) + else { + continue; + }; + let provider_error_detected = if let Some(finalizer) = self.finalizer.as_mut() { finalizer.observe_upstream_chunk(&chunk, observed_at); - } + finalizer.terminal_failure().is_some() + } else { + false + }; if let Some(client_chunk) = filter_upstream_sse_control_chunk(&mut self.upstream_control_filter, chunk) { self.prepare_client_chunk_yield(&client_chunk); + self.terminal_error_sent |= provider_error_detected; + if self + .finalizer + .as_ref() + .is_some_and(DirectPassthroughFinalizer::completed_native_anthropic_stream) + { + self.upstream.take(); + self.upstream_done = true; + } return Some((Ok(client_chunk), self)); } } + if let Some(finalizer) = self.finalizer.as_mut() { + finalizer.fail_if_anthropic_message_stop_missing(); + } + if !self.control_filter_flushed && self .finalizer @@ -2193,7 +2549,8 @@ impl DirectPassthroughInlineBodyState { if let Some(finalizer) = self.finalizer.as_mut() { if let Some(failure) = finalizer.terminal_failure() { self.terminal_error_sent = true; - match encode_terminal_sse_error_event(failure) { + match encode_terminal_sse_error_event_for_plan(&finalizer.core().plan, failure) + { Ok(error_event) => { self.prepare_client_chunk_yield(&error_event); return Some((Ok(error_event), self)); @@ -2256,6 +2613,19 @@ impl DirectPassthroughInlineBodyState { let Some(finalizer) = self.finalizer.as_mut() else { return; }; + if finalizer.completed_native_anthropic_stream() { + let core = finalizer.core(); + debug!( + event_name = "direct_passthrough_read_error_ignored_after_anthropic_stop", + log_type = "debug", + trace_id = %core.trace_id, + request_id = %core.request_id_for_log, + candidate_id = ?core.candidate_id.as_deref(), + error = %message, + "gateway ignored direct passthrough teardown error after Anthropic message_stop" + ); + return; + } let core = finalizer.core(); warn!( event_name = "direct_passthrough_body_read_error", @@ -2394,12 +2764,22 @@ async fn execute_stream_from_direct_passthrough( provider_api_format: _, stream_summary_report_context: _, prefetched_body, + stream_precommit_committed: _, response, started_at: upstream_started_at, stream_first_byte_timeout, upstream_target_permit, } = execution; + let requires_anthropic_message_stop = status_code == 200 + && response_headers_indicate_sse(&headers) + && plan + .provider_api_format + .eq_ignore_ascii_case("claude:messages") + && plan + .client_api_format + .eq_ignore_ascii_case("claude:messages"); + let request_id = plan.request_id.clone(); let candidate_id = plan.candidate_id.clone(); let request_id_for_log = short_request_id(request_id.as_str()); @@ -2514,6 +2894,7 @@ async fn execute_stream_from_direct_passthrough( provider_body_truncated: false, client_body_truncated: false, client_stream_completion_tracker: ClientVisibleStreamCompletionTracker::default(), + requires_anthropic_message_stop, client_visible_stream_completed: false, usage_stream_telemetry: None, telemetry: None, @@ -2596,6 +2977,7 @@ async fn execute_stream_from_direct_passthrough( let mut client_body_truncated = false; let mut upstream_control_filter = Some(SseControlBlockFilter::default()); let mut client_stream_completion_tracker = ClientVisibleStreamCompletionTracker::default(); + let requires_anthropic_message_stop = requires_anthropic_message_stop; let mut client_visible_stream_completed = false; let mut usage_stream_telemetry: Option = None; let telemetry: Option = None; @@ -2604,6 +2986,7 @@ async fn execute_stream_from_direct_passthrough( let mut last_client_chunk_elapsed_ms = 0u64; let mut downstream_dropped = false; let mut terminal_failure: Option = None; + let mut provider_error_forwarded_to_client = false; let mut upstream = direct_upstream_response_byte_stream(prefetched_body, response); let mut observed_first_upstream_body = false; let mut observed_first_client_send = false; @@ -2658,6 +3041,18 @@ async fn execute_stream_from_direct_passthrough( let chunk = match item { Ok(chunk) => chunk, Err(message) => { + if requires_anthropic_message_stop && client_visible_stream_completed { + debug!( + event_name = "direct_passthrough_read_error_ignored_after_anthropic_stop", + log_type = "debug", + trace_id = %trace_id_owned, + request_id = %request_id_for_report_log, + candidate_id = ?candidate_id_for_report.as_deref(), + error = %message, + "gateway ignored direct passthrough teardown error after Anthropic message_stop" + ); + break; + } warn!( event_name = "direct_passthrough_body_read_error", log_type = "ops", @@ -2679,6 +3074,19 @@ async fn execute_stream_from_direct_passthrough( if chunk.is_empty() { continue; } + if requires_anthropic_message_stop && client_visible_stream_completed { + continue; + } + + let mut provider_chunk = chunk; + if requires_anthropic_message_stop { + if let Some(terminal_end) = client_stream_completion_tracker + .observe_anthropic_message_stop_terminal_end(provider_chunk.as_ref()) + { + provider_chunk.truncate(terminal_end); + client_visible_stream_completed = true; + } + } let observed_at = Instant::now(); if !observed_first_upstream_body { @@ -2711,16 +3119,18 @@ async fn execute_stream_from_direct_passthrough( ); } - let provider_chunk = chunk.clone(); - if let Some(client_chunk) = - filter_upstream_sse_control_chunk(&mut upstream_control_filter, chunk) - { - let sent_client_chunk = forward_direct_passthrough_client_chunk( + let mut sent_client_chunk = false; + if let Some(client_chunk) = filter_upstream_sse_control_chunk( + &mut upstream_control_filter, + provider_chunk.clone(), + ) { + sent_client_chunk = forward_direct_passthrough_client_chunk( &tx, client_chunk, &mut downstream_dropped, &mut client_visible_stream_completed, &mut client_stream_completion_tracker, + !requires_anthropic_message_stop, &mut client_stream_bytes, &mut buffered_body, &mut client_body_truncated, @@ -2767,14 +3177,34 @@ async fn execute_stream_from_direct_passthrough( provider_chunk.as_ref(), ); if let Some(error_body_json) = provider_private_error_body_json { - let error_status_code = - resolve_local_sync_error_status_code(status_code, &error_body_json); + provider_error_forwarded_to_client = sent_client_chunk; + let error_status_code = resolve_provider_stream_error_status_code( + plan_for_report.provider_api_format.as_str(), + status_code, + &error_body_json, + ); terminal_failure = Some(build_stream_failure_from_provider_error_body( error_status_code, &error_body_json, )); break; } + if requires_anthropic_message_stop && client_visible_stream_completed { + break; + } + } + drop(upstream); + drop(_provider_pool_in_flight_guard); + drop(_upstream_target_permit); + + if terminal_failure.is_none() + && !downstream_dropped + && requires_anthropic_message_stop + && !client_visible_stream_completed + { + terminal_failure = Some(build_anthropic_premature_eof_failure( + "upstream Anthropic stream ended before message_stop", + )); } if terminal_failure.is_none() { @@ -2787,6 +3217,7 @@ async fn execute_stream_from_direct_passthrough( &mut downstream_dropped, &mut client_visible_stream_completed, &mut client_stream_completion_tracker, + !requires_anthropic_message_stop, &mut client_stream_bytes, &mut buffered_body, &mut client_body_truncated, @@ -2802,8 +3233,11 @@ async fn execute_stream_from_direct_passthrough( } } - if let Some(failure) = terminal_failure.as_ref().filter(|_| !downstream_dropped) { - match encode_terminal_sse_error_event(failure) { + if let Some(failure) = terminal_failure + .as_ref() + .filter(|_| !downstream_dropped && !provider_error_forwarded_to_client) + { + match encode_terminal_sse_error_event_for_plan(&plan_for_report, failure) { Ok(error_event) => { let _ = forward_direct_passthrough_client_chunk( &tx, @@ -2811,6 +3245,7 @@ async fn execute_stream_from_direct_passthrough( &mut downstream_dropped, &mut client_visible_stream_completed, &mut client_stream_completion_tracker, + true, &mut client_stream_bytes, &mut buffered_body, &mut client_body_truncated, @@ -3111,7 +3546,14 @@ async fn execute_stream_from_direct_passthrough( } }); - let body_stream = build_sse_body_stream(Vec::new(), rx, false, false, SSE_KEEPALIVE_INTERVAL); + let body_stream = build_sse_body_stream( + Vec::new(), + rx, + false, + false, + requires_anthropic_message_stop, + SSE_KEEPALIVE_INTERVAL, + ); Ok(Some(build_client_response_from_parts( status_code, &headers, @@ -3139,9 +3581,52 @@ pub(crate) fn execute_execution_runtime_stream<'a>( plan_kind, report_kind, report_context, + None, + None, )) } +#[allow(clippy::too_many_arguments)] +pub(crate) fn execute_execution_runtime_stream_with_retry_scope<'a>( + state: &'a AppState, + plan: ExecutionPlan, + trace_id: &'a str, + decision: &'a GatewayControlDecision, + plan_kind: &'a str, + report_kind: Option, + report_context: Option, +) -> Pin< + Box< + dyn Future>, GatewayError>> + + Send + + 'a, + >, +> { + Box::pin(async move { + let mut retry_scope = AiAttemptRetryScope::Candidate; + let mut fallback_response = None; + let response = execute_execution_runtime_stream_inner( + state, + plan, + trace_id, + decision, + plan_kind, + report_kind, + report_context, + Some(&mut retry_scope), + Some(&mut fallback_response), + ) + .await?; + Ok(match response { + Some(response) => AiAttemptExecutionOutcome::Responded(response), + None => AiAttemptExecutionOutcome::Retry { + scope: retry_scope, + fallback_response, + }, + }) + }) +} + async fn execute_execution_runtime_stream_inner( state: &AppState, mut plan: ExecutionPlan, @@ -3150,6 +3635,8 @@ async fn execute_execution_runtime_stream_inner( plan_kind: &str, report_kind: Option, mut report_context: Option, + mut retry_scope_out: Option<&mut AiAttemptRetryScope>, + mut retry_fallback_out: Option<&mut Option>>, ) -> Result>, GatewayError> { let stream_started_at = Instant::now(); let mut stage_trace = RequestStageTrace::from_env(); @@ -3222,7 +3709,7 @@ async fn execute_execution_runtime_stream_inner( ); match maybe_execute_grok_stream(&plan, report_context.as_ref()).await { Ok(Some(grok_stream)) => { - return execute_stream_from_frame_stream( + return execute_stream_from_frame_stream_with_retry_scope( state, plan, trace_id, @@ -3235,7 +3722,10 @@ async fn execute_execution_runtime_stream_inner( stage_trace, lifecycle_pending_recorded, grok_stream.frame_stream, + false, provider_pool_in_flight_guard.take(), + retry_scope_out.as_deref_mut(), + retry_fallback_out.as_deref_mut(), ) .await; } @@ -3276,7 +3766,7 @@ async fn execute_execution_runtime_stream_inner( } match maybe_execute_windsurf_stream(state, &plan, report_context.as_ref()).await { Ok(Some(windsurf_stream)) => { - return execute_stream_from_frame_stream( + return execute_stream_from_frame_stream_with_retry_scope( state, plan, trace_id, @@ -3289,7 +3779,10 @@ async fn execute_execution_runtime_stream_inner( stage_trace, lifecycle_pending_recorded, windsurf_stream.frame_stream, + false, provider_pool_in_flight_guard.take(), + retry_scope_out.as_deref_mut(), + retry_fallback_out.as_deref_mut(), ) .await; } @@ -3330,7 +3823,7 @@ async fn execute_execution_runtime_stream_inner( } match maybe_execute_kiro_web_search_stream(state, &plan, report_context.as_ref()).await { Ok(Some(kiro_web_search)) => { - return execute_stream_from_frame_stream( + return execute_stream_from_frame_stream_with_retry_scope( state, plan, trace_id, @@ -3343,7 +3836,10 @@ async fn execute_execution_runtime_stream_inner( stage_trace, lifecycle_pending_recorded, kiro_web_search.frame_stream, + false, provider_pool_in_flight_guard.take(), + retry_scope_out.as_deref_mut(), + retry_fallback_out.as_deref_mut(), ) .await; } @@ -3384,7 +3880,7 @@ async fn execute_execution_runtime_stream_inner( } match maybe_execute_chatgpt_web_image_stream(state, &plan, report_context.as_ref()).await { Ok(Some(chatgpt_web_image)) => { - return execute_stream_from_frame_stream( + return execute_stream_from_frame_stream_with_retry_scope( state, plan, trace_id, @@ -3397,7 +3893,10 @@ async fn execute_execution_runtime_stream_inner( stage_trace, lifecycle_pending_recorded, chatgpt_web_image.frame_stream, + false, provider_pool_in_flight_guard.take(), + retry_scope_out.as_deref_mut(), + retry_fallback_out.as_deref_mut(), ) .await; } @@ -3525,8 +4024,9 @@ async fn execute_execution_runtime_stream_inner( record_stream_pending_lifecycle(state, seed, &mut stage_trace).await; lifecycle_pending_recorded = true; } + let stream_precommit_committed = execution.stream_precommit_committed; let frame_stream = build_direct_execution_frame_stream(execution).boxed(); - return execute_stream_from_frame_stream( + return execute_stream_from_frame_stream_with_retry_scope( state, plan, trace_id, @@ -3539,7 +4039,10 @@ async fn execute_execution_runtime_stream_inner( stage_trace, lifecycle_pending_recorded, frame_stream, + stream_precommit_committed, provider_pool_in_flight_guard.take(), + retry_scope_out.as_deref_mut(), + retry_fallback_out.as_deref_mut(), ) .await; } @@ -3641,8 +4144,9 @@ async fn execute_execution_runtime_stream_inner( record_stream_pending_lifecycle(state, seed, &mut stage_trace).await; lifecycle_pending_recorded = true; } + let stream_precommit_committed = execution.stream_precommit_committed; let frame_stream = build_direct_execution_frame_stream(execution).boxed(); - return execute_stream_from_frame_stream( + return execute_stream_from_frame_stream_with_retry_scope( state, plan, trace_id, @@ -3655,7 +4159,10 @@ async fn execute_execution_runtime_stream_inner( stage_trace, lifecycle_pending_recorded, frame_stream, + stream_precommit_committed, provider_pool_in_flight_guard.take(), + retry_scope_out.as_deref_mut(), + retry_fallback_out.as_deref_mut(), ) .await; } @@ -3730,7 +4237,7 @@ async fn execute_execution_runtime_stream_inner( .bytes_stream() .map_err(|err| IoError::other(err.to_string())) .boxed(); - return execute_stream_from_frame_stream( + return execute_stream_from_frame_stream_with_retry_scope( state, plan, trace_id, @@ -3743,7 +4250,10 @@ async fn execute_execution_runtime_stream_inner( stage_trace, lifecycle_pending_recorded, frame_stream, + false, provider_pool_in_flight_guard.take(), + retry_scope_out.as_deref_mut(), + retry_fallback_out.as_deref_mut(), ) .await; } @@ -3775,6 +4285,40 @@ fn parse_prefetched_sync_json_body(body: &[u8]) -> Option { serde_json::from_slice::(stripped).ok() } +fn resolve_provider_stream_error_status_code( + provider_api_format: &str, + upstream_status_code: u16, + body_json: &Value, +) -> u16 { + if (200..300).contains(&upstream_status_code) + && provider_api_format + .trim() + .eq_ignore_ascii_case("claude:messages") + { + anthropic_error_status_code(body_json) + } else { + resolve_local_sync_error_status_code(upstream_status_code, body_json) + } +} + +fn anthropic_premature_eof_error_body(message: &str) -> Value { + serde_json::json!({ + "type": "error", + "error": { + "type": "api_error", + "message": message, + } + }) +} + +fn build_anthropic_premature_eof_failure(message: &str) -> StreamFailureReport { + let body_json = anthropic_premature_eof_error_body(message); + build_stream_failure_from_provider_error_body( + anthropic_error_status_code(&body_json), + &body_json, + ) +} + fn encode_terminal_sse_error_event(failure: &StreamFailureReport) -> Result { let payload = failure .to_json_string() @@ -3789,6 +4333,39 @@ fn encode_terminal_sse_error_event(failure: &StreamFailureReport) -> Result Result { + let payload = serde_json::to_string(&serde_json::json!({ + "type": "error", + "error": { + "type": "api_error", + "message": failure.error_message, + } + })) + .map_err(|err| IoError::other(err.to_string()))?; + Ok(Bytes::from(format!("event: error\ndata: {payload}\n\n"))) +} + +fn encode_terminal_sse_error_event_for_plan( + plan: &ExecutionPlan, + failure: &StreamFailureReport, +) -> Result { + if plan + .client_api_format + .trim() + .eq_ignore_ascii_case("claude:messages") + && plan + .provider_api_format + .trim() + .eq_ignore_ascii_case("claude:messages") + { + encode_anthropic_terminal_sse_error_event(failure) + } else { + encode_terminal_sse_error_event(failure) + } +} + fn image_stream_failed_event_name(report_context: Option<&Value>) -> &'static str { let operation = report_context .and_then(|value| value.get("image_request")) @@ -3855,15 +4432,25 @@ fn build_sse_body_stream( mut rx: mpsc::Receiver>, filter_control_blocks: bool, emit_keepalive: bool, + anthropic_message_stop_terminates_body: bool, keepalive_interval: Duration, ) -> impl futures_util::Stream> + Send + 'static { stream! { let mut upstream_control_filter = filter_control_blocks.then(SseControlBlockFilter::default); + let mut anthropic_completion_tracker = + anthropic_message_stop_terminates_body.then(ClientVisibleStreamCompletionTracker::default); let mut sent_prefetched_chunk = false; for chunk in prefetched_chunks_for_body { - if let Some(chunk) = filter_upstream_sse_control_chunk(&mut upstream_control_filter, chunk) { + if let Some(mut chunk) = filter_upstream_sse_control_chunk(&mut upstream_control_filter, chunk) { + let completed = truncate_at_anthropic_message_stop( + anthropic_completion_tracker.as_mut(), + &mut chunk, + ); sent_prefetched_chunk = true; yield Ok(chunk); + if completed { + return; + } } } @@ -3883,8 +4470,15 @@ fn build_sse_body_stream( }; match item { Ok(chunk) => { - if let Some(chunk) = filter_upstream_sse_control_chunk(&mut upstream_control_filter, chunk) { + if let Some(mut chunk) = filter_upstream_sse_control_chunk(&mut upstream_control_filter, chunk) { + let completed = truncate_at_anthropic_message_stop( + anthropic_completion_tracker.as_mut(), + &mut chunk, + ); yield Ok(chunk); + if completed { + break; + } } } Err(err) => yield Err(err), @@ -3895,27 +4489,63 @@ fn build_sse_body_stream( } } } - if let Some(chunk) = flush_upstream_sse_control_filter(&mut upstream_control_filter) { - yield Ok(chunk); + if !anthropic_completion_tracker + .as_ref() + .is_some_and(|tracker| tracker.completed) + { + if let Some(chunk) = + flush_upstream_sse_control_filter(&mut upstream_control_filter) + { + yield Ok(chunk); + } } } else { while let Some(item) = rx.recv().await { match item { Ok(chunk) => { - if let Some(chunk) = filter_upstream_sse_control_chunk(&mut upstream_control_filter, chunk) { + if let Some(mut chunk) = filter_upstream_sse_control_chunk(&mut upstream_control_filter, chunk) { + let completed = truncate_at_anthropic_message_stop( + anthropic_completion_tracker.as_mut(), + &mut chunk, + ); yield Ok(chunk); + if completed { + break; + } } } Err(err) => yield Err(err), } } - if let Some(chunk) = flush_upstream_sse_control_filter(&mut upstream_control_filter) { - yield Ok(chunk); + if !anthropic_completion_tracker + .as_ref() + .is_some_and(|tracker| tracker.completed) + { + if let Some(chunk) = + flush_upstream_sse_control_filter(&mut upstream_control_filter) + { + yield Ok(chunk); + } } } } } +fn truncate_at_anthropic_message_stop( + tracker: Option<&mut ClientVisibleStreamCompletionTracker>, + chunk: &mut Bytes, +) -> bool { + let Some(tracker) = tracker else { + return false; + }; + if let Some(terminal_end) = tracker.observe_anthropic_message_stop_terminal_end(chunk.as_ref()) + { + chunk.truncate(terminal_end); + return true; + } + false +} + #[derive(Default)] struct SseControlBlockFilter { buffered: Vec, @@ -4008,21 +4638,7 @@ fn flush_upstream_sse_control_filter(filter: &mut Option) } fn find_sse_block_boundary(buffer: &[u8]) -> Option<(usize, usize)> { - let lf = buffer - .windows(2) - .position(|window| window == b"\n\n") - .map(|index| (index, 2)); - let crlf = buffer - .windows(4) - .position(|window| window == b"\r\n\r\n") - .map(|index| (index, 4)); - - match (lf, crlf) { - (Some(lf), Some(crlf)) => Some(if lf.0 <= crlf.0 { lf } else { crlf }), - (Some(lf), None) => Some(lf), - (None, Some(crlf)) => Some(crlf), - (None, None) => None, - } + find_sse_record_boundary(buffer) } fn sse_block_has_data_line(block: &[u8]) -> bool { @@ -4030,7 +4646,7 @@ fn sse_block_has_data_line(block: &[u8]) -> bool { return true; }; - text.lines() + text.split(['\r', '\n']) .any(|line| line.trim_start().starts_with("data:")) } @@ -4039,30 +4655,71 @@ fn sse_buffer_has_data_line(buffer: &[u8]) -> bool { return true; }; - text.lines() + text.split(['\r', '\n']) .any(|line| line.trim_start().starts_with("data:")) } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum SseTerminalPolicy { + AnyKnown, + AnthropicMessageStop, +} + #[derive(Default)] struct ClientVisibleStreamCompletionTracker { line_buffer: Vec, event_type: Option, data_payload: String, has_data_payload: bool, + record_bytes: usize, + dropping_oversized_record: bool, + discarded_line_nonempty: bool, skip_next_lf: bool, completed: bool, } impl ClientVisibleStreamCompletionTracker { fn observe_chunk(&mut self, chunk: &[u8]) -> bool { + self.observe_chunk_terminal_end(chunk); + self.completed + } + + fn observe_chunk_terminal_end(&mut self, chunk: &[u8]) -> Option { + self.observe_chunk_terminal_end_with_policy(chunk, SseTerminalPolicy::AnyKnown) + } + + fn observe_anthropic_message_stop(&mut self, chunk: &[u8]) -> bool { + self.observe_anthropic_message_stop_terminal_end(chunk); + self.completed + } + + fn observe_anthropic_message_stop_terminal_end(&mut self, chunk: &[u8]) -> Option { + self.observe_chunk_terminal_end_with_policy(chunk, SseTerminalPolicy::AnthropicMessageStop) + } + + fn observe_chunk_terminal_end_with_policy( + &mut self, + chunk: &[u8], + policy: SseTerminalPolicy, + ) -> Option { if self.completed { - return true; + return None; } if chunk.is_empty() { - return false; + return None; } - for byte in chunk { + for (index, byte) in chunk.iter().enumerate() { + self.record_bytes = self.record_bytes.saturating_add(1); + if self.record_bytes > SSE_TERMINAL_DETECTOR_MAX_RECORD_BYTES + && !self.dropping_oversized_record + { + self.dropping_oversized_record = true; + self.discarded_line_nonempty = !self.line_buffer.is_empty(); + self.line_buffer.clear(); + self.reset_current_event(); + } + if self.skip_next_lf { self.skip_next_lf = false; if *byte == b'\n' { @@ -4070,29 +4727,41 @@ impl ClientVisibleStreamCompletionTracker { } } + if self.dropping_oversized_record { + match *byte { + b'\n' => self.finish_discarded_line(), + b'\r' => { + self.finish_discarded_line(); + self.skip_next_lf = true; + } + _ => self.discarded_line_nonempty = true, + } + continue; + } + match *byte { - b'\n' => self.finish_line(), + b'\n' => self.finish_line(policy), b'\r' => { - self.finish_line(); + self.finish_line(policy); self.skip_next_lf = true; } - _ => { - self.line_buffer.push(*byte); - if self.line_buffer.len() > SSE_TERMINAL_DETECTOR_MAX_LINE_BYTES { - self.line_buffer.clear(); - } - } + _ => self.line_buffer.push(*byte), } if self.completed { - break; + let terminal_end = if *byte == b'\r' && chunk.get(index + 1) == Some(&b'\n') { + index + 2 + } else { + index + 1 + }; + return Some(terminal_end); } } - self.completed + None } - fn finish_line(&mut self) { + fn finish_line(&mut self, policy: SseTerminalPolicy) { let line = std::mem::take(&mut self.line_buffer); let Ok(line) = std::str::from_utf8(&line) else { self.reset_current_event(); @@ -4101,8 +4770,9 @@ impl ClientVisibleStreamCompletionTracker { let line = line.trim(); if line.is_empty() { - self.completed = self.current_event_is_terminal(); + self.completed = self.current_event_is_terminal(policy); self.reset_current_event(); + self.record_bytes = 0; return; } @@ -4123,11 +4793,41 @@ impl ClientVisibleStreamCompletionTracker { } } - fn current_event_is_terminal(&self) -> bool { - self.event_type - .as_deref() - .is_some_and(is_terminal_sse_event_type) - || (self.has_data_payload && sse_data_payload_is_terminal(&self.data_payload)) + fn finish_discarded_line(&mut self) { + if !self.discarded_line_nonempty { + self.dropping_oversized_record = false; + self.record_bytes = 0; + self.reset_current_event(); + } + self.discarded_line_nonempty = false; + } + + fn current_event_is_terminal(&self, policy: SseTerminalPolicy) -> bool { + match policy { + SseTerminalPolicy::AnyKnown => { + self.event_type + .as_deref() + .is_some_and(is_terminal_sse_event_type) + || (self.has_data_payload && sse_data_payload_is_terminal(&self.data_payload)) + } + SseTerminalPolicy::AnthropicMessageStop => { + let payload_type = self + .has_data_payload + .then(|| serde_json::from_str::(&self.data_payload).ok()) + .flatten() + .and_then(|value| { + value + .get("type") + .and_then(serde_json::Value::as_str) + .map(ToOwned::to_owned) + }); + payload_type.as_deref() == Some("message_stop") + && self + .event_type + .as_deref() + .is_none_or(|event_type| event_type == "message_stop") + } + } } fn reset_current_event(&mut self) { @@ -4164,6 +4864,150 @@ struct ObservedStreamFrame { observed_at: Instant, } +#[derive(Clone)] +struct PostStopFrameReadBudget { + remaining: Arc, +} + +impl PostStopFrameReadBudget { + fn new() -> Self { + Self { + remaining: Arc::new(AtomicUsize::new(POST_STOP_FRAME_READ_BUDGET_INACTIVE)), + } + } + + fn activate(&self, already_buffered: usize) -> bool { + let remaining = ANTHROPIC_POST_STOP_DRAIN_MAX_BYTES.saturating_sub(already_buffered); + let activated = self + .remaining + .compare_exchange( + POST_STOP_FRAME_READ_BUDGET_INACTIVE, + remaining, + Ordering::AcqRel, + Ordering::Acquire, + ) + .is_ok(); + activated && already_buffered > ANTHROPIC_POST_STOP_DRAIN_MAX_BYTES + } +} + +struct PostStopLimitedStreamReader { + stream: S, + current: Option, + budget: PostStopFrameReadBudget, +} + +impl PostStopLimitedStreamReader { + fn new(stream: S, budget: PostStopFrameReadBudget) -> Self { + Self { + stream, + current: None, + budget, + } + } + + fn activate_post_stop_budget(&mut self, already_buffered: usize) -> bool { + let over_limit = self.budget.activate(already_buffered); + let remaining = self.budget.remaining.load(Ordering::Acquire); + self.trim_current_to_budget(remaining, true); + over_limit + } + + fn trim_current_to_budget(&mut self, remaining: usize, detach_backing: bool) { + if remaining == POST_STOP_FRAME_READ_BUDGET_INACTIVE { + return; + } + if remaining == 0 { + self.current = None; + return; + } + if let Some(current) = self.current.as_mut() { + if detach_backing || current.len() > remaining { + let retained = current.len().min(remaining); + // Detach even a small slice because it can retain a giant + // producer allocation across post-stop backpressure. + *current = Bytes::copy_from_slice(¤t[..retained]); + } + } + } +} + +impl AsyncRead for PostStopLimitedStreamReader +where + S: Stream> + Unpin, +{ + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + let this = self.get_mut(); + if buf.remaining() == 0 { + return Poll::Ready(Ok(())); + } + + let mut empty_chunks = 0usize; + loop { + let remaining = this.budget.remaining.load(Ordering::Acquire); + if remaining == 0 { + this.current = None; + return Poll::Ready(Ok(())); + } + this.trim_current_to_budget(remaining, false); + + if let Some(current) = this.current.as_mut() { + let read = current.len().min(buf.remaining()); + if read > 0 { + buf.put_slice(¤t.split_to(read)); + if remaining != POST_STOP_FRAME_READ_BUDGET_INACTIVE { + let previous = this.budget.remaining.fetch_sub(read, Ordering::AcqRel); + debug_assert!(previous != POST_STOP_FRAME_READ_BUDGET_INACTIVE); + debug_assert!(previous >= read); + } + } + if current.is_empty() { + this.current = None; + } + if read > 0 { + return Poll::Ready(Ok(())); + } + } + + match Pin::new(&mut this.stream).poll_next(cx) { + Poll::Pending => return Poll::Pending, + Poll::Ready(Some(Ok(chunk))) if chunk.is_empty() => { + empty_chunks += 1; + if empty_chunks >= POST_STOP_MAX_EMPTY_CHUNKS_PER_POLL { + cx.waker().wake_by_ref(); + return Poll::Pending; + } + } + Poll::Ready(Some(Ok(chunk))) => { + this.current = Some(chunk); + if remaining != POST_STOP_FRAME_READ_BUDGET_INACTIVE { + this.trim_current_to_budget(remaining, true); + } + } + Poll::Ready(Some(Err(err))) => return Poll::Ready(Err(err)), + Poll::Ready(None) => return Poll::Ready(Ok(())), + } + } + } +} + +fn activate_post_stop_frame_read_budget( + lines: &mut FramedRead, LinesCodec>, +) -> bool { + let already_buffered = lines.read_buffer().len(); + let over_limit = already_buffered > ANTHROPIC_POST_STOP_DRAIN_MAX_BYTES; + let reader_over_limit = lines.get_mut().activate_post_stop_budget(already_buffered); + let retained = if over_limit { 0 } else { already_buffered }; + let mut bounded = bytes::BytesMut::with_capacity(retained); + bounded.extend_from_slice(&lines.read_buffer()[..retained]); + *lines.read_buffer_mut() = bounded; + reader_over_limit || over_limit +} + async fn read_next_observed_stream_frame( lines: &mut FramedRead, ) -> Result, GatewayError> @@ -4191,6 +5035,10 @@ where read_next_observed_stream_frame(lines).await } +fn serialized_stream_frame_len(frame: &StreamFrame) -> usize { + serde_json::to_vec(frame).map_or(usize::MAX, |encoded| encoded.len()) +} + fn should_refresh_stream_usage_telemetry( previous: Option<&ExecutionTelemetry>, next: &ExecutionTelemetry, @@ -4320,36 +5168,16 @@ fn should_skip_direct_finalize_prefetch( has_local_stream_rewriter: bool, force_prefetch: bool, ) -> bool { - if direct_stream_finalize_kind.is_none() { - return false; - } - - if force_prefetch { - return false; - } - - let content_type = content_type - .map(str::trim) - .filter(|value| !value.is_empty()) - .unwrap_or_default() - .to_ascii_lowercase(); - if content_type.contains("text/event-stream") { - return true; - } - - if has_private_stream_normalizer || has_local_stream_rewriter { - return false; - } - - if !provider_api_format.eq_ignore_ascii_case(client_api_format) { - return false; - } - - if content_type.is_empty() { - return true; - } - - !(content_type.contains("json") || content_type.ends_with("+json")) + StreamCommitPolicy::for_response( + direct_stream_finalize_kind.is_some(), + content_type, + provider_api_format, + client_api_format, + has_private_stream_normalizer, + has_local_stream_rewriter, + force_prefetch, + ) + .commits_on_response_headers() } fn prefetched_openai_responses_body_has_output_boundary(body: &[u8]) -> bool { @@ -4433,10 +5261,50 @@ async fn execute_stream_from_frame_stream( report_context: Option, candidate_started_unix_secs: u64, stream_started_at: Instant, - mut stage_trace: RequestStageTrace, + stage_trace: RequestStageTrace, lifecycle_pending_recorded: bool, frame_stream: BoxStream<'static, Result>, in_flight_guard: Option, +) -> Result>, GatewayError> { + execute_stream_from_frame_stream_with_retry_scope( + state, + plan, + trace_id, + decision, + plan_kind, + report_kind, + report_context, + candidate_started_unix_secs, + stream_started_at, + stage_trace, + lifecycle_pending_recorded, + frame_stream, + false, + in_flight_guard, + None, + None, + ) + .await +} + +#[allow(clippy::too_many_arguments)] +async fn execute_stream_from_frame_stream_with_retry_scope( + state: &AppState, + plan: ExecutionPlan, + trace_id: &str, + decision: &GatewayControlDecision, + plan_kind: &str, + report_kind: Option, + report_context: Option, + candidate_started_unix_secs: u64, + stream_started_at: Instant, + mut stage_trace: RequestStageTrace, + lifecycle_pending_recorded: bool, + frame_stream: BoxStream<'static, Result>, + stream_precommit_committed: bool, + in_flight_guard: Option, + mut retry_scope_out: Option<&mut AiAttemptRetryScope>, + mut retry_fallback_out: Option<&mut Option>>, ) -> Result>, GatewayError> { let request_id = plan.request_id.as_str(); let request_id_for_log = short_request_id(request_id); @@ -4453,7 +5321,7 @@ async fn execute_stream_from_frame_stream( .and_then(|context| context.candidate_index) .map(|value| value.to_string()) .unwrap_or_else(|| "-".to_string()); - let reader = StreamReader::new(frame_stream); + let reader = PostStopLimitedStreamReader::new(frame_stream, PostStopFrameReadBudget::new()); let mut lines = FramedRead::new(reader, LinesCodec::new()); let first_frame_started_at = Instant::now(); @@ -4671,6 +5539,37 @@ async fn execute_stream_from_frame_stream( "gateway resolved execution runtime stream failover decision" ); if matches!(failover_decision, LocalFailoverDecision::RetryNextCandidate) { + let failure_disposition = classify_failure_disposition( + &plan.provider_api_format, + failover_analysis.classification, + status_code, + ); + if let Some(retry_scope) = retry_scope_out.as_deref_mut() { + *retry_scope = ai_attempt_retry_scope_from_failure_disposition(failure_disposition); + } + if failure_disposition.preserve_upstream_error { + if let Some(retry_fallback) = retry_fallback_out.as_deref_mut() { + let mut fallback_headers = headers.clone(); + apply_endpoint_response_header_rules( + state, + &plan, + &mut fallback_headers, + provider_body_json.as_ref(), + ) + .await?; + *retry_fallback = Some(attach_control_metadata_headers( + build_client_response_from_parts( + status_code, + &fallback_headers, + Body::from(provider_error_body.clone()), + trace_id, + Some(decision), + )?, + Some(request_id), + candidate_id, + )?); + } + } let terminal_unix_secs = current_request_candidate_unix_ms(); let error_trace_report_context = with_stream_error_trace_context( report_context.as_ref(), @@ -4869,8 +5768,8 @@ async fn execute_stream_from_frame_stream( let prefetch_for_cyber_failover = is_openai_responses_family_format(plan.provider_api_format.as_str()) && cyber_continue_failover_enabled(state).await; - let skip_direct_finalize_prefetch = should_skip_direct_finalize_prefetch( - direct_stream_finalize_kind.as_deref(), + let stream_commit_policy = StreamCommitPolicy::for_response( + direct_stream_finalize_kind.is_some(), upstream_content_type, plan.provider_api_format.as_str(), plan.client_api_format.as_str(), @@ -4878,16 +5777,32 @@ async fn execute_stream_from_frame_stream( local_stream_rewriter.is_some(), prefetch_for_cyber_failover, ); + let reuse_committed_precommit = + stream_precommit_committed && stream_commit_policy.is_native_anthropic(); + let skip_direct_finalize_prefetch = + stream_commit_policy.commits_on_response_headers() || reuse_committed_precommit; let limit_direct_finalize_prefetch = - should_limit_direct_finalize_prefetch(plan_kind, local_stream_rewriter.is_some()); + should_limit_direct_finalize_prefetch(plan_kind, local_stream_rewriter.is_some()) + || stream_commit_policy.requires_bounded_frame_wait(); + let mut stream_commit_gate = StreamCommitGate::new(stream_commit_policy); + let mut prefetch_client_completion_tracker = ClientVisibleStreamCompletionTracker::default(); + let mut prefetched_client_visible_stream_completed = false; + let mut prefetched_anthropic_message_stop_observed_at = None; + let mut prefetched_anthropic_post_stop_buffer_over_limit = false; + if reuse_committed_precommit { + stream_commit_gate.commit(); + } let mut prefetched_chunks: Vec = Vec::new(); let mut provider_prefetched_body = Vec::new(); + let mut provider_prefetched_body_truncated = false; let mut prefetched_body = Vec::new(); let mut prefetched_inspection_body = Vec::new(); + let mut prefetched_inspection_body_truncated = false; let mut prefetched_telemetry: Option = None; let mut prefetched_usage_telemetry: Option = None; let mut reached_eof = false; let mut sync_json_stream_bridge_active = false; + let precommit_started_at = Instant::now(); if skip_direct_finalize_prefetch { debug!( event_name = "execution_runtime_stream_prefetch_skipped", @@ -4911,18 +5826,47 @@ async fn execute_stream_from_frame_stream( .as_ref() .filter(|_| !skip_direct_finalize_prefetch) { - while prefetched_chunks.len() < MAX_STREAM_PREFETCH_FRAMES + while (stream_commit_policy.requires_bounded_frame_wait() + || prefetched_chunks.len() < MAX_STREAM_PREFETCH_FRAMES) && prefetched_inspection_body.len() < MAX_STREAM_PREFETCH_BYTES { let next_frame_result = if limit_direct_finalize_prefetch { + let prefetch_timeout = stream_commit_policy + .max_precommit_wait() + .map(|max_wait| max_wait.saturating_sub(precommit_started_at.elapsed())) + .unwrap_or(REWRITTEN_STREAM_PREFETCH_TIMEOUT); + if prefetch_timeout.is_zero() { + stream_commit_gate.commit(); + debug!( + event_name = "execution_runtime_stream_prefetch_limited", + log_type = "debug", + trace_id = %trace_id, + request_id = %request_id_for_log, + candidate_id = ?candidate_id, + plan_kind, + report_kind, + provider_name, + endpoint_id = %plan.endpoint_id, + key_id = %plan.key_id, + model_name, + candidate_index = candidate_index.as_str(), + timeout_ms = stream_commit_policy + .max_precommit_wait() + .unwrap_or(REWRITTEN_STREAM_PREFETCH_TIMEOUT) + .as_millis() as u64, + "gateway reached bounded stream precommit deadline" + ); + break; + } match tokio::time::timeout( - REWRITTEN_STREAM_PREFETCH_TIMEOUT, + prefetch_timeout, next_stream_frame(&mut buffered_frames, &mut lines), ) .await { Ok(result) => result, Err(_) => { + stream_commit_gate.commit(); debug!( event_name = "execution_runtime_stream_prefetch_limited", log_type = "debug", @@ -4936,8 +5880,11 @@ async fn execute_stream_from_frame_stream( key_id = %plan.key_id, model_name, candidate_index = candidate_index.as_str(), - timeout_ms = REWRITTEN_STREAM_PREFETCH_TIMEOUT.as_millis() as u64, - "gateway stopped rewritten stream prefetch before client-visible body" + timeout_ms = stream_commit_policy + .max_precommit_wait() + .unwrap_or(REWRITTEN_STREAM_PREFETCH_TIMEOUT) + .as_millis() as u64, + "gateway stopped bounded stream prefetch before client-visible body" ); break; } @@ -4970,6 +5917,32 @@ async fn execute_stream_from_frame_stream( .await; } }) else { + if stream_commit_policy.is_native_anthropic() && stream_commit_gate.is_uncommitted() + { + let error_body_json = anthropic_premature_eof_error_body( + "upstream Anthropic stream ended before the first semantic event", + ); + let error_status_code = anthropic_error_status_code(&error_body_json); + return handle_prefetch_provider_private_stream_error( + state, + trace_id, + decision, + &plan, + report_context, + request_id, + candidate_id, + report_kind, + headers, + prefetched_usage_telemetry.clone(), + &provider_prefetched_body, + status_code, + error_status_code, + error_body_json, + retry_scope_out.as_deref_mut(), + None, + ) + .await; + } reached_eof = true; break; }; @@ -4988,7 +5961,7 @@ async fn execute_stream_from_frame_stream( stream_elapsed_ms_at(stream_started_at, frame_observed_at), ); } - let chunk = + let mut chunk = match decode_stream_data_chunk(chunk_b64.as_deref(), text.as_deref()) { Ok(chunk) => chunk, Err(err) => { @@ -5020,38 +5993,105 @@ async fn execute_stream_from_frame_stream( if chunk.is_empty() { continue; } - - provider_prefetched_body.extend_from_slice(&chunk); - prefetched_inspection_body.extend_from_slice(&chunk); - - if let Some(error_body_json) = extract_provider_private_stream_error_body( - report_context.as_ref(), - &prefetched_inspection_body, - ) { - let error_status_code = - resolve_local_sync_error_status_code(status_code, &error_body_json); - return handle_prefetch_provider_private_stream_error( - state, - trace_id, - decision, - &plan, - report_context, - request_id, - candidate_id, - report_kind, - headers, - prefetched_usage_telemetry.clone(), - &provider_prefetched_body, - error_status_code, - error_body_json, - ) - .await; + if stream_commit_policy.is_native_anthropic() { + if prefetched_client_visible_stream_completed { + continue; + } + if let Some(terminal_end) = prefetch_client_completion_tracker + .observe_anthropic_message_stop_terminal_end(&chunk) + { + chunk.truncate(terminal_end); + prefetched_client_visible_stream_completed = true; + prefetched_anthropic_message_stop_observed_at + .get_or_insert_with(Instant::now); + prefetched_anthropic_post_stop_buffer_over_limit |= + activate_post_stop_frame_read_budget(&mut lines); + } } - let inspection = inspect_prefetched_stream_body( - &upstream_headers, - &prefetched_inspection_body, + append_stream_capture_bytes( + &mut provider_prefetched_body, + &chunk, + MAX_STREAM_PREFETCH_BYTES, + &mut provider_prefetched_body_truncated, ); + append_stream_capture_bytes( + &mut prefetched_inspection_body, + &chunk, + MAX_STREAM_PREFETCH_BYTES, + &mut prefetched_inspection_body_truncated, + ); + + let anthropic_commit_ready = + match stream_commit_gate.observe_provider_bytes(&chunk) { + StreamPrecommitObservation::Pending => false, + StreamPrecommitObservation::Commit => true, + StreamPrecommitObservation::UpstreamError { + status_code: error_status_code, + body_json: error_body_json, + } => { + return handle_prefetch_provider_private_stream_error( + state, + trace_id, + decision, + &plan, + report_context, + request_id, + candidate_id, + report_kind, + headers, + prefetched_usage_telemetry.clone(), + &provider_prefetched_body, + status_code, + error_status_code, + error_body_json, + retry_scope_out.as_deref_mut(), + retry_fallback_out.as_deref_mut(), + ) + .await; + } + }; + + if !anthropic_commit_ready { + if let Some(error_body_json) = extract_provider_private_stream_error_body( + report_context.as_ref(), + &prefetched_inspection_body, + ) { + let error_status_code = resolve_provider_stream_error_status_code( + plan.provider_api_format.as_str(), + status_code, + &error_body_json, + ); + return handle_prefetch_provider_private_stream_error( + state, + trace_id, + decision, + &plan, + report_context, + request_id, + candidate_id, + report_kind, + headers, + prefetched_usage_telemetry.clone(), + &provider_prefetched_body, + status_code, + error_status_code, + error_body_json, + retry_scope_out.as_deref_mut(), + retry_fallback_out.as_deref_mut(), + ) + .await; + } + } + + let inspection = if stream_commit_policy.is_native_anthropic() { + StreamPrefetchInspection::NeedMore + } else { + inspect_prefetched_stream_body( + &upstream_headers, + &prefetched_inspection_body, + ) + }; match inspection { StreamPrefetchInspection::EmbeddedError(body_json) => { debug!( @@ -5224,11 +6264,12 @@ async fn execute_stream_from_frame_stream( prefetched_chunks.push(Bytes::from(rewritten_chunk)); } - if matches!(inspection, StreamPrefetchInspection::NonError) - && (!prefetch_for_cyber_failover - || prefetched_openai_responses_body_has_output_boundary( - &prefetched_inspection_body, - )) + if anthropic_commit_ready + || (matches!(inspection, StreamPrefetchInspection::NonError) + && (!prefetch_for_cyber_failover + || prefetched_openai_responses_body_has_output_boundary( + &prefetched_inspection_body, + ))) { break; } @@ -5239,6 +6280,33 @@ async fn execute_stream_from_frame_stream( prefetched_telemetry = Some(frame_telemetry); } StreamFramePayload::Eof { summary } => { + if stream_commit_policy.is_native_anthropic() + && stream_commit_gate.is_uncommitted() + { + let error_body_json = anthropic_premature_eof_error_body( + "upstream Anthropic stream ended before the first semantic event", + ); + let error_status_code = anthropic_error_status_code(&error_body_json); + return handle_prefetch_provider_private_stream_error( + state, + trace_id, + decision, + &plan, + report_context, + request_id, + candidate_id, + report_kind, + headers, + prefetched_usage_telemetry.clone(), + &provider_prefetched_body, + status_code, + error_status_code, + error_body_json, + retry_scope_out.as_deref_mut(), + None, + ) + .await; + } if summary.is_some() { stream_terminal_summary = summary; } @@ -5275,6 +6343,9 @@ async fn execute_stream_from_frame_stream( } } } + if stream_commit_gate.is_uncommitted() { + stream_commit_gate.commit(); + } drop(private_stream_normalizer); drop(local_stream_rewriter); @@ -5337,8 +6408,10 @@ async fn execute_stream_from_frame_stream( let response_headers_are_sse = response_headers_indicate_sse(&headers); let emit_proxy_generated_sse_control_blocks = response_headers_are_sse && client_format_allows_proxy_generated_sse_control_blocks(&plan); + let native_anthropic_stream_for_report = stream_commit_policy.is_native_anthropic(); let plan_for_report = plan; - let emit_passthrough_sse_terminal_error = skip_direct_finalize_prefetch + let emit_passthrough_sse_terminal_error = (skip_direct_finalize_prefetch + || stream_commit_policy.is_native_anthropic()) && response_headers_indicate_sse(&upstream_headers) && !is_openai_image_stream_for_report; let plan_kind_for_report = plan_kind.to_string(); @@ -5397,13 +6470,31 @@ async fn execute_stream_from_frame_stream( &mut client_body_truncated, ); let mut client_stream_completion_tracker = ClientVisibleStreamCompletionTracker::default(); - let mut client_visible_stream_completed = - client_stream_completion_tracker.observe_chunk(&prefetched_body_for_report); + let mut client_visible_stream_completed = if native_anthropic_stream_for_report { + client_stream_completion_tracker + .observe_anthropic_message_stop(&prefetched_body_for_report) + } else { + client_stream_completion_tracker.observe_chunk(&prefetched_body_for_report) + }; + let mut anthropic_post_stop_drain_started_at = (native_anthropic_stream_for_report + && client_visible_stream_completed) + .then(|| prefetched_anthropic_message_stop_observed_at.unwrap_or_else(Instant::now)); + let mut anthropic_post_stop_buffer_over_limit = + prefetched_anthropic_post_stop_buffer_over_limit; + if anthropic_post_stop_drain_started_at.is_some() + && prefetched_anthropic_message_stop_observed_at.is_none() + { + anthropic_post_stop_buffer_over_limit |= + activate_post_stop_frame_read_budget(&mut lines); + } + let mut anthropic_post_stop_drain_frames = 0usize; + let mut anthropic_post_stop_drain_bytes = 0usize; let mut usage_stream_telemetry: Option = initial_usage_telemetry; let mut telemetry: Option = initial_telemetry; let reached_eof = initial_reached_eof; let mut downstream_dropped = false; let mut terminal_failure: Option = None; + let mut provider_error_forwarded_to_client = false; let initial_elapsed_ms = stream_started_at_for_report .elapsed() .as_millis() @@ -5548,8 +6639,12 @@ async fn execute_stream_from_frame_stream( if let Some(error_body_json) = provider_error_inspection .observe(stream_usage_report_context.as_ref(), replay_chunk) { - let error_status_code = - resolve_local_sync_error_status_code(status_code, &error_body_json); + provider_error_forwarded_to_client = !prefetched_body_for_report.is_empty(); + let error_status_code = resolve_provider_stream_error_status_code( + plan_for_report.provider_api_format.as_str(), + status_code, + &error_body_json, + ); terminal_failure = Some(build_stream_failure_from_provider_error_body( error_status_code, &error_body_json, @@ -5592,17 +6687,55 @@ async fn execute_stream_from_frame_stream( if terminal_failure.is_none() && !reached_eof { loop { - let next_frame_result = tokio::select! { - biased; - _ = tx.closed(), if !downstream_dropped => { - downstream_dropped = true; + let draining_after_anthropic_stop = anthropic_post_stop_drain_started_at.is_some(); + let next_frame_result = if let Some(drain_started_at) = + anthropic_post_stop_drain_started_at + { + if anthropic_post_stop_drain_frames >= ANTHROPIC_POST_STOP_DRAIN_MAX_FRAMES + || anthropic_post_stop_drain_bytes >= ANTHROPIC_POST_STOP_DRAIN_MAX_BYTES + || anthropic_post_stop_buffer_over_limit + { break; } - result = next_stream_frame(&mut buffered_frames, &mut lines) => result, + let remaining = ANTHROPIC_POST_STOP_DRAIN_MAX_WAIT + .saturating_sub(drain_started_at.elapsed()); + if remaining.is_zero() { + break; + } + match tokio::time::timeout( + remaining, + next_stream_frame(&mut buffered_frames, &mut lines), + ) + .await + { + Ok(result) => result, + Err(_) => break, + } + } else { + tokio::select! { + biased; + _ = tx.closed(), if !downstream_dropped => { + downstream_dropped = true; + break; + } + result = next_stream_frame(&mut buffered_frames, &mut lines) => result, + } }; let next_frame = match next_frame_result { Ok(frame) => frame, Err(err) => { + if native_anthropic_stream_for_report && client_visible_stream_completed { + debug!( + event_name = "stream_execution_frame_decode_ignored_after_anthropic_stop", + log_type = "debug", + trace_id = %trace_id_owned, + request_id = %request_id_for_report_log, + candidate_id = ?candidate_id_for_report.as_deref(), + error = ?err, + "gateway ignored execution runtime teardown error after Anthropic message_stop" + ); + break; + } warn!( event_name = "stream_execution_frame_decode_failed", log_type = "ops", @@ -5623,15 +6756,36 @@ async fn execute_stream_from_frame_stream( let Some(observed_frame) = next_frame else { if tx.is_closed() { downstream_dropped = true; + } else if native_anthropic_stream_for_report && !client_visible_stream_completed + { + terminal_failure = Some(build_anthropic_premature_eof_failure( + "upstream Anthropic stream ended before message_stop", + )); } break; }; + if draining_after_anthropic_stop { + anthropic_post_stop_drain_frames = + anthropic_post_stop_drain_frames.saturating_add(1); + let frame_bytes = serialized_stream_frame_len(&observed_frame.frame); + if frame_bytes + > ANTHROPIC_POST_STOP_DRAIN_MAX_BYTES + .saturating_sub(anthropic_post_stop_drain_bytes) + { + break; + } + anthropic_post_stop_drain_bytes = + anthropic_post_stop_drain_bytes.saturating_add(frame_bytes); + } let frame_observed_at = observed_frame.observed_at; let frame_elapsed_ms = stream_elapsed_ms_at(stream_started_at_for_report, frame_observed_at); last_upstream_frame_elapsed_ms.store(frame_elapsed_ms, Ordering::Relaxed); match observed_frame.frame.payload { StreamFramePayload::Data { chunk_b64, text } => { + if native_anthropic_stream_for_report && client_visible_stream_completed { + continue; + } let first_data_before = usage_stream_telemetry .as_ref() .and_then(|telemetry| telemetry.ttfb_ms) @@ -5666,7 +6820,7 @@ async fn execute_stream_from_frame_stream( if sync_json_stream_bridge_active_for_report { continue; } - let chunk = + let mut chunk = match decode_stream_data_chunk(chunk_b64.as_deref(), text.as_deref()) { Ok(chunk) => chunk, Err(err) => { @@ -5693,6 +6847,19 @@ async fn execute_stream_from_frame_stream( if chunk.is_empty() { continue; } + if native_anthropic_stream_for_report { + if let Some(terminal_end) = client_stream_completion_tracker + .observe_anthropic_message_stop_terminal_end(&chunk) + { + chunk.truncate(terminal_end); + client_visible_stream_completed = true; + if anthropic_post_stop_drain_started_at.is_none() { + anthropic_post_stop_drain_started_at = Some(Instant::now()); + anthropic_post_stop_buffer_over_limit |= + activate_post_stop_frame_read_budget(&mut lines); + } + } + } provider_stream_bytes.fetch_add( u64::try_from(chunk.len()).unwrap_or(u64::MAX), @@ -5771,7 +6938,8 @@ async fn execute_stream_from_frame_stream( if rewritten_chunk.is_empty() { if let Some(error_body_json) = provider_private_error_body_json { - let error_status_code = resolve_local_sync_error_status_code( + let error_status_code = resolve_provider_stream_error_status_code( + plan_for_report.provider_api_format.as_str(), status_code, &error_body_json, ); @@ -5809,8 +6977,10 @@ async fn execute_stream_from_frame_stream( downstream_dropped = true; break; } else { - client_visible_stream_completed |= client_stream_completion_tracker - .observe_chunk(rewritten_chunk.as_ref()); + if !native_anthropic_stream_for_report { + client_visible_stream_completed |= client_stream_completion_tracker + .observe_chunk(rewritten_chunk.as_ref()); + } client_stream_bytes.fetch_add(rewritten_chunk_len, Ordering::Relaxed); last_client_chunk_elapsed_ms.store( stream_started_at_for_report @@ -5820,10 +6990,15 @@ async fn execute_stream_from_frame_stream( as u64, Ordering::Relaxed, ); + provider_error_forwarded_to_client = + provider_private_error_body_json.is_some(); } if let Some(error_body_json) = provider_private_error_body_json { - let error_status_code = - resolve_local_sync_error_status_code(status_code, &error_body_json); + let error_status_code = resolve_provider_stream_error_status_code( + plan_for_report.provider_api_format.as_str(), + status_code, + &error_body_json, + ); terminal_failure = Some(build_stream_failure_from_provider_error_body( error_status_code, &error_body_json, @@ -5853,9 +7028,26 @@ async fn execute_stream_from_frame_stream( StreamFramePayload::Eof { summary } => { stream_terminal_summary = merge_stream_terminal_summary(stream_terminal_summary.take(), summary); + if native_anthropic_stream_for_report && !client_visible_stream_completed { + terminal_failure = Some(build_anthropic_premature_eof_failure( + "upstream Anthropic stream ended before message_stop", + )); + } break; } StreamFramePayload::Error { error } => { + if native_anthropic_stream_for_report && client_visible_stream_completed { + debug!( + event_name = "stream_execution_error_frame_ignored_after_anthropic_stop", + log_type = "debug", + trace_id = %trace_id_owned, + request_id = %request_id_for_report_log, + candidate_id = ?candidate_id_for_report.as_deref(), + error = %error.message, + "gateway ignored execution runtime error frame after Anthropic message_stop" + ); + continue; + } warn!( event_name = "stream_execution_error_frame", log_type = "ops", @@ -5872,6 +7064,9 @@ async fn execute_stream_from_frame_stream( } } } + drop(lines); + drop(buffered_frames); + drop(_provider_pool_in_flight_guard); if downstream_dropped { debug!( @@ -5968,8 +7163,11 @@ async fn execute_stream_from_frame_stream( } } if let Some(error_body_json) = provider_private_error_body_json { - let error_status_code = - resolve_local_sync_error_status_code(status_code, &error_body_json); + let error_status_code = resolve_provider_stream_error_status_code( + plan_for_report.provider_api_format.as_str(), + status_code, + &error_body_json, + ); terminal_failure.get_or_insert_with(|| { build_stream_failure_from_provider_error_body( error_status_code, @@ -6067,8 +7265,12 @@ async fn execute_stream_from_frame_stream( report_context_owned.as_ref(), failure, )) - } else if emit_passthrough_sse_terminal_error { - Some(encode_terminal_sse_error_event(failure)) + } else if emit_passthrough_sse_terminal_error && !provider_error_forwarded_to_client + { + Some(encode_terminal_sse_error_event_for_plan( + &plan_for_report, + failure, + )) } else { None }; @@ -6431,6 +7633,7 @@ async fn execute_stream_from_frame_stream( rx, response_headers_are_sse, emit_proxy_generated_sse_control_blocks, + native_anthropic_stream_for_report, SSE_KEEPALIVE_INTERVAL, ); @@ -6462,6 +7665,7 @@ mod tests { }; use std::time::{Duration, Instant}; + use aether_ai_serving::{AiAttemptExecutionOutcome, AiAttemptRetryScope}; use aether_contracts::{ ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionPlan, ExecutionStreamTerminalSummary, ExecutionTelemetry, ExecutionTimeouts, RequestBody, @@ -6504,24 +7708,28 @@ mod tests { use tokio::sync::{mpsc, watch, Notify}; use super::{ - build_sse_body_stream, build_stream_sync_payload, + activate_post_stop_frame_read_budget, build_direct_execution_frame_stream, + build_sse_body_stream, build_stream_failure_report, build_stream_sync_payload, client_format_allows_proxy_generated_sse_control_blocks, - direct_upstream_response_byte_stream, + direct_upstream_response_byte_stream, encode_terminal_sse_error_event_for_plan, ensure_stream_terminal_summary_for_missing_observed_finish, execute_execution_runtime_stream, execute_in_process_stream_with_oauth_retry, - execute_stream_from_frame_stream, maybe_apply_kiro_prompt_cache_usage_to_stream_summary, - merge_stream_terminal_summary, parse_direct_passthrough_mode, - prefetch_direct_stream_error_body, prefetched_openai_responses_body_has_output_boundary, + execute_stream_from_frame_stream, execute_stream_from_frame_stream_with_retry_scope, + maybe_apply_kiro_prompt_cache_usage_to_stream_summary, merge_stream_terminal_summary, + parse_direct_passthrough_mode, prefetch_direct_stream_error_body, + prefetched_openai_responses_body_has_output_boundary, record_sync_terminal_usage_with_handoff, - record_sync_terminal_usage_with_handoff_after_spawn, should_limit_direct_finalize_prefetch, - should_probe_success_failover_before_stream, should_skip_direct_finalize_prefetch, - stream_chunk_contains_sse_done, stream_requires_observed_terminal_event, - stream_terminal_summary_missing_observed_finish, + record_sync_terminal_usage_with_handoff_after_spawn, + resolve_provider_stream_error_status_code, select_direct_anthropic_prefetch_wait, + should_limit_direct_finalize_prefetch, should_probe_success_failover_before_stream, + should_skip_direct_finalize_prefetch, stream_chunk_contains_sse_done, + stream_requires_observed_terminal_event, stream_terminal_summary_missing_observed_finish, stream_terminal_summary_missing_observed_finish_with_requirement, stream_terminal_summary_represents_failure_with_requirement, ClientVisibleStreamCompletionTracker, DirectPassthroughFinalizer, DirectPassthroughFinalizerCore, DirectPassthroughInlineBodyState, DirectPassthroughMode, - ProviderStreamErrorInspection, + PostStopFrameReadBudget, PostStopLimitedStreamReader, ProviderStreamErrorInspection, + ANTHROPIC_POST_STOP_DRAIN_MAX_BYTES, POST_STOP_MAX_EMPTY_CHUNKS_PER_POLL, }; use crate::control::GatewayControlDecision; use crate::stage_metrics::RequestStageTrace; @@ -6716,6 +7924,15 @@ mod tests { }) } + fn generic_oauth_test_auth_config(provider_type: &str) -> Value { + json!({ + "provider_type": provider_type, + "access_token": "stale-access-token", + "refresh_token": "refresh-token", + "expires_at": 4_102_444_800_u64 + }) + } + async fn collect_direct_execution_body( mut execution: crate::execution_runtime::DirectUpstreamStreamExecution, ) -> Result, String> { @@ -6841,6 +8058,207 @@ mod tests { .expect("execution should succeed") } + fn native_anthropic_stream_plan(request_id: &str) -> ExecutionPlan { + ExecutionPlan { + request_id: request_id.to_string(), + candidate_id: Some(format!("candidate-{request_id}")), + provider_name: Some("custom".to_string()), + provider_id: format!("provider-{request_id}"), + endpoint_id: format!("endpoint-{request_id}"), + key_id: format!("key-{request_id}"), + method: "POST".to_string(), + url: "https://api.anthropic.com/v1/messages".to_string(), + headers: BTreeMap::from([ + ("content-type".to_string(), "application/json".to_string()), + ("accept".to_string(), "text/event-stream".to_string()), + ]), + content_type: Some("application/json".to_string()), + content_encoding: None, + body: RequestBody::from_json(json!({ + "model": "claude-sonnet-4-6", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 32, + "stream": true + })), + stream: true, + client_api_format: "claude:messages".to_string(), + provider_api_format: "claude:messages".to_string(), + model_name: Some("claude-sonnet-4-6".to_string()), + proxy: None, + transport_profile: None, + timeouts: None, + } + } + + struct StreamDropFlag(Arc); + + impl Drop for StreamDropFlag { + fn drop(&mut self) { + self.0.store(true, Ordering::SeqCst); + } + } + + fn direct_anthropic_test_finalizer(request_id: &str) -> DirectPassthroughFinalizer { + let state = AppState::new().expect("app state should build"); + let plan = native_anthropic_stream_plan(request_id); + let lifecycle_seed = aether_usage_runtime::build_lifecycle_usage_seed(&plan, None); + DirectPassthroughFinalizer::new(DirectPassthroughFinalizerCore { + state, + trace_id: format!("trace-{request_id}"), + report_kind: None, + report_context: None, + lifecycle_seed, + direct_stream_finalize_kind: None, + stream_started_at: Instant::now(), + stage_trace: RequestStageTrace::from_env(), + request_diagnostics: None, + request_id_for_log: request_id.to_string(), + candidate_id: plan.candidate_id.clone(), + request_candidate_status_snapshot: None, + deferred_request_candidate_status_record: None, + candidate_started_unix_secs: crate::clock::current_unix_ms(), + status_code: 200, + headers: BTreeMap::from([( + "content-type".to_string(), + "text/event-stream".to_string(), + )]), + stream_usage_report_context: None, + stream_usage_observer: None, + stream_usage_observer_buffered: Vec::new(), + provider_error_inspection: ProviderStreamErrorInspection::default(), + provider_buffered_body: Vec::new(), + buffered_body: Vec::new(), + provider_body_truncated: false, + client_body_truncated: false, + client_stream_completion_tracker: ClientVisibleStreamCompletionTracker::default(), + requires_anthropic_message_stop: true, + client_visible_stream_completed: false, + usage_stream_telemetry: None, + telemetry: None, + provider_stream_bytes: 0, + client_stream_bytes: 0, + last_client_chunk_elapsed_ms: 0, + pending_recorded: false, + stream_started_recorded: false, + terminal_failure: None, + _provider_pool_in_flight_guard: None, + _upstream_target_permit: None, + plan, + }) + } + + fn discard_direct_test_finalizer(state: &mut DirectPassthroughInlineBodyState) { + if let Some(mut finalizer) = state.finalizer.take() { + finalizer.core.take(); + } + } + + fn direct_anthropic_inline_state( + request_id: &str, + items: Vec>, + ) -> DirectPassthroughInlineBodyState { + DirectPassthroughInlineBodyState { + finalizer: Some(direct_anthropic_test_finalizer(request_id)), + upstream: Some(futures_util::stream::iter(items).boxed()), + upstream_control_filter: Some(super::SseControlBlockFilter::default()), + upstream_started_at: Instant::now(), + stream_first_byte_timeout: None, + observed_first_body_poll: false, + observed_first_client_yield: false, + upstream_done: false, + control_filter_flushed: false, + terminal_error_sent: false, + finalized: false, + } + } + + async fn execute_native_anthropic_prefetch_stream( + request_id: &str, + chunks: Vec, + ) -> AiAttemptExecutionOutcome> { + execute_native_anthropic_prefetch_stream_with_terminal_error(request_id, chunks, None).await + } + + async fn execute_native_anthropic_prefetch_stream_with_terminal_error( + request_id: &str, + chunks: Vec, + terminal_error: Option, + ) -> AiAttemptExecutionOutcome> { + let plan = native_anthropic_stream_plan(request_id); + let provider_catalog = provider_catalog_for_plan(&plan, None); + let data_state = crate::data::GatewayDataState::with_provider_transport_reader_for_tests( + Arc::new(provider_catalog), + "development-key", + ); + let state = AppState::new() + .expect("app state should build") + .with_data_state_for_tests(data_state); + let frame_stream = stream! { + yield Ok::(ndjson_frame(StreamFrame { + frame_type: StreamFrameType::Headers, + payload: StreamFramePayload::Headers { + status_code: 200, + headers: BTreeMap::from([( + "content-type".to_string(), + "text/event-stream".to_string(), + )]), + }, + })); + for chunk in chunks { + yield Ok::(ndjson_frame(StreamFrame { + frame_type: StreamFrameType::Data, + payload: StreamFramePayload::Data { + chunk_b64: None, + text: Some(chunk), + }, + })); + } + if let Some(error) = terminal_error { + yield Err::(std::io::Error::other(error)); + } else { + yield Ok::(ndjson_frame(StreamFrame::eof())); + } + } + .boxed(); + + let mut retry_scope = AiAttemptRetryScope::Candidate; + let mut fallback_response = None; + let response = execute_stream_from_frame_stream_with_retry_scope( + &state, + plan, + &format!("trace-{request_id}"), + &test_decision(), + "claude_chat_stream", + Some("claude_chat_stream_success".to_string()), + Some(json!({ + "request_id": request_id, + "candidate_id": format!("candidate-{request_id}"), + "candidate_index": 0, + "retry_index": 0, + "provider_api_format": "claude:messages", + "client_api_format": "claude:messages" + })), + crate::clock::current_unix_ms(), + Instant::now(), + RequestStageTrace::from_env(), + true, + frame_stream, + false, + None, + Some(&mut retry_scope), + Some(&mut fallback_response), + ) + .await + .expect("native Anthropic stream execution should succeed"); + match response { + Some(response) => AiAttemptExecutionOutcome::Responded(response), + None => AiAttemptExecutionOutcome::Retry { + scope: retry_scope, + fallback_response, + }, + } + } + fn test_decision() -> GatewayControlDecision { GatewayControlDecision::synthetic( "/v1/chat/completions", @@ -7176,6 +8594,425 @@ mod tests { server.abort(); } + #[tokio::test] + async fn native_anthropic_embedded_auth_error_refreshes_oauth_and_retries_once() { + let upstream_hits = Arc::new(AtomicUsize::new(0)); + let upstream_hits_for_server = Arc::clone(&upstream_hits); + let refresh_hits = Arc::new(AtomicUsize::new(0)); + let refresh_hits_for_server = Arc::clone(&refresh_hits); + let observed_authorization = Arc::new(Mutex::new(Vec::::new())); + let observed_authorization_for_server = Arc::clone(&observed_authorization); + let listener = crate::test_support::bind_loopback_listener() + .await + .expect("listener should bind"); + let addr = listener.local_addr().expect("address should resolve"); + let server = tokio::spawn(async move { + let app = Router::new() + .route( + "/v1/messages", + any(move |request: Request| { + let hits = Arc::clone(&upstream_hits_for_server); + let authorizations = Arc::clone(&observed_authorization_for_server); + async move { + authorizations + .lock() + .expect("authorization mutex should lock") + .push( + request + .headers() + .get(header::AUTHORIZATION) + .and_then(|value| value.to_str().ok()) + .unwrap_or_default() + .to_string(), + ); + let body = if hits.fetch_add(1, Ordering::SeqCst) == 0 { + concat!( + "event: error\n", + "data: {\"type\":\"error\",\"error\":{\"type\":\"authentication_error\",\"message\":\"expired token\"}}\n\n", + ) + } else { + concat!( + "event: message_start\n", + "data: {\"type\":\"message_start\",\"message\":{}}\n\n", + "event: message_stop\n", + "data: {\"type\":\"message_stop\"}\n\n", + ) + }; + ( + StatusCode::OK, + [(header::CONTENT_TYPE, "text/event-stream")], + body, + ) + .into_response() + } + }), + ) + .route( + "/oauth/token", + any(move || { + let hits = Arc::clone(&refresh_hits_for_server); + async move { + hits.fetch_add(1, Ordering::SeqCst); + Json(json!({ + "access_token": "fresh-access-token", + "refresh_token": "fresh-refresh-token", + "expires_in": 3600, + "token_type": "Bearer" + })) + } + }), + ); + axum::serve(listener, app) + .await + .expect("server should start"); + }); + + let mut plan = native_anthropic_stream_plan("anthropic-embedded-oauth-refresh"); + plan.url = format!("http://{addr}/v1/messages"); + plan.provider_name = Some("claude_code".to_string()); + plan.headers.insert( + "authorization".to_string(), + "Bearer stale-access-token".to_string(), + ); + let repository = Arc::new(provider_catalog_for_stream_auth_plan( + &plan, + "claude_code", + "oauth", + Some(generic_oauth_test_auth_config("claude_code")), + )); + let oauth_refresh = + aether_provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![ + Arc::new( + aether_provider_transport::GenericOAuthRefreshAdapter::default() + .with_token_url_for_tests( + "claude_code", + format!("http://{addr}/oauth/token"), + ), + ), + ]); + let state = AppState::new() + .expect("state should build") + .with_data_state_for_tests( + crate::data::GatewayDataState::with_provider_catalog_repository_for_tests( + repository, + ) + .with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY), + ) + .with_oauth_refresh_coordinator_for_tests(oauth_refresh); + + let execution = execute_in_process_stream_with_oauth_retry( + &state, + &mut plan, + "trace-anthropic-embedded-oauth-refresh", + None, + ) + .await + .expect("embedded authentication error should recover"); + let replayed = collect_direct_execution_body(execution) + .await + .expect("retried response should read"); + + assert_eq!(upstream_hits.load(Ordering::SeqCst), 2); + assert_eq!(refresh_hits.load(Ordering::SeqCst), 1); + assert!(String::from_utf8_lossy(&replayed).contains("event: message_start")); + assert_eq!( + observed_authorization + .lock() + .expect("authorization mutex should lock") + .as_slice(), + [ + "Bearer stale-access-token".to_string(), + "Bearer fresh-access-token".to_string(), + ] + ); + server.abort(); + } + + #[tokio::test] + async fn native_anthropic_http_permission_error_does_not_refresh_oauth() { + let upstream_hits = Arc::new(AtomicUsize::new(0)); + let upstream_hits_for_server = Arc::clone(&upstream_hits); + let refresh_hits = Arc::new(AtomicUsize::new(0)); + let refresh_hits_for_server = Arc::clone(&refresh_hits); + let permission_body = concat!( + "{\"type\":\"error\",\"error\":{", + "\"type\":\"permission_error\",", + "\"message\":\"this token is not authorized for the workspace\"}}", + ); + let listener = crate::test_support::bind_loopback_listener() + .await + .expect("listener should bind"); + let addr = listener.local_addr().expect("address should resolve"); + let server = tokio::spawn(async move { + let app = Router::new() + .route( + "/v1/messages", + any(move || { + let hits = Arc::clone(&upstream_hits_for_server); + async move { + hits.fetch_add(1, Ordering::SeqCst); + ( + StatusCode::FORBIDDEN, + [(header::CONTENT_TYPE, "application/json")], + permission_body, + ) + } + }), + ) + .route( + "/oauth/token", + any(move || { + let hits = Arc::clone(&refresh_hits_for_server); + async move { + hits.fetch_add(1, Ordering::SeqCst); + Json(json!({ + "access_token": "unexpected-access-token", + "refresh_token": "unexpected-refresh-token", + "expires_in": 3600, + "token_type": "Bearer" + })) + } + }), + ); + axum::serve(listener, app) + .await + .expect("server should start"); + }); + + let mut plan = native_anthropic_stream_plan("anthropic-http-oauth-permission"); + plan.url = format!("http://{addr}/v1/messages"); + plan.provider_name = Some("claude_code".to_string()); + plan.headers.insert( + "authorization".to_string(), + "Bearer stale-access-token".to_string(), + ); + let repository = Arc::new(provider_catalog_for_stream_auth_plan( + &plan, + "claude_code", + "oauth", + Some(generic_oauth_test_auth_config("claude_code")), + )); + let oauth_refresh = + aether_provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![ + Arc::new( + aether_provider_transport::GenericOAuthRefreshAdapter::default() + .with_token_url_for_tests( + "claude_code", + format!("http://{addr}/oauth/token"), + ), + ), + ]); + let state = AppState::new() + .expect("state should build") + .with_data_state_for_tests( + crate::data::GatewayDataState::with_provider_catalog_repository_for_tests( + repository, + ) + .with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY), + ) + .with_oauth_refresh_coordinator_for_tests(oauth_refresh); + + let execution = execute_in_process_stream_with_oauth_retry( + &state, + &mut plan, + "trace-anthropic-http-oauth-permission", + None, + ) + .await + .expect("permission response should remain available"); + assert_eq!(execution.status_code, StatusCode::FORBIDDEN.as_u16()); + let replayed = collect_direct_execution_body(execution) + .await + .expect("permission response should replay"); + + assert_eq!(upstream_hits.load(Ordering::SeqCst), 1); + assert_eq!(refresh_hits.load(Ordering::SeqCst), 0); + assert_eq!(replayed, permission_body.as_bytes()); + server.abort(); + } + + #[test] + fn native_anthropic_oauth_prefetch_respects_short_first_byte_timeout() { + let now = Instant::now(); + let precommit_started_at = now + .checked_sub(Duration::from_millis(10)) + .expect("precommit start should be representable"); + let upstream_started_at = now + .checked_sub(Duration::from_millis(90)) + .expect("upstream start should be representable"); + + let first_byte_wait = select_direct_anthropic_prefetch_wait( + precommit_started_at, + Duration::from_millis(750), + upstream_started_at, + Some(Duration::from_millis(100)), + false, + now, + ); + assert_eq!(first_byte_wait.remaining, Duration::from_millis(10)); + assert!(!first_byte_wait.commit_on_timeout); + + let precommit_wait = select_direct_anthropic_prefetch_wait( + now.checked_sub(Duration::from_millis(750)) + .expect("precommit start should be representable"), + Duration::from_millis(750), + upstream_started_at, + Some(Duration::from_secs(5)), + false, + now, + ); + assert!(precommit_wait.remaining.is_zero()); + assert!(precommit_wait.commit_on_timeout); + } + + #[tokio::test] + async fn native_anthropic_oauth_pending_events_do_not_start_a_second_precommit_wait() { + let request_id = "anthropic-oauth-single-precommit"; + let plan = native_anthropic_stream_plan(request_id); + let provider_catalog = provider_catalog_for_plan(&plan, None); + let data_state = crate::data::GatewayDataState::with_provider_transport_reader_for_tests( + Arc::new(provider_catalog), + DEVELOPMENT_ENCRYPTION_KEY, + ); + let state = AppState::new() + .expect("app state should build") + .with_data_state_for_tests(data_state); + let frame_stream = stream! { + yield Ok::(ndjson_frame(StreamFrame { + frame_type: StreamFrameType::Headers, + payload: StreamFramePayload::Headers { + status_code: 200, + headers: BTreeMap::from([( + "content-type".to_string(), + "text/event-stream".to_string(), + )]), + }, + })); + yield Ok::(ndjson_frame(StreamFrame { + frame_type: StreamFrameType::Data, + payload: StreamFramePayload::Data { + chunk_b64: None, + text: Some(": ping\n\n".to_string()), + }, + })); + yield Ok::(ndjson_frame(StreamFrame { + frame_type: StreamFrameType::Data, + payload: StreamFramePayload::Data { + chunk_b64: None, + text: Some( + "event: future_event\ndata: {\"type\":\"future_event\",\"value\":1}\n\n" + .to_string(), + ), + }, + })); + tokio::time::sleep(Duration::from_secs(2)).await; + yield Ok::(ndjson_frame(StreamFrame::eof())); + } + .boxed(); + let response = tokio::time::timeout( + Duration::from_millis(300), + execute_stream_from_frame_stream_with_retry_scope( + &state, + plan, + "trace-anthropic-oauth-single-precommit", + &test_decision(), + "claude_chat_stream", + Some("claude_chat_stream_success".to_string()), + Some(json!({ + "request_id": request_id, + "candidate_id": format!("candidate-{request_id}"), + "candidate_index": 0, + "retry_index": 0, + "provider_api_format": "claude:messages", + "client_api_format": "claude:messages" + })), + crate::clock::current_unix_ms(), + Instant::now(), + RequestStageTrace::from_env(), + true, + frame_stream, + true, + None, + None, + None, + ), + ) + .await + .expect("the committed OAuth prefetch must not be followed by another 750 ms wait") + .expect("frame stream execution should resolve"); + + assert!(response.is_some()); + } + + #[tokio::test] + async fn native_anthropic_embedded_auth_error_does_not_refresh_api_key() { + let upstream_hits = Arc::new(AtomicUsize::new(0)); + let upstream_hits_for_server = Arc::clone(&upstream_hits); + let listener = crate::test_support::bind_loopback_listener() + .await + .expect("listener should bind"); + let addr = listener.local_addr().expect("address should resolve"); + let upstream_body = concat!( + "event: error\n", + "data: {\"type\":\"error\",\"error\":{\"type\":\"authentication_error\",\"message\":\"invalid key\"}}\n\n", + ); + let server = tokio::spawn(async move { + let app = Router::new().route( + "/v1/messages", + any(move || { + let hits = Arc::clone(&upstream_hits_for_server); + async move { + hits.fetch_add(1, Ordering::SeqCst); + ( + StatusCode::OK, + [(header::CONTENT_TYPE, "text/event-stream")], + upstream_body, + ) + } + }), + ); + axum::serve(listener, app) + .await + .expect("server should start"); + }); + + let mut plan = native_anthropic_stream_plan("anthropic-embedded-api-key"); + plan.url = format!("http://{addr}/v1/messages"); + plan.provider_name = Some("claude_code".to_string()); + plan.headers + .insert("x-api-key".to_string(), "invalid-api-key".to_string()); + let repository = Arc::new(provider_catalog_for_stream_auth_plan( + &plan, + "claude_code", + "api_key", + None, + )); + let state = AppState::new() + .expect("state should build") + .with_data_state_for_tests( + crate::data::GatewayDataState::with_provider_catalog_repository_for_tests( + repository, + ) + .with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY), + ); + + let execution = execute_in_process_stream_with_oauth_retry( + &state, + &mut plan, + "trace-anthropic-embedded-api-key", + None, + ) + .await + .expect("API-key error response should remain available"); + let replayed = collect_direct_execution_body(execution) + .await + .expect("prefetched API-key error should replay"); + + assert_eq!(upstream_hits.load(Ordering::SeqCst), 1); + assert_eq!(replayed, upstream_body.as_bytes()); + server.abort(); + } + struct BlockingStreamingRequestCandidateRepository { inner: InMemoryRequestCandidateRepository, block_streaming: AtomicBool, @@ -7408,6 +9245,7 @@ mod tests { provider_body_truncated: false, client_body_truncated: false, client_stream_completion_tracker: ClientVisibleStreamCompletionTracker::default(), + requires_anthropic_message_stop: false, client_visible_stream_completed: false, usage_stream_telemetry: Some(ExecutionTelemetry { ttfb_ms: Some(7), @@ -7759,6 +9597,64 @@ mod tests { .observe_chunk(b"data: {\"type\":\"response.completed\",\"response\":{}}\r\n\r\n")); } + #[test] + fn client_visible_terminal_tracker_reports_the_exact_record_boundary() { + let message_stop = b"event: message_stop\r\ndata: {\"type\":\"message_stop\"}\r\n\r\n"; + let trailing_error = + b"event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"api_error\"}}\n\n"; + let chunk = [message_stop.as_slice(), trailing_error.as_slice()].concat(); + let mut tracker = ClientVisibleStreamCompletionTracker::default(); + + assert_eq!( + tracker.observe_chunk_terminal_end(&chunk), + Some(message_stop.len()) + ); + assert!(tracker.completed); + } + + #[test] + fn anthropic_terminal_tracker_ignores_non_message_stop_terminals() { + let mut tracker = ClientVisibleStreamCompletionTracker::default(); + + assert!(!tracker.observe_anthropic_message_stop(b"data: [DONE]\n\n")); + assert!(!tracker.observe_anthropic_message_stop( + b"event: response.completed\ndata: {\"type\":\"response.completed\"}\n\n" + )); + assert!(tracker.observe_anthropic_message_stop( + b"event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n" + )); + } + + #[test] + fn terminal_tracker_caps_multiline_record_and_resumes_after_boundary() { + let line = b"data: short-payload\r\n"; + let repeated = super::SSE_TERMINAL_DETECTOR_MAX_RECORD_BYTES / line.len() + 2; + let oversized_record = line.repeat(repeated); + let mut tracker = ClientVisibleStreamCompletionTracker::default(); + + assert!(!tracker.observe_anthropic_message_stop(&oversized_record)); + assert!(tracker.dropping_oversized_record); + assert!(tracker.line_buffer.is_empty()); + assert!(tracker.data_payload.is_empty()); + + assert!(!tracker.observe_anthropic_message_stop(b"\r\n")); + assert!(!tracker.dropping_oversized_record); + assert!(tracker.observe_anthropic_message_stop( + b"event: message_stop\r\ndata: {\"type\":\"message_stop\"}\r\n\r\n" + )); + } + + #[test] + fn stream_capture_hard_caps_a_single_oversized_chunk() { + let mut buffer = vec![1, 2]; + let mut truncated = false; + + super::append_stream_capture_bytes(&mut buffer, &[3, 4, 5, 6], 4, &mut truncated); + + assert_eq!(buffer, vec![1, 2, 3, 4]); + assert!(truncated); + } + #[test] fn provider_error_inspection_detects_response_failed_at_every_chunk_boundary() { let body = concat!( @@ -7848,6 +9744,100 @@ mod tests { Bytes::from(bytes) } + #[test] + fn post_stop_reader_yields_after_bounded_empty_chunks() { + let polls = Arc::new(AtomicUsize::new(0)); + let polls_for_stream = Arc::clone(&polls); + let stream = futures_util::stream::poll_fn(move |_| { + let poll = polls_for_stream.fetch_add(1, Ordering::SeqCst); + if poll < POST_STOP_MAX_EMPTY_CHUNKS_PER_POLL * 2 { + std::task::Poll::Ready(Some(Ok::(Bytes::new()))) + } else { + std::task::Poll::Ready(Some(Ok::(Bytes::from_static(b"x")))) + } + }); + let mut reader = Box::pin(PostStopLimitedStreamReader::new( + stream, + PostStopFrameReadBudget::new(), + )); + let waker = futures_util::task::noop_waker(); + let mut context = std::task::Context::from_waker(&waker); + let mut storage = [0u8; 1]; + let mut read_buf = tokio::io::ReadBuf::new(&mut storage); + + let result = tokio::io::AsyncRead::poll_read(reader.as_mut(), &mut context, &mut read_buf); + + assert!(result.is_pending()); + assert_eq!( + polls.load(Ordering::SeqCst), + POST_STOP_MAX_EMPTY_CHUNKS_PER_POLL + ); + } + + #[tokio::test] + async fn post_stop_activation_trims_prefetched_current_item_immediately() { + const GIANT_TAIL_BYTES: usize = 4 * 1024 * 1024; + + let mut combined = ndjson_frame(StreamFrame { + frame_type: StreamFrameType::Data, + payload: StreamFramePayload::Data { + chunk_b64: None, + text: Some( + "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n".to_string(), + ), + }, + }) + .to_vec(); + combined.resize(combined.len() + GIANT_TAIL_BYTES, b'x'); + let frame_stream = + futures_util::stream::iter([Ok::(Bytes::from(combined))]); + let reader = PostStopLimitedStreamReader::new(frame_stream, PostStopFrameReadBudget::new()); + let mut lines = + tokio_util::codec::FramedRead::new(reader, tokio_util::codec::LinesCodec::new()); + + super::read_next_frame(&mut lines) + .await + .expect("frame should decode") + .expect("data frame should exist"); + assert!(lines + .get_ref() + .current + .as_ref() + .is_some_and(|current| current.len() > ANTHROPIC_POST_STOP_DRAIN_MAX_BYTES)); + + let already_buffered = lines.read_buffer().len(); + lines.read_buffer_mut().reserve(GIANT_TAIL_BYTES); + assert!(lines.read_buffer().capacity() >= GIANT_TAIL_BYTES); + assert!(!activate_post_stop_frame_read_budget(&mut lines)); + let retained = lines + .get_ref() + .current + .as_ref() + .map(Bytes::len) + .unwrap_or_default(); + assert_eq!( + retained, + ANTHROPIC_POST_STOP_DRAIN_MAX_BYTES.saturating_sub(already_buffered) + ); + assert_eq!(lines.read_buffer().len(), already_buffered); + assert!(lines.read_buffer().capacity() <= ANTHROPIC_POST_STOP_DRAIN_MAX_BYTES); + } + + #[test] + fn post_stop_activation_releases_over_limit_framed_buffer() { + let frame_stream = futures_util::stream::empty::>(); + let reader = PostStopLimitedStreamReader::new(frame_stream, PostStopFrameReadBudget::new()); + let mut lines = + tokio_util::codec::FramedRead::new(reader, tokio_util::codec::LinesCodec::new()); + lines + .read_buffer_mut() + .resize(ANTHROPIC_POST_STOP_DRAIN_MAX_BYTES + 1, b'x'); + + assert!(activate_post_stop_frame_read_budget(&mut lines)); + assert!(lines.read_buffer().is_empty()); + assert_eq!(lines.read_buffer().capacity(), 0); + } + #[test] fn merge_stream_terminal_summary_prefers_more_complete_observed_usage() { let mut runtime_usage = StandardizedUsage::new(); @@ -8708,8 +10698,8 @@ mod tests { } #[test] - fn skips_prefetch_for_same_format_passthrough_event_streams() { - assert!(should_skip_direct_finalize_prefetch( + fn native_anthropic_event_stream_uses_bounded_precommit() { + assert!(!should_skip_direct_finalize_prefetch( Some("claude_cli_sync_finalize"), Some("text/event-stream"), "claude:messages", @@ -8720,6 +10710,648 @@ mod tests { )); } + #[test] + fn native_anthropic_terminal_error_uses_anthropic_sse_shape() { + let plan = native_anthropic_stream_plan("anthropic-terminal-error-shape"); + let failure = build_stream_failure_report( + "execution_runtime_stream_read_error", + "upstream disconnected", + 502, + ); + + let event = encode_terminal_sse_error_event_for_plan(&plan, &failure) + .expect("terminal event should encode"); + let event = String::from_utf8(event.to_vec()).expect("event should be utf8"); + + assert!(event.starts_with("event: error\ndata: ")); + assert!(event.contains("\"type\":\"error\"")); + assert!(event.contains("\"type\":\"api_error\"")); + assert!(event.contains("upstream disconnected")); + assert!(!event.contains("[DONE]")); + } + + #[tokio::test] + async fn native_anthropic_error_before_semantic_event_allows_failover() { + let unknown = "event: future_event\ndata: {\"type\":\"future_event\",\"value\":1}\n\n"; + let upstream_error = "event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",\"message\":\"Overloaded\"}}\n\n"; + let outcome = execute_native_anthropic_prefetch_stream( + "req-anthropic-precommit-error", + vec![unknown.to_string(), upstream_error.to_string()], + ) + .await; + + let AiAttemptExecutionOutcome::Retry { + scope, + fallback_response: Some(fallback_response), + } = outcome + else { + panic!("precommit 529 should retry with the upstream response preserved") + }; + assert_eq!(scope, AiAttemptRetryScope::Provider); + assert_eq!(fallback_response.status(), StatusCode::OK); + let fallback_body = to_bytes(fallback_response.into_body(), usize::MAX) + .await + .expect("fallback response body should read"); + assert_eq!( + fallback_body.as_ref(), + format!("{unknown}{upstream_error}").as_bytes() + ); + } + + #[tokio::test] + async fn native_anthropic_eof_before_semantic_event_allows_failover() { + let outcome = execute_native_anthropic_prefetch_stream( + "req-anthropic-precommit-eof", + vec![": ping\n\n".to_string()], + ) + .await; + + let AiAttemptExecutionOutcome::Retry { + scope, + fallback_response, + } = outcome + else { + panic!("EOF before the first semantic event should retry another endpoint") + }; + assert_eq!(scope, AiAttemptRetryScope::Endpoint); + assert!(fallback_response.is_none()); + } + + #[tokio::test] + async fn native_anthropic_auth_error_moves_to_the_next_credential() { + let upstream_error = concat!( + "event: error\n", + "data: {\"type\":\"error\",\"error\":{\"type\":\"authentication_error\",\"message\":\"invalid credential\"}}\n\n", + ); + let outcome = execute_native_anthropic_prefetch_stream( + "req-anthropic-precommit-auth-error", + vec![upstream_error.to_string()], + ) + .await; + + let AiAttemptExecutionOutcome::Retry { + scope, + fallback_response: Some(fallback_response), + } = outcome + else { + panic!("precommit authentication error should retry another credential") + }; + assert_eq!(scope, AiAttemptRetryScope::Credential); + let fallback_body = to_bytes(fallback_response.into_body(), usize::MAX) + .await + .expect("fallback response body should read"); + assert_eq!(fallback_body.as_ref(), upstream_error.as_bytes()); + } + + #[tokio::test] + async fn native_anthropic_semantic_event_commits_before_later_error() { + let raw = concat!( + "event: message_start\n", + "data: {\"type\":\"message_start\",\"message\":{}}\n\n", + "event: error\n", + "data: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",\"message\":\"late\"}}\n\n", + ); + let outcome = execute_native_anthropic_prefetch_stream( + "req-anthropic-postcommit-error", + vec![raw.to_string()], + ) + .await; + let AiAttemptExecutionOutcome::Responded(response) = outcome else { + panic!("a semantic event should commit the selected candidate") + }; + let body = to_bytes(response.into_body(), usize::MAX) + .await + .expect("response body should read"); + + assert_eq!(body.as_ref(), raw.as_bytes()); + } + + #[tokio::test] + async fn native_anthropic_frame_error_after_commit_emits_anthropic_terminal_event() { + let message_start = concat!( + "event: message_start\n", + "data: {\"type\":\"message_start\",\"message\":{}}\n\n", + ); + let original_error = "upstream disconnected after message_start"; + let outcome = execute_native_anthropic_prefetch_stream_with_terminal_error( + "req-anthropic-postcommit-frame-error", + vec![message_start.to_string()], + Some(original_error.to_string()), + ) + .await; + let AiAttemptExecutionOutcome::Responded(response) = outcome else { + panic!("a semantic event should commit the selected candidate") + }; + let body = to_bytes(response.into_body(), usize::MAX) + .await + .expect("response body should read"); + let body = String::from_utf8(body.to_vec()).expect("body should be utf8"); + + assert!(body.starts_with(message_start)); + assert!(body.contains("event: error\ndata: {\"type\":\"error\"")); + assert!(body.contains("\"type\":\"api_error\"")); + assert!(body.contains(original_error)); + assert!(!body.contains("[DONE]")); + } + + #[tokio::test] + async fn native_anthropic_eof_after_commit_emits_anthropic_terminal_event() { + let message_start = concat!( + "event: message_start\n", + "data: {\"type\":\"message_start\",\"message\":{}}\n\n", + ); + let outcome = execute_native_anthropic_prefetch_stream( + "req-anthropic-postcommit-eof", + vec![message_start.to_string()], + ) + .await; + let AiAttemptExecutionOutcome::Responded(response) = outcome else { + panic!("a semantic event should commit the selected candidate") + }; + let body = to_bytes(response.into_body(), usize::MAX) + .await + .expect("response body should read"); + let body = String::from_utf8(body.to_vec()).expect("body should be utf8"); + + assert!(body.starts_with(message_start)); + assert!(body.contains("event: error\ndata: {\"type\":\"error\"")); + assert!(body.contains("ended before message_stop")); + assert!(!body.contains("[DONE]")); + } + + #[tokio::test] + async fn native_anthropic_done_marker_does_not_replace_message_stop() { + let message_start = concat!( + "event: message_start\n", + "data: {\"type\":\"message_start\",\"message\":{}}\n\n", + ); + let done = "data: [DONE]\n\n"; + let outcome = execute_native_anthropic_prefetch_stream( + "req-anthropic-done-without-message-stop", + vec![message_start.to_string(), done.to_string()], + ) + .await; + let AiAttemptExecutionOutcome::Responded(response) = outcome else { + panic!("message_start should commit the selected candidate") + }; + let body = to_bytes(response.into_body(), usize::MAX) + .await + .expect("response body should read"); + let body = String::from_utf8(body.to_vec()).expect("body should be utf8"); + + assert!(body.starts_with(message_start)); + assert!(body.contains(done)); + assert!(body.contains("event: error\ndata: {\"type\":\"error\"")); + assert!(body.contains("ended before message_stop")); + } + + #[tokio::test] + async fn native_anthropic_hanging_tail_does_not_delay_body_eof_and_is_bounded() { + let request_id = "req-anthropic-hanging-tail"; + let plan = native_anthropic_stream_plan(request_id); + let provider_catalog = provider_catalog_for_plan(&plan, None); + let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); + let state = AppState::new() + .expect("app state should build") + .with_data_state_for_tests( + crate::data::GatewayDataState::with_usage_repository_for_tests(Arc::clone( + &usage_repository, + )) + .with_provider_catalog_reader(Arc::new(provider_catalog)) + .with_encryption_key_for_tests("development-key"), + ) + .with_usage_runtime_for_tests(UsageRuntimeConfig { + enabled: true, + ..UsageRuntimeConfig::default() + }); + let message_start = concat!( + "event: message_start\n", + "data: {\"type\":\"message_start\",\"message\":{}}\n\n", + ); + let message_stop = "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"; + let stream_dropped = Arc::new(AtomicBool::new(false)); + let drop_flag = StreamDropFlag(Arc::clone(&stream_dropped)); + let frame_stream = stream! { + let _drop_flag = drop_flag; + yield Ok::(ndjson_frame(StreamFrame { + frame_type: StreamFrameType::Headers, + payload: StreamFramePayload::Headers { + status_code: 200, + headers: BTreeMap::from([( + "content-type".to_string(), + "text/event-stream".to_string(), + )]), + }, + })); + yield Ok::(ndjson_frame(StreamFrame { + frame_type: StreamFrameType::Data, + payload: StreamFramePayload::Data { + chunk_b64: None, + text: Some(message_start.to_string()), + }, + })); + yield Ok::(ndjson_frame(StreamFrame { + frame_type: StreamFrameType::Data, + payload: StreamFramePayload::Data { + chunk_b64: None, + text: Some(message_stop.to_string()), + }, + })); + std::future::pending::<()>().await; + } + .boxed(); + + let response = execute_stream_from_frame_stream( + &state, + plan, + "trace-anthropic-hanging-tail", + &test_decision(), + "claude_chat_stream", + Some("claude_chat_stream_success".to_string()), + Some(json!({ + "request_id": request_id, + "candidate_id": format!("candidate-{request_id}"), + "candidate_index": 0, + "retry_index": 0, + "provider_api_format": "claude:messages", + "client_api_format": "claude:messages" + })), + crate::clock::current_unix_ms(), + Instant::now(), + RequestStageTrace::from_env(), + false, + frame_stream, + None, + ) + .await + .expect("stream execution should succeed") + .expect("stream execution should return a response"); + + let body = tokio::time::timeout( + Duration::from_secs(1), + to_bytes(response.into_body(), usize::MAX), + ) + .await + .expect("client body EOF must not wait for the hanging producer tail") + .expect("client body should read"); + assert_eq!( + body.as_ref(), + format!("{message_start}{message_stop}").as_bytes() + ); + + tokio::time::timeout(Duration::from_secs(1), async { + while !stream_dropped.load(Ordering::SeqCst) { + tokio::task::yield_now().await; + } + }) + .await + .expect("producer tail should be dropped after the bounded drain window"); + + let stored_usage = tokio::time::timeout(Duration::from_secs(2), async { + loop { + let usage = usage_repository + .find_by_request_id(request_id) + .await + .expect("usage should read"); + if usage + .as_ref() + .is_some_and(|usage| usage.status == "completed") + { + break usage.expect("completed usage should exist"); + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("bounded drain timeout should still settle usage successfully"); + assert_eq!(stored_usage.status_code, Some(200)); + } + + #[tokio::test] + async fn native_anthropic_unterminated_oversized_tail_respects_read_budget() { + const TAIL_CHUNK_BYTES: usize = 4 * 1024 * 1024; + + let request_id = "req-anthropic-oversized-tail"; + let plan = native_anthropic_stream_plan(request_id); + let provider_catalog = provider_catalog_for_plan(&plan, None); + let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); + let state = AppState::new() + .expect("app state should build") + .with_data_state_for_tests( + crate::data::GatewayDataState::with_usage_repository_for_tests(Arc::clone( + &usage_repository, + )) + .with_provider_catalog_reader(Arc::new(provider_catalog)) + .with_encryption_key_for_tests("development-key"), + ) + .with_usage_runtime_for_tests(UsageRuntimeConfig { + enabled: true, + ..UsageRuntimeConfig::default() + }); + let message_start = concat!( + "event: message_start\n", + "data: {\"type\":\"message_start\",\"message\":{}}\n\n", + ); + let message_stop = "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"; + let stream_dropped = Arc::new(AtomicBool::new(false)); + let drop_flag = StreamDropFlag(Arc::clone(&stream_dropped)); + let tail_chunks_polled = Arc::new(AtomicUsize::new(0)); + let tail_chunks_polled_for_stream = Arc::clone(&tail_chunks_polled); + let tail_chunk = Bytes::from(vec![b'x'; TAIL_CHUNK_BYTES]); + let tail_chunk_count = 32; + let frame_stream = stream! { + let _drop_flag = drop_flag; + yield Ok::(ndjson_frame(StreamFrame { + frame_type: StreamFrameType::Headers, + payload: StreamFramePayload::Headers { + status_code: 200, + headers: BTreeMap::from([( + "content-type".to_string(), + "text/event-stream".to_string(), + )]), + }, + })); + yield Ok::(ndjson_frame(StreamFrame { + frame_type: StreamFrameType::Data, + payload: StreamFramePayload::Data { + chunk_b64: None, + text: Some(message_start.to_string()), + }, + })); + yield Ok::(ndjson_frame(StreamFrame { + frame_type: StreamFrameType::Data, + payload: StreamFramePayload::Data { + chunk_b64: None, + text: Some(message_stop.to_string()), + }, + })); + for _ in 0..tail_chunk_count { + tail_chunks_polled_for_stream.fetch_add(1, Ordering::SeqCst); + yield Ok::(tail_chunk.clone()); + } + std::future::pending::<()>().await; + } + .boxed(); + + let response = execute_stream_from_frame_stream( + &state, + plan, + "trace-anthropic-oversized-tail", + &test_decision(), + "claude_chat_stream", + Some("claude_chat_stream_success".to_string()), + Some(json!({ + "request_id": request_id, + "candidate_id": format!("candidate-{request_id}"), + "candidate_index": 0, + "retry_index": 0, + "provider_api_format": "claude:messages", + "client_api_format": "claude:messages" + })), + crate::clock::current_unix_ms(), + Instant::now(), + RequestStageTrace::from_env(), + false, + frame_stream, + None, + ) + .await + .expect("stream execution should succeed") + .expect("stream execution should return a response"); + + let body = tokio::time::timeout( + Duration::from_secs(1), + to_bytes(response.into_body(), usize::MAX), + ) + .await + .expect("client body EOF must not wait for the oversized unterminated tail") + .expect("client body should read"); + assert_eq!( + body.as_ref(), + format!("{message_start}{message_stop}").as_bytes() + ); + + tokio::time::timeout(Duration::from_secs(1), async { + while !stream_dropped.load(Ordering::SeqCst) { + tokio::task::yield_now().await; + } + }) + .await + .expect("oversized tail producer should be released by the read budget"); + assert!( + tail_chunks_polled.load(Ordering::SeqCst) <= 1, + "post-stop drain must retain at most one atomic upstream stream item" + ); + + let stored_usage = tokio::time::timeout(Duration::from_secs(2), async { + loop { + let usage = usage_repository + .find_by_request_id(request_id) + .await + .expect("usage should read"); + if usage + .as_ref() + .is_some_and(|usage| usage.status == "completed") + { + break usage.expect("completed usage should exist"); + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("bounded oversized tail should still settle usage successfully"); + assert_eq!(stored_usage.status_code, Some(200)); + } + + #[tokio::test] + async fn native_anthropic_frame_error_after_message_stop_is_ignored() { + let message_start = concat!( + "event: message_start\n", + "data: {\"type\":\"message_start\",\"message\":{}}\n\n", + ); + let message_stop = "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"; + let outcome = execute_native_anthropic_prefetch_stream_with_terminal_error( + "req-anthropic-error-after-message-stop", + vec![message_start.to_string(), message_stop.to_string()], + Some("connection reset after message_stop".to_string()), + ) + .await; + let AiAttemptExecutionOutcome::Responded(response) = outcome else { + panic!("message_stop should keep the selected candidate committed") + }; + let body = to_bytes(response.into_body(), usize::MAX) + .await + .expect("response body should read"); + + assert_eq!( + body.as_ref(), + format!("{message_start}{message_stop}").as_bytes() + ); + } + + #[tokio::test] + async fn native_anthropic_same_chunk_stops_at_message_stop_record() { + let message_start = concat!( + "event: message_start\n", + "data: {\"type\":\"message_start\",\"message\":{}}\n\n", + ); + let message_stop = "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"; + let trailing_error = concat!( + "event: error\n", + "data: {\"type\":\"error\",\"error\":{\"type\":\"api_error\",\"message\":\"after stop\"}}\n\n", + ); + let outcome = execute_native_anthropic_prefetch_stream( + "req-anthropic-same-chunk-message-stop", + vec![format!("{message_start}{message_stop}{trailing_error}")], + ) + .await; + let AiAttemptExecutionOutcome::Responded(response) = outcome else { + panic!("message_stop should complete the selected candidate") + }; + let body = to_bytes(response.into_body(), usize::MAX) + .await + .expect("response body should read"); + + assert_eq!( + body.as_ref(), + format!("{message_start}{message_stop}").as_bytes() + ); + } + + #[tokio::test] + async fn direct_anthropic_stops_at_message_stop_and_ignores_teardown_error() { + let message_start = Bytes::from_static( + b"event: message_start\ndata: {\"type\":\"message_start\",\"message\":{}}\n\n", + ); + let message_stop = "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"; + let trailing_error = concat!( + "event: error\n", + "data: {\"type\":\"error\",\"error\":{\"type\":\"api_error\",\"message\":\"after stop\"}}\n\n", + ); + let state = direct_anthropic_inline_state( + "req-direct-anthropic-message-stop", + vec![ + Ok(message_start.clone()), + Ok(Bytes::from(format!("{message_stop}{trailing_error}"))), + Err("connection reset after message_stop".to_string()), + ], + ); + + let (first, state) = state + .next_item() + .await + .expect("message_start should stream"); + assert_eq!(first.expect("message_start should succeed"), message_start); + let (second, state) = state.next_item().await.expect("message_stop should stream"); + assert_eq!( + second.expect("message_stop should succeed").as_ref(), + message_stop.as_bytes() + ); + assert!(state.next_item().await.is_none()); + } + + #[tokio::test] + async fn direct_anthropic_clean_eof_after_message_start_emits_one_error() { + let message_start = Bytes::from_static( + b"event: message_start\ndata: {\"type\":\"message_start\",\"message\":{}}\n\n", + ); + let state = direct_anthropic_inline_state( + "req-direct-anthropic-premature-eof", + vec![Ok(message_start.clone())], + ); + + let (first, state) = state + .next_item() + .await + .expect("message_start should stream"); + assert_eq!(first.expect("message_start should succeed"), message_start); + let (error, mut state) = state + .next_item() + .await + .expect("premature EOF should emit an Anthropic error event"); + let error = String::from_utf8(error.expect("error event should succeed").to_vec()) + .expect("error event should be utf8"); + assert!(error.starts_with("event: error\ndata: ")); + assert!(error.contains("ended before message_stop")); + discard_direct_test_finalizer(&mut state); + } + + #[tokio::test] + async fn direct_anthropic_provider_error_is_not_duplicated() { + let message_start = Bytes::from_static( + b"event: message_start\ndata: {\"type\":\"message_start\",\"message\":{}}\n\n", + ); + let provider_error = Bytes::from_static( + b"event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",\"message\":\"busy\"}}\n\n", + ); + let state = direct_anthropic_inline_state( + "req-direct-anthropic-provider-error", + vec![Ok(message_start.clone()), Ok(provider_error.clone())], + ); + + let (first, state) = state + .next_item() + .await + .expect("message_start should stream"); + assert_eq!(first.expect("message_start should succeed"), message_start); + let (second, state) = state + .next_item() + .await + .expect("provider error should stream"); + assert_eq!( + second.expect("provider error should succeed"), + provider_error + ); + assert!(state.terminal_error_sent); + assert!(state.next_item().await.is_none()); + } + + #[test] + fn postcommit_anthropic_errors_use_the_precommit_status_taxonomy() { + for (error_type, expected_status) in [ + ("request_too_large", 413), + ("overloaded_error", 529), + ("api_error", 500), + ] { + let body = json!({ + "type": "error", + "error": { "type": error_type, "message": "upstream failure" } + }); + assert_eq!( + resolve_provider_stream_error_status_code("claude:messages", 200, &body), + expected_status, + ); + } + } + + #[tokio::test] + async fn native_anthropic_unknown_event_is_replayed_byte_for_byte() { + let unknown = "event: future_event\ndata: {\"type\":\"future_event\",\"value\":1}\n\n"; + let message_start = concat!( + "event: message_start\n", + "data: {\"type\":\"message_start\",\"message\":{}}\n\n", + ); + let message_stop = "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"; + let expected = format!("{unknown}{message_start}{message_stop}"); + let outcome = execute_native_anthropic_prefetch_stream( + "req-anthropic-unknown-replay", + vec![ + unknown.to_string(), + message_start.to_string(), + message_stop.to_string(), + ], + ) + .await; + let AiAttemptExecutionOutcome::Responded(response) = outcome else { + panic!("unknown event should not terminate the stream") + }; + let body = to_bytes(response.into_body(), usize::MAX) + .await + .expect("response body should read"); + + assert_eq!(body.as_ref(), expected.as_bytes()); + } + #[test] fn skips_prefetch_for_same_format_passthrough_streams_without_content_type() { assert!(should_skip_direct_finalize_prefetch( @@ -8868,6 +11500,38 @@ mod tests { )); } + #[tokio::test] + async fn native_anthropic_sse_body_ends_at_message_stop_while_sender_is_alive() { + let (tx, rx) = mpsc::channel::>(1); + let message_stop = + Bytes::from_static(b"event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"); + tx.send(Ok(message_stop.clone())) + .await + .expect("message_stop should send"); + let mut body_stream = Box::pin(build_sse_body_stream( + Vec::new(), + rx, + true, + false, + true, + Duration::from_secs(60), + )); + + let chunk = tokio::time::timeout(Duration::from_millis(50), body_stream.next()) + .await + .expect("message_stop should arrive immediately") + .expect("stream should yield message_stop") + .expect("message_stop should be successful"); + assert_eq!(chunk, message_stop); + assert!( + tokio::time::timeout(Duration::from_millis(50), body_stream.next()) + .await + .expect("body EOF must not wait for the producer") + .is_none() + ); + assert!(tx.is_closed(), "body EOF should drop the receiver"); + } + #[tokio::test] async fn sse_body_stream_emits_initial_and_periodic_keepalive_without_business_chunks() { let (_tx, rx) = mpsc::channel::>(1); @@ -8876,6 +11540,7 @@ mod tests { rx, true, true, + false, Duration::from_millis(10), )); @@ -8902,6 +11567,7 @@ mod tests { rx, true, false, + false, Duration::from_millis(10), )); @@ -8942,6 +11608,7 @@ mod tests { rx, true, true, + false, Duration::from_secs(60), )); @@ -8971,6 +11638,7 @@ mod tests { rx, true, true, + false, Duration::from_secs(60), )); @@ -8994,6 +11662,7 @@ mod tests { rx, true, true, + false, Duration::from_secs(60), )); @@ -9050,6 +11719,7 @@ mod tests { rx, true, true, + false, Duration::from_secs(60), )); diff --git a/apps/aether-gateway/src/execution_runtime/stream/execution_failures.rs b/apps/aether-gateway/src/execution_runtime/stream/execution_failures.rs index 33c8af2d3..57e1df1b2 100644 --- a/apps/aether-gateway/src/execution_runtime/stream/execution_failures.rs +++ b/apps/aether-gateway/src/execution_runtime/stream/execution_failures.rs @@ -1,3 +1,4 @@ +use aether_ai_serving::AiAttemptRetryScope; use aether_contracts::{ExecutionError, ExecutionPlan, ExecutionTelemetry}; use aether_data_contracts::repository::candidates::RequestCandidateStatus; use aether_scheduler_core::SchedulerRequestCandidateStatusUpdate; @@ -14,16 +15,17 @@ use tracing::warn; use crate::api::response::attach_control_metadata_headers; use crate::clock::current_unix_ms as current_request_candidate_unix_ms; use crate::control::GatewayControlDecision; +use crate::execution_runtime::ai_attempt_retry_scope_from_failure_disposition; use crate::execution_runtime::submission::{ resolve_core_error_background_report_kind, submit_local_core_error_or_sync_finalize, }; use crate::log_ids::short_request_id; use crate::orchestration::{ - apply_local_execution_effect, resolve_local_failover_analysis_for_attempt, - with_upstream_response_report_context, LocalAdaptiveRateLimitEffect, LocalAttemptFailureEffect, - LocalExecutionEffect, LocalExecutionEffectContext, LocalFailoverAnalysis, - LocalFailoverDecision, LocalHealthFailureEffect, LocalOAuthInvalidationEffect, - LocalPoolErrorEffect, + apply_local_execution_effect, classify_failure_disposition, + resolve_local_failover_analysis_for_attempt, with_upstream_response_report_context, + LocalAdaptiveRateLimitEffect, LocalAttemptFailureEffect, LocalExecutionEffect, + LocalExecutionEffectContext, LocalFailoverAnalysis, LocalFailoverDecision, + LocalHealthFailureEffect, LocalOAuthInvalidationEffect, LocalPoolErrorEffect, }; use crate::request_candidate_runtime::record_report_request_candidate_status; use crate::request_diagnostics::attach_current_request_diagnostics_to_report_context; @@ -409,9 +411,13 @@ pub(super) async fn handle_prefetch_provider_private_stream_error( mut headers: std::collections::BTreeMap, telemetry: Option, buffered_body: &[u8], + upstream_status_code: u16, status_code: u16, body_json: Value, + retry_scope_out: Option<&mut AiAttemptRetryScope>, + retry_fallback_out: Option<&mut Option>>, ) -> Result>, GatewayError> { + let upstream_headers = headers.clone(); headers.remove("content-encoding"); headers.remove("content-length"); headers.insert("content-type".to_string(), "application/json".to_string()); @@ -441,6 +447,29 @@ pub(super) async fn handle_prefetch_provider_private_stream_error( failure_analysis.decision, LocalFailoverDecision::RetryNextCandidate ) { + let failure_disposition = classify_failure_disposition( + &plan.provider_api_format, + failure_analysis.classification, + status_code, + ); + if let Some(retry_scope) = retry_scope_out { + *retry_scope = ai_attempt_retry_scope_from_failure_disposition(failure_disposition); + } + if failure_disposition.preserve_upstream_error { + if let Some(retry_fallback) = retry_fallback_out { + *retry_fallback = Some(attach_control_metadata_headers( + crate::api::response::build_client_response_from_parts( + upstream_status_code, + &upstream_headers, + Body::from(buffered_body.to_vec()), + trace_id, + Some(decision), + )?, + Some(request_id), + candidate_id, + )?); + } + } warn!( event_name = "local_stream_candidate_retry_scheduled", log_type = "event", diff --git a/apps/aether-gateway/src/execution_runtime/stream/mod.rs b/apps/aether-gateway/src/execution_runtime/stream/mod.rs index 486c6403f..c45389853 100644 --- a/apps/aether-gateway/src/execution_runtime/stream/mod.rs +++ b/apps/aether-gateway/src/execution_runtime/stream/mod.rs @@ -1,4 +1,7 @@ +mod commit_policy; mod error; mod execution; -pub(crate) use execution::execute_execution_runtime_stream; +pub(crate) use execution::{ + execute_execution_runtime_stream, execute_execution_runtime_stream_with_retry_scope, +}; diff --git a/apps/aether-gateway/src/execution_runtime/stream_pump.rs b/apps/aether-gateway/src/execution_runtime/stream_pump.rs index adf70bb57..3f24c10ed 100644 --- a/apps/aether-gateway/src/execution_runtime/stream_pump.rs +++ b/apps/aether-gateway/src/execution_runtime/stream_pump.rs @@ -21,12 +21,14 @@ use crate::ai_serving::api::{ }; use crate::execution_runtime::ndjson::encode_stream_frame_ndjson; use crate::execution_runtime::transport::{ - format_hyper_error_chain, format_wreq_upstream_request_error, - stream_first_byte_timeout_message, DirectUpstreamResponse, + append_upstream_response_body_chunk, decode_response_body_bytes, format_hyper_error_chain, + format_wreq_upstream_request_error, stream_first_byte_timeout_message, DirectUpstreamResponse, }; use crate::execution_runtime::DirectUpstreamStreamExecution; use crate::GatewayError; +const STREAM_USAGE_OBSERVER_MAX_LINE_BYTES: usize = 1024 * 1024; + pub(crate) fn build_direct_execution_frame_stream( execution: DirectUpstreamStreamExecution, ) -> impl Stream> + Send + 'static { @@ -39,6 +41,7 @@ pub(crate) fn build_direct_execution_frame_stream( provider_api_format, stream_summary_report_context, prefetched_body, + stream_precommit_committed: _, response, started_at, stream_first_byte_timeout, @@ -712,6 +715,23 @@ struct BufferedUpstreamBodyError { first_byte_timeout: Option, } +fn append_buffered_upstream_body_chunk( + body_bytes: &mut Vec, + chunk: &[u8], + ttfb_ms: Option, + upstream_bytes: &mut u64, +) -> Result<(), BufferedUpstreamBodyError> { + *upstream_bytes = upstream_bytes.saturating_add(chunk.len() as u64); + append_upstream_response_body_chunk(body_bytes, chunk).map_err(|error| { + BufferedUpstreamBodyError { + message: error.to_string(), + ttfb_ms, + upstream_bytes: *upstream_bytes, + first_byte_timeout: None, + } + }) +} + fn response_headers_indicate_sse(headers: &BTreeMap) -> bool { headers .get("content-type") @@ -770,8 +790,12 @@ async fn buffer_non_sse_upstream_body( if ttfb_ms.is_none() { ttfb_ms = Some(started_at.elapsed().as_millis() as u64); } - upstream_bytes += chunk.len() as u64; - body_bytes.extend_from_slice(&chunk); + append_buffered_upstream_body_chunk( + &mut body_bytes, + &chunk, + ttfb_ms, + &mut upstream_bytes, + )?; } Err(message) => { return Err(BufferedUpstreamBodyError { @@ -817,8 +841,12 @@ async fn buffer_non_sse_upstream_body( if ttfb_ms.is_none() { ttfb_ms = Some(started_at.elapsed().as_millis() as u64); } - upstream_bytes += chunk.len() as u64; - body_bytes.extend_from_slice(&chunk); + append_buffered_upstream_body_chunk( + &mut body_bytes, + &chunk, + ttfb_ms, + &mut upstream_bytes, + )?; } Err(err) => { let message = format_error_chain(&err); @@ -871,8 +899,12 @@ async fn buffer_non_sse_upstream_body( if ttfb_ms.is_none() { ttfb_ms = Some(started_at.elapsed().as_millis() as u64); } - upstream_bytes += chunk.len() as u64; - body_bytes.extend_from_slice(&chunk); + append_buffered_upstream_body_chunk( + &mut body_bytes, + &chunk, + ttfb_ms, + &mut upstream_bytes, + )?; } Err(err) => { let message = format_hyper_error_chain(&err); @@ -925,8 +957,12 @@ async fn buffer_non_sse_upstream_body( if ttfb_ms.is_none() { ttfb_ms = Some(started_at.elapsed().as_millis() as u64); } - upstream_bytes += chunk.len() as u64; - body_bytes.extend_from_slice(&chunk); + append_buffered_upstream_body_chunk( + &mut body_bytes, + &chunk, + ttfb_ms, + &mut upstream_bytes, + )?; } Err(err) => { let message = format_wreq_upstream_request_error(&err); @@ -974,8 +1010,12 @@ async fn buffer_non_sse_upstream_body( if ttfb_ms.is_none() { ttfb_ms = Some(started_at.elapsed().as_millis() as u64); } - upstream_bytes += chunk.len() as u64; - body_bytes.extend_from_slice(&chunk); + append_buffered_upstream_body_chunk( + &mut body_bytes, + &chunk, + ttfb_ms, + &mut upstream_bytes, + )?; } Ok(None) => break, Err(message) => { @@ -1015,13 +1055,13 @@ fn maybe_bridge_non_sse_sync_json_to_stream( return Ok(None); } - let decoded_body_bytes = decode_non_sse_response_body_bytes(headers, body_bytes) - .unwrap_or_else(|| body_bytes.to_vec()); - if !response_body_is_json(headers, &decoded_body_bytes) { + let decoded_body_bytes = decode_response_body_bytes(headers, body_bytes) + .map_err(|error| GatewayError::Internal(error.to_string()))?; + if !response_body_is_json(headers, decoded_body_bytes.as_ref()) { return Ok(None); } - let body_json: Value = serde_json::from_slice(&decoded_body_bytes) + let body_json: Value = serde_json::from_slice(decoded_body_bytes.as_ref()) .map_err(|err| GatewayError::Internal(err.to_string()))?; let client_api_format = report_context .get("client_api_format") @@ -1046,33 +1086,6 @@ fn rewrite_headers_for_bridged_sse_response( rewritten } -fn decode_non_sse_response_body_bytes( - headers: &BTreeMap, - body_bytes: &[u8], -) -> Option> { - let encoding = headers - .get("content-encoding") - .map(String::as_str) - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(|value| value.to_ascii_lowercase()); - match encoding.as_deref() { - Some("gzip") => { - let mut decoder = flate2::read::GzDecoder::new(body_bytes); - let mut out = Vec::new(); - std::io::Read::read_to_end(&mut decoder, &mut out).ok()?; - Some(out) - } - Some("deflate") => { - let mut decoder = flate2::read::DeflateDecoder::new(body_bytes); - let mut out = Vec::new(); - std::io::Read::read_to_end(&mut decoder, &mut out).ok()?; - Some(out) - } - _ => None, - } -} - fn response_body_is_json(headers: &BTreeMap, body_bytes: &[u8]) -> bool { if headers .get("content-type") @@ -1159,16 +1172,39 @@ fn observe_normalized_bytes( observer_buffered: &mut Vec, normalized: &[u8], ) { - if normalized.is_empty() { + if normalized.is_empty() + || observer + .latest_summary() + .and_then(|summary| summary.parser_error.as_deref()) + .is_some() + { return; } - observer_buffered.extend_from_slice(normalized); - while let Some(line_end) = observer_buffered.iter().position(|byte| *byte == b'\n') { - let line = observer_buffered.drain(..=line_end).collect::>(); - if let Err(err) = observer.push_line(report_context, line) { - observer.disable_with_error(err.to_string()); + + let mut remaining = normalized; + while !remaining.is_empty() { + let line_part_len = remaining + .iter() + .position(|byte| *byte == b'\n') + .map_or(remaining.len(), |index| index + 1); + if observer_buffered.len().saturating_add(line_part_len) + > STREAM_USAGE_OBSERVER_MAX_LINE_BYTES + { + observer.disable_with_error(format!( + "stream usage event exceeded {STREAM_USAGE_OBSERVER_MAX_LINE_BYTES} bytes" + )); observer_buffered.clear(); - break; + return; + } + observer_buffered.extend_from_slice(&remaining[..line_part_len]); + remaining = &remaining[line_part_len..]; + if observer_buffered.last() == Some(&b'\n') { + let line = std::mem::take(observer_buffered); + if let Err(err) = observer.push_line(report_context, line) { + observer.disable_with_error(err.to_string()); + observer_buffered.clear(); + return; + } } } } @@ -1193,9 +1229,11 @@ mod tests { use tokio::sync::watch; use super::{ - build_direct_execution_frame_stream, should_buffer_non_stream_response, - should_treat_upstream_response_as_stream, + build_direct_execution_frame_stream, observe_normalized_bytes, + should_buffer_non_stream_response, should_treat_upstream_response_as_stream, + STREAM_USAGE_OBSERVER_MAX_LINE_BYTES, }; + use crate::ai_serving::api::StreamingStandardTerminalObserver; use crate::execution_runtime::transport::{ execute_stream_plan_via_local_tunnel, DirectSyncExecutionRuntime, DirectUpstreamResponse, }; @@ -1260,6 +1298,27 @@ mod tests { )); } + #[test] + fn oversized_usage_line_disables_observation_without_retaining_the_line() { + let mut observer = StreamingStandardTerminalObserver::default(); + let report_context = serde_json::json!({ + "provider_api_format": "claude:messages", + "client_api_format": "claude:messages", + }); + let mut buffered = Vec::new(); + let oversized = vec![b'x'; STREAM_USAGE_OBSERVER_MAX_LINE_BYTES + 1]; + + observe_normalized_bytes(&mut observer, &report_context, &mut buffered, &oversized); + + assert!(buffered.is_empty()); + assert!(observer + .latest_summary() + .and_then(|summary| summary.parser_error.as_deref()) + .is_some_and(|error| error.contains("stream usage event exceeded"))); + observe_normalized_bytes(&mut observer, &report_context, &mut buffered, b"ignored"); + assert!(buffered.is_empty()); + } + #[tokio::test] async fn direct_execution_frame_stream_reports_ttfb_after_first_upstream_chunk() { let listener = crate::test_support::bind_loopback_listener() diff --git a/apps/aether-gateway/src/execution_runtime/submission.rs b/apps/aether-gateway/src/execution_runtime/submission.rs index 02b2b9010..861b5792b 100644 --- a/apps/aether-gateway/src/execution_runtime/submission.rs +++ b/apps/aether-gateway/src/execution_runtime/submission.rs @@ -496,6 +496,15 @@ fn classify_local_sync_error_kind( { return LocalCoreSyncErrorKind::RateLimit; } + if status_code == 413 + || fingerprint.contains("request_too_large") + || fingerprint.contains("request too large") + || fingerprint.contains("payload_too_large") + || fingerprint.contains("payload too large") + || fingerprint.contains("request entity too large") + { + return LocalCoreSyncErrorKind::RequestTooLarge; + } if fingerprint.contains("contextlength") || fingerprint.contains("contentlengthexceeded") || fingerprint.contains("context window") @@ -534,6 +543,7 @@ fn default_status_code_for_local_sync_error_kind(kind: LocalCoreSyncErrorKind) - LocalCoreSyncErrorKind::InvalidRequest | LocalCoreSyncErrorKind::ContextLengthExceeded => { 400 } + LocalCoreSyncErrorKind::RequestTooLarge => 413, LocalCoreSyncErrorKind::Authentication => 401, LocalCoreSyncErrorKind::PermissionDenied => 403, LocalCoreSyncErrorKind::NotFound => 404, @@ -792,6 +802,76 @@ mod tests { ); } + #[tokio::test] + async fn local_core_error_maps_request_too_large_without_changing_openai_shape() { + let claude_payload = core_finalize_payload( + "claude_chat_sync_finalize", + "claude:messages", + "openai:chat", + 413, + json!({ + "error": { + "type": "invalid_request_error", + "message": "request body is too large" + } + }), + ); + let claude_response = maybe_build_local_core_error_response( + "trace-sync-claude-too-large", + &test_decision(), + &claude_payload, + ) + .expect("response build should not error") + .expect("response should exist"); + assert_eq!( + claude_response.status(), + http::StatusCode::PAYLOAD_TOO_LARGE + ); + let claude_body: serde_json::Value = serde_json::from_slice( + &to_bytes(claude_response.into_body(), usize::MAX) + .await + .expect("body should read"), + ) + .expect("body should decode"); + assert_eq!(claude_body["type"], "error"); + assert_eq!(claude_body["error"]["type"], "request_too_large"); + + let openai_payload = core_finalize_payload( + "openai_chat_sync_finalize", + "openai:chat", + "claude:messages", + 200, + json!({ + "type": "error", + "error": { + "type": "request_too_large", + "message": "request body is too large" + } + }), + ); + let openai_response = maybe_build_local_core_error_response( + "trace-sync-openai-too-large", + &test_decision(), + &openai_payload, + ) + .expect("response build should not error") + .expect("response should exist"); + assert_eq!( + openai_response.status(), + http::StatusCode::PAYLOAD_TOO_LARGE + ); + let openai_body: serde_json::Value = serde_json::from_slice( + &to_bytes(openai_response.into_body(), usize::MAX) + .await + .expect("body should read"), + ) + .expect("body should decode"); + assert_eq!( + openai_body["error"]["type"], "context_length_exceeded", + "OpenAI compatibility shape should remain unchanged" + ); + } + #[tokio::test] async fn local_core_sync_finalize_rejects_gemini_http_200_without_visible_output() { let mut payload = core_finalize_payload( diff --git a/apps/aether-gateway/src/execution_runtime/sync/execution.rs b/apps/aether-gateway/src/execution_runtime/sync/execution.rs index 4a712d56c..61d35b812 100644 --- a/apps/aether-gateway/src/execution_runtime/sync/execution.rs +++ b/apps/aether-gateway/src/execution_runtime/sync/execution.rs @@ -3,7 +3,7 @@ use std::io::Error as IoError; use std::sync::Arc; use std::time::{Duration, Instant}; -use aether_ai_serving::UPSTREAM_IS_STREAM_KEY; +use aether_ai_serving::{AiAttemptExecutionOutcome, AiAttemptRetryScope, UPSTREAM_IS_STREAM_KEY}; use aether_contracts::{ ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionPlan, ExecutionResult, ExecutionTelemetry, @@ -55,17 +55,18 @@ use crate::execution_runtime::submission::{ resolve_local_sync_error_status_code, submit_local_core_error_or_sync_finalize, }; use crate::execution_runtime::transport::{ - build_execution_response_body, build_request_body, collect_response_headers, - decode_response_body_bytes, format_hyper_error_chain, format_upstream_request_error, - format_wreq_upstream_request_error, response_body_is_json, send_request, DirectHttpResponse, - DirectSyncExecutionRuntime, ExecutionRuntimeTransportError, + append_upstream_response_body_chunk, build_execution_response_body, build_request_body, + collect_response_headers, decode_response_body_bytes, execution_response_body_mode, + format_hyper_error_chain, format_upstream_request_error, format_wreq_upstream_request_error, + response_body_is_json, send_request, DirectHttpResponse, DirectSyncExecutionRuntime, + ExecutionRuntimeTransportError, }; use crate::execution_runtime::windsurf::maybe_execute_windsurf_sync; use crate::execution_runtime::{ - analyze_local_candidate_failover_sync, apply_endpoint_response_header_rules, - attach_provider_response_headers_to_report_context, local_failover_response_text, - resolve_core_sync_error_finalize_report_kind, should_fallback_to_control_sync, - should_finalize_sync_response, LocalFailoverDecision, + ai_attempt_retry_scope_from_failure_disposition, analyze_local_candidate_failover_sync, + apply_endpoint_response_header_rules, attach_provider_response_headers_to_report_context, + local_failover_response_text, resolve_core_sync_error_finalize_report_kind, + should_fallback_to_control_sync, should_finalize_sync_response, LocalFailoverDecision, }; use crate::log_ids::short_request_id; use crate::orchestration::{ @@ -113,6 +114,29 @@ struct SyncExecutionFailure { message: String, status_code: Option, latency_ms: Option, + fallback_kind: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum SyncExecutionFailureFallbackKind { + UpstreamResponseTooLarge, + UpstreamResponseDecode, +} + +impl SyncExecutionFailureFallbackKind { + fn error_type(self) -> &'static str { + match self { + Self::UpstreamResponseTooLarge => "upstream_response_too_large", + Self::UpstreamResponseDecode => "upstream_response_decode_failed", + } + } + + fn client_message(self) -> &'static str { + match self { + Self::UpstreamResponseTooLarge => "Upstream response too large", + Self::UpstreamResponseDecode => "Failed to decode upstream response", + } + } } struct SyncAttemptTerminalGuard { @@ -269,11 +293,23 @@ async fn record_sync_attempt_forced_terminal_state( impl SyncExecutionFailure { fn from_transport(err: ExecutionRuntimeTransportError) -> Self { + let fallback_kind = match &err { + ExecutionRuntimeTransportError::UpstreamResponseTooLarge { .. } => { + Some(SyncExecutionFailureFallbackKind::UpstreamResponseTooLarge) + } + ExecutionRuntimeTransportError::UpstreamResponseDecode { .. } => { + Some(SyncExecutionFailureFallbackKind::UpstreamResponseDecode) + } + _ => None, + }; Self { - error_type: "execution_runtime_unavailable", + error_type: fallback_kind + .map(SyncExecutionFailureFallbackKind::error_type) + .unwrap_or("execution_runtime_unavailable"), message: err.to_string(), - status_code: None, + status_code: fallback_kind.map(|_| StatusCode::BAD_GATEWAY.as_u16()), latency_ms: None, + fallback_kind, } } @@ -285,10 +321,94 @@ impl SyncExecutionFailure { ), status_code: Some(StatusCode::GATEWAY_TIMEOUT.as_u16()), latency_ms: Some(elapsed_ms), + fallback_kind: None, } } } +fn build_sync_execution_failure_fallback_body( + client_api_format: &str, + kind: SyncExecutionFailureFallbackKind, +) -> Value { + let message = kind.client_message(); + let error_type = kind.error_type(); + match crate::ai_serving::normalize_api_format_alias(client_api_format).as_str() { + "claude:messages" => json!({ + "type": "error", + "error": { + "type": "upstream_error", + "message": message, + } + }), + "gemini:generate_content" => json!({ + "error": { + "code": StatusCode::BAD_GATEWAY.as_u16(), + "message": message, + "status": "BAD_GATEWAY", + } + }), + _ => json!({ + "error": { + "type": "upstream_error", + "message": message, + "code": error_type, + } + }), + } +} + +fn build_sync_execution_failure_fallback_response( + failure: &SyncExecutionFailure, + plan: &ExecutionPlan, + trace_id: &str, + decision: &GatewayControlDecision, +) -> Result>, GatewayError> { + let Some(kind) = failure.fallback_kind else { + return Ok(None); + }; + let body_json = build_sync_execution_failure_fallback_body(&plan.client_api_format, kind); + let body_bytes = serde_json::to_vec(&body_json) + .map_err(|error| GatewayError::Internal(error.to_string()))?; + let headers = BTreeMap::from([ + ("content-type".to_string(), "application/json".to_string()), + ("content-length".to_string(), body_bytes.len().to_string()), + ]); + let response = build_client_response_from_parts( + StatusCode::BAD_GATEWAY.as_u16(), + &headers, + Body::from(body_bytes), + trace_id, + Some(decision), + )?; + attach_control_metadata_headers( + response, + Some(plan.request_id.as_str()), + plan.candidate_id.as_deref(), + ) + .map(Some) +} + +fn maybe_store_sync_execution_failure_fallback( + failure: &SyncExecutionFailure, + plan: &ExecutionPlan, + trace_id: &str, + decision: &GatewayControlDecision, + retry_scope_out: &mut Option<&mut AiAttemptRetryScope>, + retry_fallback_out: &mut Option<&mut Option>>, +) -> Result<(), GatewayError> { + if failure.fallback_kind.is_none() { + return Ok(()); + } + if let Some(retry_scope) = retry_scope_out.as_deref_mut() { + *retry_scope = AiAttemptRetryScope::Candidate; + } + if let Some(retry_fallback) = retry_fallback_out.as_deref_mut() { + *retry_fallback = + build_sync_execution_failure_fallback_response(failure, plan, trace_id, decision)?; + } + Ok(()) +} + struct ImplicitSyncFinalizeOutcome { payload: GatewaySyncReportRequest, outcome: LocalCoreSyncFinalizeOutcome, @@ -1297,11 +1417,12 @@ async fn execute_openai_image_sync_upstream_sse_candidate( ), ) })?; + append_upstream_response_body_chunk(&mut body_bytes, &chunk) + .map_err(SyncExecutionFailure::from_transport)?; let elapsed_ms = started_at.elapsed().as_millis() as u64; progress .observe_chunk(&chunk, status_code, elapsed_ms) .await; - body_bytes.extend_from_slice(&chunk); } } DirectHttpResponse::HyperH2c(response) => { @@ -1314,11 +1435,12 @@ async fn execute_openai_image_sync_upstream_sse_candidate( )), ) })?; + append_upstream_response_body_chunk(&mut body_bytes, &chunk) + .map_err(SyncExecutionFailure::from_transport)?; let elapsed_ms = started_at.elapsed().as_millis() as u64; progress .observe_chunk(&chunk, status_code, elapsed_ms) .await; - body_bytes.extend_from_slice(&chunk); } } DirectHttpResponse::BrowserWreq(response) => { @@ -1331,24 +1453,30 @@ async fn execute_openai_image_sync_upstream_sse_candidate( ), ) })?; + append_upstream_response_body_chunk(&mut body_bytes, &chunk) + .map_err(SyncExecutionFailure::from_transport)?; let elapsed_ms = started_at.elapsed().as_millis() as u64; progress .observe_chunk(&chunk, status_code, elapsed_ms) .await; - body_bytes.extend_from_slice(&chunk); } } } - let decoded_body_bytes = - decode_response_body_bytes(&headers, &body_bytes).unwrap_or_else(|| body_bytes.clone()); + let decoded_body_bytes = decode_response_body_bytes(&headers, &body_bytes) + .map_err(SyncExecutionFailure::from_transport)?; let elapsed_ms = started_at.elapsed().as_millis() as u64; let upstream_bytes = body_bytes.len() as u64; progress.finish(status_code, elapsed_ms).await; - let body = - build_execution_response_body(&headers, &body_bytes, &decoded_body_bytes, plan.stream) - .map_err(SyncExecutionFailure::from_transport)?; + let body = build_execution_response_body( + &headers, + &body_bytes, + decoded_body_bytes.as_ref(), + plan.stream, + execution_response_body_mode(plan), + ) + .map_err(SyncExecutionFailure::from_transport)?; Ok(ExecutionResult { request_id: plan.request_id.clone(), @@ -1451,6 +1579,8 @@ fn build_openai_image_sync_json_heartbeat_response( report_context, false, Some(progress_snapshot), + None, + None, ) .await, ) @@ -1631,10 +1761,49 @@ pub(crate) async fn execute_execution_runtime_sync( report_context, true, None, + None, + None, ) .await } +#[allow(clippy::too_many_arguments)] +pub(crate) async fn execute_execution_runtime_sync_with_retry_scope( + state: &AppState, + request_path: &str, + plan: ExecutionPlan, + trace_id: &str, + decision: &GatewayControlDecision, + plan_kind: &str, + report_kind: Option, + report_context: Option, +) -> Result>, GatewayError> { + let mut retry_scope = AiAttemptRetryScope::Candidate; + let mut fallback_response = None; + let response = execute_execution_runtime_sync_impl( + state, + request_path, + plan, + trace_id, + decision, + plan_kind, + report_kind, + report_context, + true, + None, + Some(&mut retry_scope), + Some(&mut fallback_response), + ) + .await?; + Ok(match response { + Some(response) => AiAttemptExecutionOutcome::Responded(response), + None => AiAttemptExecutionOutcome::Retry { + scope: retry_scope, + fallback_response, + }, + }) +} + #[allow(clippy::too_many_arguments)] // internal function, grouping would add unnecessary indirection async fn execute_execution_runtime_sync_impl( state: &AppState, @@ -1647,6 +1816,8 @@ async fn execute_execution_runtime_sync_impl( mut report_context: Option, allow_json_heartbeat: bool, progress_snapshot: Option>>, + mut retry_scope_out: Option<&mut AiAttemptRetryScope>, + mut retry_fallback_out: Option<&mut Option>>, ) -> Result>, GatewayError> { if allow_json_heartbeat && should_enable_openai_image_sync_json_heartbeat(plan_kind, &plan, report_context.as_ref()) @@ -1751,6 +1922,14 @@ async fn execute_execution_runtime_sync_impl( { Ok(result) => result, Err(err) => { + maybe_store_sync_execution_failure_fallback( + &err, + &plan, + trace_id, + decision, + &mut retry_scope_out, + &mut retry_fallback_out, + )?; warn!( event_name = "sync_execution_runtime_unavailable", log_type = "ops", @@ -1932,6 +2111,14 @@ async fn execute_execution_runtime_sync_impl( { Ok(result) => result, Err(err) => { + maybe_store_sync_execution_failure_fallback( + &err, + &plan, + trace_id, + decision, + &mut retry_scope_out, + &mut retry_fallback_out, + )?; warn!( event_name = "sync_execution_runtime_unavailable", log_type = "ops", @@ -2256,6 +2443,38 @@ async fn execute_execution_runtime_sync_impl( local_failover_analysis.decision, LocalFailoverDecision::RetryNextCandidate ) { + let failure_disposition = crate::orchestration::classify_failure_disposition( + &plan.provider_api_format, + local_failover_analysis.classification, + result.status_code, + ); + if let Some(retry_scope) = retry_scope_out.as_deref_mut() { + *retry_scope = + ai_attempt_retry_scope_from_failure_disposition(failure_disposition); + } + if failure_disposition.preserve_upstream_error { + if let Some(retry_fallback) = retry_fallback_out.as_deref_mut() { + let mut fallback_headers = headers.clone(); + apply_endpoint_response_header_rules( + state, + &plan, + &mut fallback_headers, + body_json.as_ref(), + ) + .await?; + *retry_fallback = Some(attach_control_metadata_headers( + build_client_response_from_parts( + result.status_code, + &fallback_headers, + Body::from(body_bytes.clone()), + trace_id, + Some(decision), + )?, + Some(plan.request_id.as_str()), + plan.candidate_id.as_deref(), + )?); + } + } let terminal_unix_secs = current_request_candidate_unix_ms(); let error_trace_report_context = with_sync_error_trace_context( report_context.as_ref(), @@ -2909,6 +3128,60 @@ mod tests { .with_execution_runtime_candidate(true) } + #[tokio::test] + async fn oversized_upstream_response_builds_claude_502_retry_fallback() { + let mut plan = test_openai_image_plan(false); + plan.client_api_format = "claude:messages".to_string(); + plan.provider_api_format = "claude:messages".to_string(); + let decision = GatewayControlDecision::synthetic( + "/v1/messages", + Some("ai_public".to_string()), + Some("claude".to_string()), + Some("messages".to_string()), + Some("claude:messages".to_string()), + ) + .with_execution_runtime_candidate(true); + let failure = SyncExecutionFailure::from_transport( + ExecutionRuntimeTransportError::UpstreamResponseTooLarge { + phase: crate::execution_runtime::transport::UpstreamResponseBodyPhase::Wire, + limit_bytes: 8, + }, + ); + + assert_eq!(failure.status_code, Some(StatusCode::BAD_GATEWAY.as_u16())); + assert_eq!( + failure.fallback_kind, + Some(SyncExecutionFailureFallbackKind::UpstreamResponseTooLarge) + ); + let mut retry_scope = AiAttemptRetryScope::Provider; + let mut retry_fallback = None; + { + let mut retry_scope_out = Some(&mut retry_scope); + let mut retry_fallback_out = Some(&mut retry_fallback); + maybe_store_sync_execution_failure_fallback( + &failure, + &plan, + "trace-too-large", + &decision, + &mut retry_scope_out, + &mut retry_fallback_out, + ) + .expect("fallback response should build"); + } + + assert_eq!(retry_scope, AiAttemptRetryScope::Candidate); + let response = retry_fallback.expect("oversized response should provide a fallback"); + + assert_eq!(response.status(), StatusCode::BAD_GATEWAY); + let body = to_bytes(response.into_body(), 1024) + .await + .expect("fallback body should read"); + let body: Value = serde_json::from_slice(&body).expect("fallback body should be json"); + assert_eq!(body["type"], "error"); + assert_eq!(body["error"]["type"], "upstream_error"); + assert_eq!(body["error"]["message"], "Upstream response too large"); + } + fn test_kiro_sync_plan() -> ExecutionPlan { ExecutionPlan { request_id: "req-kiro-sync-cache-1".to_string(), diff --git a/apps/aether-gateway/src/execution_runtime/sync/execution/policy.rs b/apps/aether-gateway/src/execution_runtime/sync/execution/policy.rs index 36368d0a6..f865ba025 100644 --- a/apps/aether-gateway/src/execution_runtime/sync/execution/policy.rs +++ b/apps/aether-gateway/src/execution_runtime/sync/execution/policy.rs @@ -14,8 +14,19 @@ pub(super) fn decode_execution_result_body( let Some(body) = body else { return Ok((Vec::new(), None, None)); }; + let ResponseBody { + json_body, + body_bytes_b64, + } = body; - if let Some(json_body) = body.json_body { + if let Some(body_bytes_b64) = body_bytes_b64 { + let bytes = base64::engine::general_purpose::STANDARD + .decode(&body_bytes_b64) + .map_err(|err| GatewayError::Internal(err.to_string()))?; + return Ok((bytes, json_body, Some(body_bytes_b64))); + } + + if let Some(json_body) = json_body { remove_header_case_insensitive(headers, "content-encoding"); remove_header_case_insensitive(headers, "content-length"); headers @@ -27,13 +38,6 @@ pub(super) fn decode_execution_result_body( return Ok((bytes, Some(json_body), None)); } - if let Some(body_bytes_b64) = body.body_bytes_b64 { - let bytes = base64::engine::general_purpose::STANDARD - .decode(&body_bytes_b64) - .map_err(|err| GatewayError::Internal(err.to_string()))?; - return Ok((bytes, None, Some(body_bytes_b64))); - } - Ok((Vec::new(), None, None)) } @@ -52,6 +56,7 @@ mod tests { use std::collections::BTreeMap; use aether_contracts::ResponseBody; + use base64::Engine as _; use serde_json::json; use super::decode_execution_result_body; @@ -81,4 +86,36 @@ mod tests { Some(body_bytes.len().to_string()) ); } + + #[test] + fn dual_body_prefers_wire_bytes_and_retains_parsed_json() { + let raw = br#"{ "unknown": true, "ok": true }"#; + let encoded = base64::engine::general_purpose::STANDARD.encode(raw); + let raw_len = raw.len().to_string(); + let mut headers = BTreeMap::from([ + ("content-encoding".to_string(), "gzip".to_string()), + ("content-length".to_string(), raw_len.clone()), + ]); + + let (body_bytes, body_json, body_base64) = decode_execution_result_body( + Some(ResponseBody { + json_body: Some(json!({"unknown": true, "ok": true})), + body_bytes_b64: Some(encoded.clone()), + }), + &mut headers, + ) + .expect("body should decode"); + + assert_eq!(body_bytes, raw); + assert_eq!(body_json, Some(json!({"unknown": true, "ok": true}))); + assert_eq!(body_base64.as_deref(), Some(encoded.as_str())); + assert_eq!( + headers.get("content-encoding").map(String::as_str), + Some("gzip") + ); + assert_eq!( + headers.get("content-length").map(String::as_str), + Some(raw_len.as_str()) + ); + } } diff --git a/apps/aether-gateway/src/execution_runtime/sync/mod.rs b/apps/aether-gateway/src/execution_runtime/sync/mod.rs index 72a4ca5e4..d8ab07cbd 100644 --- a/apps/aether-gateway/src/execution_runtime/sync/mod.rs +++ b/apps/aether-gateway/src/execution_runtime/sync/mod.rs @@ -3,6 +3,7 @@ mod execution; pub(crate) use execution::{ build_openai_image_sync_json_whitespace_heartbeat_stream, build_sync_json_whitespace_heartbeat_stream, execute_execution_runtime_sync, + execute_execution_runtime_sync_with_retry_scope, }; #[allow(unused_imports)] diff --git a/apps/aether-gateway/src/execution_runtime/tests.rs b/apps/aether-gateway/src/execution_runtime/tests.rs index 5a862f4ce..8ac659404 100644 --- a/apps/aether-gateway/src/execution_runtime/tests.rs +++ b/apps/aether-gateway/src/execution_runtime/tests.rs @@ -390,7 +390,7 @@ fn build_best_effort_local_core_error_body_converts_sync_errors_across_standard_ "type": "error", "error": { "message": "backend busy", - "type": "api_error", + "type": "overloaded_error", "code": "UNAVAILABLE" } }), diff --git a/apps/aether-gateway/src/execution_runtime/transport.rs b/apps/aether-gateway/src/execution_runtime/transport.rs index 8dc3884a7..3c2107515 100644 --- a/apps/aether-gateway/src/execution_runtime/transport.rs +++ b/apps/aether-gateway/src/execution_runtime/transport.rs @@ -1,3 +1,4 @@ +use std::borrow::Cow; use std::collections::{BTreeMap, HashMap, HashSet, VecDeque}; use std::error::Error as _; use std::future::Future; @@ -8,11 +9,12 @@ use std::sync::{Arc, LazyLock, Mutex as StdMutex, OnceLock, RwLock as StdRwLock} use std::time::{Duration, Instant}; use aether_contracts::{ - ExecutionPlan, ExecutionResult, ExecutionTelemetry, ProxySnapshot, ResolvedTransportProfile, - ResponseBody, EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER, + ExecutionPlan, ExecutionResponseBodyMode, ExecutionResult, ExecutionTelemetry, ProxySnapshot, + ResolvedTransportProfile, ResponseBody, EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER, EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER, EXECUTION_REQUEST_HTTP1_ONLY_HEADER, - TRANSPORT_BACKEND_BROWSER_WREQ, TRANSPORT_BACKEND_REQWEST_RUSTLS, - TRANSPORT_HTTP_MODE_H2C_PRIOR_KNOWLEDGE, TRANSPORT_HTTP_MODE_HTTP1_ONLY, + EXECUTION_RESPONSE_BODY_MODE_HEADER, TRANSPORT_BACKEND_BROWSER_WREQ, + TRANSPORT_BACKEND_REQWEST_RUSTLS, TRANSPORT_HTTP_MODE_H2C_PRIOR_KNOWLEDGE, + TRANSPORT_HTTP_MODE_HTTP1_ONLY, }; use aether_data::repository::proxy_nodes::ProxyNodeTrafficMutation; use aether_http::{apply_http_client_config, HttpClientConfig}; @@ -35,6 +37,7 @@ use reqwest::redirect::Policy; use serde::Serialize; use serde_json::json; use serde_json::Value; +use sha2::Digest as _; use thiserror::Error; use tokio::net::TcpStream; use tokio::sync::OnceCell as TokioOnceCell; @@ -107,6 +110,7 @@ type DirectHyperH2cSenderCacheCell = TokioOnceCell, + pool_partition: Option, connect_timeout_ms: Option, proxy_url: Option, follow_redirects: bool, @@ -444,8 +448,11 @@ pub(crate) fn format_upstream_request_error(err: &reqwest::Error) -> String { } if let Some(url) = err.url() { + let (sanitized_detail, sanitized_url) = + sanitize_upstream_request_error_detail(&detail, url.as_str()); + detail = sanitized_detail; detail.push_str(" [url="); - detail.push_str(url.as_str()); + detail.push_str(&sanitized_url); detail.push(']'); } if !kinds.is_empty() { @@ -457,6 +464,25 @@ pub(crate) fn format_upstream_request_error(err: &reqwest::Error) -> String { detail } +fn sanitize_upstream_request_error_detail(detail: &str, upstream_url: &str) -> (String, String) { + let sanitized_url = sanitize_upstream_url_text(upstream_url); + (detail.replace(upstream_url, &sanitized_url), sanitized_url) +} + +fn sanitize_upstream_url_text(upstream_url: &str) -> String { + if let Ok(mut parsed_url) = reqwest::Url::parse(upstream_url) { + parsed_url.set_query(None); + parsed_url.set_fragment(None); + return parsed_url.to_string(); + } + + let suffix_offset = upstream_url + .char_indices() + .find_map(|(offset, character)| matches!(character, '?' | '#').then_some(offset)) + .unwrap_or(upstream_url.len()); + upstream_url[..suffix_offset].to_string() +} + pub(crate) fn format_wreq_upstream_request_error(err: &wreq::Error) -> String { let mut kinds = Vec::new(); if err.is_connect() { @@ -490,8 +516,12 @@ pub(crate) fn format_wreq_upstream_request_error(err: &wreq::Error) -> String { } if let Some(uri) = err.uri() { + let uri = uri.to_string(); + let (sanitized_detail, sanitized_uri) = + sanitize_upstream_request_error_detail(&detail, &uri); + detail = sanitized_detail; detail.push_str(" [uri="); - detail.push_str(&uri.to_string()); + detail.push_str(&sanitized_uri); detail.push(']'); } if !kinds.is_empty() { @@ -547,12 +577,60 @@ pub(crate) enum ExecutionRuntimeTransportError { BrowserBody(String), #[error("failed to execute upstream request: {0}")] UpstreamRequest(String), + #[error("upstream response {phase} body exceeds {limit_bytes} bytes")] + UpstreamResponseTooLarge { + phase: UpstreamResponseBodyPhase, + limit_bytes: usize, + }, + #[error("failed to decode upstream response body with content-encoding {encoding}: {message}")] + UpstreamResponseDecode { encoding: String, message: String }, #[error("hub relay request failed: {0}")] RelayError(String), #[error("upstream response is not valid JSON: {0}")] InvalidJson(serde_json::Error), } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum UpstreamResponseBodyPhase { + Wire, + Decoded, +} + +impl std::fmt::Display for UpstreamResponseBodyPhase { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str(match self { + Self::Wire => "wire", + Self::Decoded => "decoded", + }) + } +} + +pub(crate) fn append_upstream_response_body_chunk( + body: &mut Vec, + chunk: &[u8], +) -> Result<(), ExecutionRuntimeTransportError> { + append_upstream_response_body_chunk_with_limit( + body, + chunk, + crate::headers::max_internal_buffered_body_bytes(), + ) +} + +fn append_upstream_response_body_chunk_with_limit( + body: &mut Vec, + chunk: &[u8], + limit_bytes: usize, +) -> Result<(), ExecutionRuntimeTransportError> { + if body.len() > limit_bytes || chunk.len() > limit_bytes.saturating_sub(body.len()) { + return Err(ExecutionRuntimeTransportError::UpstreamResponseTooLarge { + phase: UpstreamResponseBodyPhase::Wire, + limit_bytes, + }); + } + body.extend_from_slice(chunk); + Ok(()) +} + #[derive(Debug, Serialize)] struct RelayRequestMeta { provider_id: String, @@ -608,6 +686,7 @@ pub(crate) struct DirectUpstreamStreamExecution { pub(crate) provider_api_format: String, pub(crate) stream_summary_report_context: Value, pub(crate) prefetched_body: VecDeque>, + pub(crate) stream_precommit_committed: bool, pub(crate) response: DirectUpstreamResponse, pub(crate) started_at: Instant, pub(crate) stream_first_byte_timeout: Option, @@ -654,16 +733,16 @@ impl DirectSyncExecutionRuntime { }); let (body_bytes, stream_ttfb_ms) = response.bytes_with_stream_timeout(plan, started_at).await?; - let decoded_body_bytes = decode_response_body_bytes(&headers, &body_bytes) - .unwrap_or_else(|| body_bytes.to_vec()); + let decoded_body_bytes = decode_response_body_bytes(&headers, &body_bytes)?; let elapsed_ms = started_at.elapsed().as_millis() as u64; let upstream_bytes = body_bytes.len() as u64; let body = build_execution_response_body( &headers, &body_bytes, - &decoded_body_bytes, + decoded_body_bytes.as_ref(), plan.stream, + execution_response_body_mode(plan), )?; Ok(ExecutionResult { @@ -713,6 +792,7 @@ impl DirectSyncExecutionRuntime { provider_api_format: plan.provider_api_format.clone(), stream_summary_report_context, prefetched_body: VecDeque::new(), + stream_precommit_committed: false, response: response.into_direct_upstream_response(), started_at, stream_first_byte_timeout: resolve_stream_first_byte_timeout(plan), @@ -827,6 +907,7 @@ pub(crate) async fn execute_stream_plan_via_local_tunnel( provider_api_format: plan.provider_api_format.clone(), stream_summary_report_context: build_stream_summary_report_context(plan), prefetched_body: VecDeque::new(), + stream_precommit_committed: false, response: DirectUpstreamResponse::LocalTunnel(response), started_at, stream_first_byte_timeout: resolve_stream_first_byte_timeout(plan), @@ -962,8 +1043,7 @@ async fn execute_sync_plan_via_local_tunnel_inner( let proxy_timing = execution_header_for_log(&headers, "x-proxy-timing").unwrap_or("-"); let (body_bytes, stream_ttfb_ms) = collect_local_tunnel_response_body(response, plan, started_at).await?; - let decoded_body_bytes = - decode_response_body_bytes(&headers, &body_bytes).unwrap_or_else(|| body_bytes.clone()); + let decoded_body_bytes = decode_response_body_bytes(&headers, &body_bytes)?; let elapsed_ms = started_at.elapsed().as_millis() as u64; let upstream_bytes = body_bytes.len() as u64; if status_code >= 400 { @@ -1000,8 +1080,13 @@ async fn execute_sync_plan_via_local_tunnel_inner( ); } - let body = - build_execution_response_body(&headers, &body_bytes, &decoded_body_bytes, plan.stream)?; + let body = build_execution_response_body( + &headers, + &body_bytes, + decoded_body_bytes.as_ref(), + plan.stream, + execution_response_body_mode(plan), + )?; Ok(ExecutionResult { request_id: plan.request_id.clone(), @@ -1044,7 +1129,7 @@ async fn collect_local_tunnel_response_body( if plan.stream && first_byte_ms.is_none() && !chunk.is_empty() { first_byte_ms = Some(started_at.elapsed().as_millis() as u64); } - body_bytes.extend_from_slice(&chunk); + append_upstream_response_body_chunk(&mut body_bytes, &chunk)?; } Ok((body_bytes, first_byte_ms)) @@ -1154,6 +1239,7 @@ async fn send_request_inner( let client_select_started_at = Instant::now(); let client = build_client( &plan.url, + &plan.key_id, plan.timeouts.as_ref(), plan.proxy.as_ref(), plan.transport_profile.as_ref(), @@ -1204,23 +1290,23 @@ impl DirectHttpResponse { } pub(crate) async fn bytes(self) -> Result { + let started_at = Instant::now(); match self { - DirectHttpResponse::Reqwest(response) => response.bytes().await.map_err(|err| { - ExecutionRuntimeTransportError::UpstreamRequest(format_upstream_request_error(&err)) - }), - DirectHttpResponse::HyperH2c(response) => response - .into_body() - .collect() - .await - .map(|collected| collected.to_bytes()) - .map_err(|err| { - ExecutionRuntimeTransportError::UpstreamRequest(format_hyper_error_chain(&err)) - }), - DirectHttpResponse::BrowserWreq(response) => response.bytes().await.map_err(|err| { - ExecutionRuntimeTransportError::BrowserBody(format_wreq_upstream_request_error( - &err, - )) - }), + DirectHttpResponse::Reqwest(response) => { + collect_reqwest_stream_body(response, started_at, None) + .await + .map(|(body, _)| body) + } + DirectHttpResponse::HyperH2c(response) => { + collect_hyper_stream_body(response, started_at, None) + .await + .map(|(body, _)| body) + } + DirectHttpResponse::BrowserWreq(response) => { + collect_wreq_stream_body(response, started_at, None) + .await + .map(|(body, _)| body) + } } } @@ -1308,7 +1394,7 @@ async fn collect_reqwest_stream_body( if first_byte_ms.is_none() && !chunk.is_empty() { first_byte_ms = Some(started_at.elapsed().as_millis() as u64); } - body_bytes.extend_from_slice(&chunk); + append_upstream_response_body_chunk(&mut body_bytes, &chunk)?; } Ok((Bytes::from(body_bytes), first_byte_ms)) @@ -1338,7 +1424,7 @@ async fn collect_hyper_stream_body( if first_byte_ms.is_none() && !chunk.is_empty() { first_byte_ms = Some(started_at.elapsed().as_millis() as u64); } - body_bytes.extend_from_slice(&chunk); + append_upstream_response_body_chunk(&mut body_bytes, &chunk)?; } Ok((Bytes::from(body_bytes), first_byte_ms)) @@ -1368,7 +1454,7 @@ async fn collect_wreq_stream_body( if first_byte_ms.is_none() && !chunk.is_empty() { first_byte_ms = Some(started_at.elapsed().as_millis() as u64); } - body_bytes.extend_from_slice(&chunk); + append_upstream_response_body_chunk(&mut body_bytes, &chunk)?; } Ok((Bytes::from(body_bytes), first_byte_ms)) @@ -2555,6 +2641,7 @@ fn resolve_local_tunnel_node_id(state: &AppState, proxy: Option<&ProxySnapshot>) fn build_client( request_url: &str, + key_id: &str, timeouts: Option<&aether_contracts::ExecutionTimeouts>, proxy: Option<&ProxySnapshot>, transport_profile: Option<&ResolvedTransportProfile>, @@ -2564,6 +2651,7 @@ fn build_client( let resolved_proxy_url = resolve_proxy_url(proxy)?; let cache_key = direct_reqwest_client_cache_key( request_url, + key_id, timeouts, resolved_proxy_url, transport_profile, @@ -2610,7 +2698,10 @@ pub(crate) fn prewarm_direct_reqwest_client_cache_for_plan(plan: &ExecutionPlan) candidate_id = ?plan.candidate_id, provider_id = %plan.provider_id, endpoint_id = %plan.endpoint_id, - key_id = %plan.key_id, + key_partition = ?direct_reqwest_pool_partition( + plan.transport_profile.as_ref(), + &plan.key_id, + ), "gateway direct reqwest client prewarm skipped" ); } @@ -2638,6 +2729,7 @@ fn try_prewarm_direct_reqwest_client_cache_for_plan( let resolved_proxy_url = resolve_proxy_url(plan.proxy.as_ref())?; let cache_key = direct_reqwest_client_cache_key( &plan.url, + &plan.key_id, plan.timeouts.as_ref(), resolved_proxy_url, plan.transport_profile.as_ref(), @@ -2877,6 +2969,7 @@ fn mark_direct_reqwest_client_cache_not_warming(cache_key: &DirectReqwestClientC fn direct_reqwest_client_cache_key( request_url: &str, + key_id: &str, timeouts: Option<&aether_contracts::ExecutionTimeouts>, proxy_url: Option, transport_profile: Option<&ResolvedTransportProfile>, @@ -2886,6 +2979,7 @@ fn direct_reqwest_client_cache_key( upstream_origin: direct_reqwest_cache_per_origin() .then(|| direct_reqwest_upstream_origin(request_url)) .flatten(), + pool_partition: direct_reqwest_pool_partition(transport_profile, key_id), connect_timeout_ms: timeouts.and_then(|timeouts| timeouts.connect_ms), proxy_url, follow_redirects: transport_controls.follow_redirects == Some(true), @@ -2895,6 +2989,17 @@ fn direct_reqwest_client_cache_key( } } +fn direct_reqwest_pool_partition( + transport_profile: Option<&ResolvedTransportProfile>, + key_id: &str, +) -> Option { + let key_id = key_id.trim(); + transport_profile + .filter(|profile| profile.pool_scope.trim().eq_ignore_ascii_case("key")) + .filter(|_| !key_id.is_empty()) + .map(|_| format!("{:x}", sha2::Sha256::digest(key_id.as_bytes()))) +} + fn direct_reqwest_cache_per_origin() -> bool { std::env::var(DIRECT_REQWEST_CACHE_PER_ORIGIN_ENV) .ok() @@ -3719,11 +3824,13 @@ pub(crate) fn build_request_headers( } for (key, value) in headers { let normalized_key = key.trim().to_ascii_lowercase(); - if is_hop_by_hop_header(&normalized_key) + if crate::headers::should_skip_request_header(&normalized_key) + || is_hop_by_hop_header(&normalized_key) || normalized_key == "content-encoding" || normalized_key == EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER || normalized_key == EXECUTION_REQUEST_HTTP1_ONLY_HEADER || normalized_key == EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER + || normalized_key == EXECUTION_RESPONSE_BODY_MODE_HEADER { continue; } @@ -3766,6 +3873,23 @@ fn resolve_execution_transport_controls( } } +pub(crate) fn execution_response_body_mode(plan: &ExecutionPlan) -> ExecutionResponseBodyMode { + if plan.stream + || plan.body.body_bytes_b64.is_none() + || !plan + .client_api_format + .trim() + .eq_ignore_ascii_case(plan.provider_api_format.trim()) + { + return ExecutionResponseBodyMode::StructuredJson; + } + + ExecutionResponseBodyMode::from_header_value(execution_transport_header_value( + &plan.headers, + EXECUTION_RESPONSE_BODY_MODE_HEADER, + )) +} + fn execution_transport_header_value<'a>( headers: &'a BTreeMap, target: &str, @@ -3844,10 +3968,22 @@ fn execution_log_url_host(url: &str) -> String { .unwrap_or_else(|| "-".to_string()) } -pub(crate) fn decode_response_body_bytes( +pub(crate) fn decode_response_body_bytes<'a>( headers: &BTreeMap, - body_bytes: &[u8], -) -> Option> { + body_bytes: &'a [u8], +) -> Result, ExecutionRuntimeTransportError> { + decode_response_body_bytes_with_limit( + headers, + body_bytes, + crate::headers::max_internal_buffered_body_bytes(), + ) +} + +fn decode_response_body_bytes_with_limit<'a>( + headers: &BTreeMap, + body_bytes: &'a [u8], + limit_bytes: usize, +) -> Result, ExecutionRuntimeTransportError> { let encoding = headers .get("content-encoding") .map(String::as_str) @@ -3857,20 +3993,43 @@ pub(crate) fn decode_response_body_bytes( match encoding.as_deref() { Some("gzip") => { let mut decoder = GzDecoder::new(body_bytes); - let mut out = Vec::new(); - decoder.read_to_end(&mut out).ok()?; - Some(out) + read_upstream_response_decoder_with_limit("gzip", &mut decoder, limit_bytes) + .map(Cow::Owned) } Some("deflate") => { let mut decoder = DeflateDecoder::new(body_bytes); - let mut out = Vec::new(); - decoder.read_to_end(&mut out).ok()?; - Some(out) + read_upstream_response_decoder_with_limit("deflate", &mut decoder, limit_bytes) + .map(Cow::Owned) } - _ => None, + _ => Ok(Cow::Borrowed(body_bytes)), } } +fn read_upstream_response_decoder_with_limit( + encoding: &str, + decoder: &mut impl Read, + limit_bytes: usize, +) -> Result, ExecutionRuntimeTransportError> { + let read_limit = u64::try_from(limit_bytes) + .unwrap_or(u64::MAX) + .saturating_add(1); + let mut limited = decoder.take(read_limit); + let mut out = Vec::new(); + limited.read_to_end(&mut out).map_err(|error| { + ExecutionRuntimeTransportError::UpstreamResponseDecode { + encoding: encoding.to_string(), + message: error.to_string(), + } + })?; + if out.len() > limit_bytes { + return Err(ExecutionRuntimeTransportError::UpstreamResponseTooLarge { + phase: UpstreamResponseBodyPhase::Decoded, + limit_bytes, + }); + } + Ok(out) +} + pub(crate) fn response_body_is_json(headers: &BTreeMap, body_bytes: &[u8]) -> bool { let content_type = headers .get("content-type") @@ -3893,6 +4052,7 @@ pub(crate) fn build_execution_response_body( body_bytes: &[u8], decoded_body_bytes: &[u8], stream: bool, + response_body_mode: ExecutionResponseBodyMode, ) -> Result, ExecutionRuntimeTransportError> { if body_bytes.is_empty() { return Ok(None); @@ -3903,7 +4063,8 @@ pub(crate) fn build_execution_response_body( .map_err(ExecutionRuntimeTransportError::InvalidJson)?; return Ok(Some(ResponseBody { json_body: Some(body_json), - body_bytes_b64: None, + body_bytes_b64: (response_body_mode == ExecutionResponseBodyMode::PreserveBytes) + .then(|| base64::engine::general_purpose::STANDARD.encode(body_bytes)), })); } @@ -3932,12 +4093,13 @@ pub(crate) fn build_execution_response_body( #[cfg(test)] mod tests { use std::collections::BTreeMap; - use std::io::Read; + use std::io::{Read, Write}; use std::sync::{Arc, Mutex, MutexGuard, OnceLock}; use aether_contracts::{ - ExecutionPlan, ExecutionTimeouts, ProxySnapshot, RequestBody, ResolvedTransportProfile, - EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER, EXECUTION_REQUEST_HTTP1_ONLY_HEADER, + ExecutionPlan, ExecutionResponseBodyMode, ExecutionTimeouts, ProxySnapshot, RequestBody, + ResolvedTransportProfile, EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER, + EXECUTION_REQUEST_HTTP1_ONLY_HEADER, EXECUTION_RESPONSE_BODY_MODE_HEADER, TRANSPORT_BACKEND_BROWSER_WREQ, TRANSPORT_BACKEND_REQWEST_RUSTLS, TRANSPORT_HTTP_MODE_AUTO, TRANSPORT_HTTP_MODE_H2C_PRIOR_KNOWLEDGE, TRANSPORT_HTTP_MODE_HTTP1_ONLY, }; @@ -3956,13 +4118,14 @@ mod tests { use tokio::sync::watch; use super::{ - build_browser_wreq_client, build_client, build_direct_tunnel_request_meta, - build_execution_response_body, build_request_headers, execute_sync_plan, + append_upstream_response_body_chunk_with_limit, build_browser_wreq_client, build_client, + build_direct_tunnel_request_meta, build_execution_response_body, build_request_headers, + decode_response_body_bytes_with_limit, execute_sync_plan, execution_response_body_mode, record_manual_proxy_request_failure, record_manual_proxy_request_outcome, record_manual_proxy_request_success, record_manual_proxy_stream_error, resolve_execution_transport_controls, resolve_non_stream_total_timeout, resolve_stream_first_byte_timeout, response_body_is_json, DirectSyncExecutionRuntime, - ExecutionRuntimeTransportError, ExecutionTransportControls, + ExecutionRuntimeTransportError, ExecutionTransportControls, UpstreamResponseBodyPhase, }; use crate::constants::{ EXECUTION_RUNTIME_LOOP_GUARD_HEADER, EXECUTION_RUNTIME_LOOP_GUARD_VIA_TOKEN, @@ -3976,6 +4139,106 @@ mod tests { const LOCAL_HTTP_SUCCESS_TIMEOUT_MS: u64 = 15_000; + #[test] + fn upstream_error_url_sanitization_removes_secrets_everywhere() { + let upstream_url = + "https://api.example.test/v1/messages?key=query-secret&alt=sse#fragment-secret"; + let detail = format!( + "error sending request for url ({upstream_url}); source repeated {upstream_url}" + ); + + let (sanitized_detail, sanitized_url) = + super::sanitize_upstream_request_error_detail(&detail, upstream_url); + + assert_eq!(sanitized_url, "https://api.example.test/v1/messages"); + assert_eq!( + sanitized_detail, + "error sending request for url (https://api.example.test/v1/messages); source repeated https://api.example.test/v1/messages" + ); + assert!(!sanitized_detail.contains("query-secret")); + assert!(!sanitized_detail.contains("fragment-secret")); + } + + #[test] + fn request_header_materialization_strips_all_aether_internal_headers() { + let headers = BTreeMap::from([ + ("authorization".to_string(), "Bearer upstream".to_string()), + ("x-aether-grok-runtime".to_string(), "1".to_string()), + ("x-aether-future-control".to_string(), "private".to_string()), + ]); + + let materialized = build_request_headers(&headers, None, false) + .expect("provider request headers should materialize"); + + assert_eq!( + materialized + .get("authorization") + .and_then(|value| value.to_str().ok()), + Some("Bearer upstream") + ); + assert!(!materialized.contains_key("x-aether-grok-runtime")); + assert!(!materialized.contains_key("x-aether-future-control")); + } + + #[test] + fn upstream_response_wire_limit_allows_exact_body_and_rejects_next_byte() { + let mut body = Vec::new(); + append_upstream_response_body_chunk_with_limit(&mut body, b"1234", 5) + .expect("chunk below limit should append"); + append_upstream_response_body_chunk_with_limit(&mut body, b"5", 5) + .expect("body exactly at limit should append"); + + let error = append_upstream_response_body_chunk_with_limit(&mut body, b"6", 5) + .expect_err("body above limit should fail"); + + assert_eq!(body, b"12345"); + assert!(matches!( + error, + ExecutionRuntimeTransportError::UpstreamResponseTooLarge { + phase: UpstreamResponseBodyPhase::Wire, + limit_bytes: 5, + } + )); + } + + #[test] + fn upstream_response_gzip_decode_limit_rejects_decompression_bomb() { + let payload = vec![b'x'; 9]; + let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default()); + encoder + .write_all(&payload) + .expect("gzip payload should encode"); + let encoded = encoder.finish().expect("gzip payload should finish"); + let headers = BTreeMap::from([("content-encoding".to_string(), "gzip".to_string())]); + + let error = decode_response_body_bytes_with_limit(&headers, &encoded, 8) + .expect_err("decoded body above limit should fail"); + + assert!(matches!( + error, + ExecutionRuntimeTransportError::UpstreamResponseTooLarge { + phase: UpstreamResponseBodyPhase::Decoded, + limit_bytes: 8, + } + )); + } + + #[test] + fn upstream_response_gzip_decode_limit_allows_exact_body() { + let payload = b"12345678"; + let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default()); + encoder + .write_all(payload) + .expect("gzip payload should encode"); + let encoded = encoder.finish().expect("gzip payload should finish"); + let headers = BTreeMap::from([("content-encoding".to_string(), "gzip".to_string())]); + + let decoded = decode_response_body_bytes_with_limit(&headers, &encoded, payload.len()) + .expect("decoded body exactly at limit should pass"); + + assert_eq!(decoded.as_ref(), payload); + } + #[test] fn gateway_frontdoor_self_loop_guard_matches_loopback_public_ai_route() { assert!(gateway_frontdoor_self_loop_guard_matches_with_port( @@ -4034,6 +4297,7 @@ mod tests { for proxy_url in ["socks5://127.0.0.1:1080", "socks5h://127.0.0.1:1080"] { build_client( "https://api.example.test/v1/chat/completions", + "key-test", Some(&timeouts), Some(&aether_contracts::ProxySnapshot { enabled: Some(true), @@ -4103,6 +4367,7 @@ mod tests { let left = super::direct_reqwest_client_cache_key( "http://127.0.0.1:18184/v1/chat/completions", + "key-1", Some(&timeouts), None, Some(&h2c_profile), @@ -4110,6 +4375,7 @@ mod tests { ); let right = super::direct_reqwest_client_cache_key( "http://127.0.0.1:18184/v1/responses", + "key-1", Some(&timeouts), None, Some(&same_h2c_profile), @@ -4117,6 +4383,7 @@ mod tests { ); let different_mode = super::direct_reqwest_client_cache_key( "http://127.0.0.1:18184/v1/chat/completions", + "key-1", Some(&timeouts), None, Some(&http1_profile), @@ -4124,6 +4391,7 @@ mod tests { ); let different_proxy = super::direct_reqwest_client_cache_key( "http://127.0.0.1:18184/v1/chat/completions", + "key-1", Some(&timeouts), Some("http://127.0.0.1:8080".into()), Some(&h2c_profile), @@ -4138,6 +4406,72 @@ mod tests { )); } + #[test] + fn direct_reqwest_client_cache_key_partitions_key_scoped_pools_by_hashed_key_id() { + let profile = ResolvedTransportProfile { + profile_id: "key-scoped-profile".into(), + backend: TRANSPORT_BACKEND_REQWEST_RUSTLS.into(), + http_mode: TRANSPORT_HTTP_MODE_AUTO.into(), + pool_scope: " key ".into(), + header_fingerprint: None, + extra: None, + }; + let first_key_id = "plain-key-identity-alpha"; + let second_key_id = "plain-key-identity-beta"; + let cache_key = |key_id| { + super::direct_reqwest_client_cache_key( + "https://api.example.test/v1/messages", + key_id, + None, + None, + Some(&profile), + ExecutionTransportControls::default(), + ) + }; + + let first = cache_key(first_key_id); + let first_key_id_with_whitespace = format!(" {first_key_id} "); + let first_with_whitespace = cache_key(&first_key_id_with_whitespace); + let second = cache_key(second_key_id); + let empty = cache_key(" "); + + assert_eq!(first, first_with_whitespace); + assert_ne!(first, second); + assert_eq!(first.pool_partition.as_deref().map(str::len), Some(64)); + assert!(empty.pool_partition.is_none()); + let debug = format!("{first:?} {second:?}"); + assert!(!debug.contains(first_key_id)); + assert!(!debug.contains(second_key_id)); + } + + #[test] + fn direct_reqwest_client_cache_key_shares_non_key_scoped_pools() { + let profile = ResolvedTransportProfile { + profile_id: "provider-scoped-profile".into(), + backend: TRANSPORT_BACKEND_REQWEST_RUSTLS.into(), + http_mode: TRANSPORT_HTTP_MODE_AUTO.into(), + pool_scope: "provider".into(), + header_fingerprint: None, + extra: None, + }; + let cache_key = |key_id| { + super::direct_reqwest_client_cache_key( + "https://api.example.test/v1/messages", + key_id, + None, + None, + Some(&profile), + ExecutionTransportControls::default(), + ) + }; + + let first = cache_key("plain-key-identity-alpha"); + let second = cache_key("plain-key-identity-beta"); + + assert_eq!(first, second); + assert!(first.pool_partition.is_none()); + } + #[test] fn direct_reqwest_client_cache_key_splits_origin_only_when_enabled() { let _guard = direct_reqwest_env_lock(); @@ -4152,6 +4486,7 @@ mod tests { let shared_left = super::direct_reqwest_client_cache_key( "http://127.0.0.1:18184/v1/chat/completions", + "key-1", None, None, Some(&profile), @@ -4159,6 +4494,7 @@ mod tests { ); let shared_right = super::direct_reqwest_client_cache_key( "http://127.0.0.1:18185/v1/chat/completions", + "key-1", None, None, Some(&profile), @@ -4169,6 +4505,7 @@ mod tests { let _per_origin = set_test_env_var(super::DIRECT_REQWEST_CACHE_PER_ORIGIN_ENV, "true"); let split_left = super::direct_reqwest_client_cache_key( "http://127.0.0.1:18184/v1/chat/completions", + "key-1", None, None, Some(&profile), @@ -4176,6 +4513,7 @@ mod tests { ); let split_right = super::direct_reqwest_client_cache_key( "http://127.0.0.1:18185/v1/chat/completions", + "key-1", None, None, Some(&profile), @@ -4201,6 +4539,7 @@ mod tests { let auto_key = super::direct_reqwest_client_cache_key( "http://127.0.0.1:18184/v1/chat/completions", + "key-1", None, None, Some(&auto_profile), @@ -4208,6 +4547,7 @@ mod tests { ); let h2c_key = super::direct_reqwest_client_cache_key( "http://127.0.0.1:18184/v1/chat/completions", + "key-1", None, None, Some(&h2c_profile), @@ -4547,6 +4887,7 @@ mod tests { let cache_key = super::direct_reqwest_client_cache_key( &plan.url, + &plan.key_id, plan.timeouts.as_ref(), None, Some(&profile), @@ -4610,6 +4951,7 @@ mod tests { let cache_key = super::direct_reqwest_client_cache_key( &plan.url, + &plan.key_id, plan.timeouts.as_ref(), None, Some(&profile), @@ -4664,6 +5006,7 @@ mod tests { let cache_key = super::direct_reqwest_client_cache_key( &plan.url, + &plan.key_id, plan.timeouts.as_ref(), None, Some(&profile), @@ -4782,6 +5125,57 @@ mod tests { .is_none()); } + #[test] + fn response_body_mode_control_header_is_never_forwarded_upstream() { + let headers = BTreeMap::from([ + ("content-type".into(), "application/json".into()), + ( + EXECUTION_RESPONSE_BODY_MODE_HEADER.into(), + ExecutionResponseBodyMode::PreserveBytes + .as_str() + .to_string(), + ), + ]); + + let forwarded = build_request_headers(&headers, None, true) + .expect("headers should build after stripping internal controls"); + + assert!(forwarded.get("content-type").is_some()); + assert!(forwarded.get(EXECUTION_RESPONSE_BODY_MODE_HEADER).is_none()); + } + + #[test] + fn response_body_mode_requires_same_format_raw_sync_plan() { + let mut plan = tunnel_timeout_plan(false); + plan.headers.insert( + EXECUTION_RESPONSE_BODY_MODE_HEADER.to_string(), + ExecutionResponseBodyMode::PreserveBytes + .as_str() + .to_string(), + ); + + assert_eq!( + execution_response_body_mode(&plan), + ExecutionResponseBodyMode::StructuredJson + ); + + plan.body = RequestBody { + json_body: None, + body_bytes_b64: Some("e30=".to_string()), + body_ref: None, + }; + assert_eq!( + execution_response_body_mode(&plan), + ExecutionResponseBodyMode::PreserveBytes + ); + + plan.provider_api_format = "claude:messages".to_string(); + assert_eq!( + execution_response_body_mode(&plan), + ExecutionResponseBodyMode::StructuredJson + ); + } + #[test] fn tunnel_request_meta_uses_total_timeout_for_non_stream_requests() { let plan = tunnel_timeout_plan(false); @@ -6465,6 +6859,7 @@ mod tests { ); let cache_key = super::direct_reqwest_client_cache_key( &plan.url, + &plan.key_id, plan.timeouts.as_ref(), None, Some(&profile), @@ -6646,6 +7041,7 @@ mod tests { let error = match build_client( "https://api.example.test/v1/chat/completions", + "key-test", None, None, Some(&profile), @@ -6673,6 +7069,51 @@ mod tests { assert!(!response_body_is_json(&headers, &body)); } + #[test] + fn structured_json_response_does_not_duplicate_body_bytes() { + let headers = + BTreeMap::from([("content-type".to_string(), "application/json".to_string())]); + let body_bytes = br#"{ "unknown": true, "ok": true }"#; + + let body = build_execution_response_body( + &headers, + body_bytes, + body_bytes, + false, + ExecutionResponseBodyMode::StructuredJson, + ) + .expect("body should build") + .expect("body should be present"); + + assert!(body.json_body.is_some()); + assert!(body.body_bytes_b64.is_none()); + } + + #[test] + fn preserve_bytes_json_response_keeps_parsed_and_wire_representations() { + let headers = + BTreeMap::from([("content-type".to_string(), "application/json".to_string())]); + let body_bytes = br#"{ "unknown": true, "ok": true }"#; + + let body = build_execution_response_body( + &headers, + body_bytes, + body_bytes, + false, + ExecutionResponseBodyMode::PreserveBytes, + ) + .expect("body should build") + .expect("body should be present"); + + assert_eq!(body.json_body, Some(json!({"unknown": true, "ok": true}))); + assert_eq!( + base64::engine::general_purpose::STANDARD + .decode(body.body_bytes_b64.expect("wire bytes should be present")) + .expect("wire body should decode"), + body_bytes + ); + } + #[test] fn connect_json_error_response_is_decoded_for_stream_sync_body() { let headers = BTreeMap::from([( @@ -6684,9 +7125,15 @@ mod tests { body_bytes.extend_from_slice(&(payload.len() as u32).to_be_bytes()); body_bytes.extend_from_slice(payload); - let body = build_execution_response_body(&headers, &body_bytes, &body_bytes, true) - .expect("body should build") - .expect("body should be present"); + let body = build_execution_response_body( + &headers, + &body_bytes, + &body_bytes, + true, + ExecutionResponseBodyMode::StructuredJson, + ) + .expect("body should build") + .expect("body should be present"); assert_eq!( body.json_body diff --git a/apps/aether-gateway/src/executor/candidate_loop.rs b/apps/aether-gateway/src/executor/candidate_loop.rs index 9447928e9..4b0628cfe 100644 --- a/apps/aether-gateway/src/executor/candidate_loop.rs +++ b/apps/aether-gateway/src/executor/candidate_loop.rs @@ -1,7 +1,8 @@ use std::collections::{BTreeMap, BTreeSet}; use aether_ai_serving::{ - run_ai_attempt_loop, AiAttemptLoopOutcome, AiAttemptLoopPort, AiExecutionAttempt, + run_ai_attempt_loop, AiAttemptExecutionOutcome, AiAttemptLoopOutcome, AiAttemptLoopPort, + AiAttemptRetryScope, AiExecutionAttempt, }; use aether_data_contracts::repository::candidates::RequestCandidateStatus; use aether_runtime::ConcurrencyPermit; @@ -18,8 +19,13 @@ use tracing::{debug, warn, Instrument}; use crate::ai_serving::LocalExecutionAttemptSource; use crate::clock::current_unix_ms; use crate::control::GatewayControlDecision; -use crate::execution_runtime::{execute_execution_runtime_stream, execute_execution_runtime_sync}; -use crate::executor::{build_local_execution_exhaustion, LocalExecutionRequestOutcome}; +use crate::execution_runtime::{ + execute_execution_runtime_stream_with_retry_scope, + execute_execution_runtime_sync_with_retry_scope, +}; +use crate::executor::{ + build_local_execution_exhaustion, mark_deferred_upstream_response, LocalExecutionRequestOutcome, +}; use crate::handlers::shared::provider_pool::release_admin_provider_pool_key_lease; use crate::log_ids::short_request_id; use crate::orchestration::{ @@ -124,6 +130,9 @@ where AiAttemptLoopOutcome::Responded(response) => { Ok(LocalExecutionRequestOutcome::responded(response)) } + AiAttemptLoopOutcome::Deferred(response) => Ok( + LocalExecutionRequestOutcome::responded(mark_deferred_upstream_response(response)), + ), AiAttemptLoopOutcome::Exhausted(exhaustion) => { Ok(LocalExecutionRequestOutcome::Exhausted(exhaustion)) } @@ -251,7 +260,10 @@ where Ok(()) } - async fn execute_attempt(&self, attempt: &T) -> Result, Self::Error> { + async fn execute_attempt( + &self, + attempt: &T, + ) -> Result, Self::Error> { let plan = attempt.execution_plan(); let report_context = attempt.report_context(); if let Some(response) = execution_plan_balance_capacity_response( @@ -263,12 +275,12 @@ where ) .await? { - return Ok(Some(response)); + return Ok(AiAttemptExecutionOutcome::Responded(response)); } prewarm_direct_reqwest_candidate_client(plan); let _permit = acquire_upstream_execution_gate(self.state, self.trace_id).await?; let upstream_execution_gate_held_started_at = std::time::Instant::now(); - let mut response = execute_execution_runtime_sync( + let mut execution = execute_execution_runtime_sync_with_retry_scope( self.state, self.parts.uri.path(), plan.clone(), @@ -285,10 +297,18 @@ where .elapsed() .as_millis() as u64, ); - if let Some(response) = response.as_mut() { - attach_redaction_execution_candidate(response, plan.candidate_id.as_deref()); + match &mut execution { + AiAttemptExecutionOutcome::Responded(response) + | AiAttemptExecutionOutcome::Retry { + fallback_response: Some(response), + .. + } => attach_redaction_execution_candidate(response, plan.candidate_id.as_deref()), + AiAttemptExecutionOutcome::Retry { + fallback_response: None, + .. + } => {} } - Ok(response) + Ok(execution) } async fn mark_unused_attempts(&self, attempts: Vec) -> Result<(), Self::Error> { @@ -389,6 +409,9 @@ where AiAttemptLoopOutcome::Responded(response) => { Ok(LocalExecutionRequestOutcome::responded(response)) } + AiAttemptLoopOutcome::Deferred(response) => Ok( + LocalExecutionRequestOutcome::responded(mark_deferred_upstream_response(response)), + ), AiAttemptLoopOutcome::Exhausted(exhaustion) => { Ok(LocalExecutionRequestOutcome::Exhausted(exhaustion)) } @@ -760,6 +783,7 @@ where Attempt: AiExecutionAttempt + Send + Sync + 'static, { let mut last_attempted = None; + let mut fallback_response = None; loop { let next_started_at = std::time::Instant::now(); @@ -781,8 +805,8 @@ where } port.record_attempt_started(&attempt).await?; let execute_started_at = std::time::Instant::now(); - let response = match port.execute_attempt(&attempt).await { - Ok(response) => response, + let execution = match port.execute_attempt(&attempt).await { + Ok(execution) => execution, Err(err) => { let remaining = source.drain_execution_attempts().await?; port.mark_unused_attempts(remaining).await?; @@ -793,15 +817,26 @@ where "stream_candidate_execute", execute_started_at.elapsed().as_millis() as u64, ); - if let Some(response) = response { - let remaining = source.drain_execution_attempts().await?; - let unused_started_at = std::time::Instant::now(); - port.mark_unused_attempts(remaining).await?; - observe_gateway_stage_ms( - "stream_candidate_unused", - unused_started_at.elapsed().as_millis() as u64, - ); - return Ok(LocalExecutionRequestOutcome::responded(response)); + match execution { + AiAttemptExecutionOutcome::Responded(response) => { + let remaining = source.drain_execution_attempts().await?; + let unused_started_at = std::time::Instant::now(); + port.mark_unused_attempts(remaining).await?; + observe_gateway_stage_ms( + "stream_candidate_unused", + unused_started_at.elapsed().as_millis() as u64, + ); + return Ok(LocalExecutionRequestOutcome::responded(response)); + } + AiAttemptExecutionOutcome::Retry { + scope, + fallback_response: attempt_fallback_response, + } => { + if attempt_fallback_response.is_some() { + fallback_response = attempt_fallback_response; + } + apply_attempt_retry_scope(source, &attempt, scope).await?; + } } port.record_attempt_failed(&attempt).await?; @@ -816,6 +851,12 @@ where last_attempted = Some((attempt.execution_plan().clone(), attempt.report_context())); } + if let Some(response) = fallback_response { + return Ok(LocalExecutionRequestOutcome::responded( + mark_deferred_upstream_response(response), + )); + } + let Some((last_plan, last_report_context)) = last_attempted else { return Ok(LocalExecutionRequestOutcome::NoPath); }; @@ -826,6 +867,24 @@ where )) } +async fn apply_attempt_retry_scope( + source: &mut Source, + attempt: &Attempt, + scope: AiAttemptRetryScope, +) -> Result<(), GatewayError> +where + Source: LocalExecutionAttemptSource, + Attempt: AiExecutionAttempt, +{ + let plan = attempt.execution_plan(); + match scope { + AiAttemptRetryScope::Candidate => Ok(()), + AiAttemptRetryScope::Credential => source.skip_credential(plan.key_id.as_str()).await, + AiAttemptRetryScope::Endpoint => source.skip_endpoint(plan.endpoint_id.as_str()).await, + AiAttemptRetryScope::Provider => source.skip_provider(plan.provider_id.as_str()).await, + } +} + async fn next_execution_attempt_with_timeout( source: &mut Source, trace_id: &str, @@ -901,7 +960,10 @@ where Ok(()) } - async fn execute_attempt(&self, attempt: &T) -> Result, Self::Error> { + async fn execute_attempt( + &self, + attempt: &T, + ) -> Result, Self::Error> { let plan = attempt.execution_plan(); let report_context = attempt.report_context(); let candidate_index = parse_request_candidate_report_context(report_context.as_ref()) @@ -931,7 +993,7 @@ where ) .await? { - return Ok(Some(response)); + return Ok(AiAttemptExecutionOutcome::Responded(response)); } prewarm_direct_reqwest_candidate_client(plan); // The attempt owns the canonical report context. Borrow it for the @@ -951,14 +1013,14 @@ where let execution_decision = self.decision.clone(); let execution_report_kind = attempt.report_kind(); let execution_plan = plan.clone(); - let mut response = execute_stream_candidate_with_watchdog( + let mut execution = execute_stream_candidate_with_watchdog( self.state, self.trace_id, self.plan_kind, plan, watchdog_report_context, move || async move { - execute_execution_runtime_stream( + execute_execution_runtime_stream_with_retry_scope( &execution_state, execution_plan, execution_trace_id.as_str(), @@ -971,10 +1033,18 @@ where }, ) .await?; - if let Some(response) = response.as_mut() { - attach_redaction_execution_candidate(response, plan.candidate_id.as_deref()); + match &mut execution { + AiAttemptExecutionOutcome::Responded(response) + | AiAttemptExecutionOutcome::Retry { + fallback_response: Some(response), + .. + } => attach_redaction_execution_candidate(response, plan.candidate_id.as_deref()), + AiAttemptExecutionOutcome::Retry { + fallback_response: None, + .. + } => {} } - Ok(response) + Ok(execution) } async fn mark_unused_attempts(&self, attempts: Vec) -> Result<(), Self::Error> { @@ -1234,9 +1304,11 @@ async fn execute_stream_candidate_with_watchdog( plan: &aether_contracts::ExecutionPlan, report_context: Option<&serde_json::Value>, execute: impl FnOnce() -> Fut, -) -> Result>, GatewayError> +) -> Result>, GatewayError> where - Fut: std::future::Future>, GatewayError>> + Send, + Fut: std::future::Future< + Output = Result>, GatewayError>, + > + Send, { let timeout_duration = resolve_stream_candidate_watchdog_timeout(plan, report_context); let candidate_started_unix_ms = current_unix_ms(); @@ -1252,7 +1324,9 @@ where ) .await; log_stream_candidate_admission_timeout(trace_id, plan_kind, plan, report_context, &err); - return Ok(None); + return Ok(AiAttemptExecutionOutcome::retry( + AiAttemptRetryScope::Candidate, + )); } Err(err) => return Err(err), }; @@ -1300,7 +1374,9 @@ where timeout_ms, "gateway local stream candidate watchdog timed out" ); - Ok(None) + Ok(AiAttemptExecutionOutcome::retry( + AiAttemptRetryScope::Candidate, + )) } }; observe_gateway_stage_ms( @@ -1308,7 +1384,21 @@ where watchdog_started_at.elapsed().as_millis() as u64, ); match outcome { - Ok(response) => Ok(maybe_hold_upstream_execution_permit(response, permit_hold)), + Ok(AiAttemptExecutionOutcome::Responded(response)) => { + let response = maybe_hold_upstream_execution_permit(Some(response), permit_hold) + .expect("responded stream attempt must retain its response"); + Ok(AiAttemptExecutionOutcome::Responded(response)) + } + Ok(AiAttemptExecutionOutcome::Retry { + scope, + fallback_response, + }) => { + drop(permit_hold); + Ok(AiAttemptExecutionOutcome::Retry { + scope, + fallback_response, + }) + } Err(err) if is_candidate_level_admission_timeout(&err) => { drop(permit_hold); if should_record_candidate_admission_timeout(&err) { @@ -1322,7 +1412,9 @@ where .await; } log_stream_candidate_admission_timeout(trace_id, plan_kind, plan, report_context, &err); - Ok(None) + Ok(AiAttemptExecutionOutcome::retry( + AiAttemptRetryScope::Candidate, + )) } Err(err) => { drop(permit_hold); @@ -1605,6 +1697,14 @@ mod tests { Ok(Vec::new()) } + async fn skip_credential(&mut self, _key_id: &str) -> Result<(), GatewayError> { + Ok(()) + } + + async fn skip_endpoint(&mut self, _endpoint_id: &str) -> Result<(), GatewayError> { + Ok(()) + } + async fn skip_provider(&mut self, _provider_id: &str) -> Result<(), GatewayError> { Ok(()) } @@ -1638,6 +1738,7 @@ mod tests { struct TransferTestPort<'a> { state: &'a AppState, tracker: ProviderTransferTracker, + retry_scope: AiAttemptRetryScope, executed: StdMutex>, unused: StdMutex>, } @@ -1651,6 +1752,17 @@ mod tests { Self { state, tracker, + retry_scope: AiAttemptRetryScope::Candidate, + executed: StdMutex::new(Vec::new()), + unused: StdMutex::new(Vec::new()), + } + } + + fn with_retry_scope(state: &'a AppState, retry_scope: AiAttemptRetryScope) -> Self { + Self { + state, + tracker: ProviderTransferTracker::default(), + retry_scope, executed: StdMutex::new(Vec::new()), unused: StdMutex::new(Vec::new()), } @@ -1702,9 +1814,13 @@ mod tests { async fn execute_attempt( &self, attempt: &TransferTestAttempt, - ) -> Result, Self::Error> { + ) -> Result, Self::Error> { self.executed.lock().unwrap().push(attempt.label); - Ok((attempt.plan.provider_id == "provider-b").then(|| Response::new(Body::from("ok")))) + Ok(if attempt.plan.provider_id == "provider-b" { + AiAttemptExecutionOutcome::Responded(Response::new(Body::from("ok"))) + } else { + AiAttemptExecutionOutcome::retry(self.retry_scope) + }) } async fn mark_unused_attempts( @@ -1751,6 +1867,18 @@ mod tests { Ok(self.attempts.drain(..).collect()) } + async fn skip_credential(&mut self, key_id: &str) -> Result<(), GatewayError> { + self.attempts + .retain(|attempt| attempt.plan.key_id != key_id); + Ok(()) + } + + async fn skip_endpoint(&mut self, endpoint_id: &str) -> Result<(), GatewayError> { + self.attempts + .retain(|attempt| attempt.plan.endpoint_id != endpoint_id); + Ok(()) + } + async fn skip_provider(&mut self, provider_id: &str) -> Result<(), GatewayError> { self.skipped_providers.push(provider_id.to_string()); self.attempts @@ -1872,6 +2000,36 @@ mod tests { assert_eq!(source.skipped_providers, ["provider-a"]); } + #[tokio::test] + async fn dynamic_loop_applies_provider_scoped_retry_to_candidate_source() { + let state = AppState::new().expect("state should build"); + let port = TransferTestPort::with_retry_scope(&state, AiAttemptRetryScope::Provider); + let mut source = TransferTestAttemptSource { + attempts: transfer_test_attempts().into(), + skipped_providers: Vec::new(), + }; + + let outcome = run_dynamic_attempt_loop( + &port, + &mut source, + "trace-provider-scope-test", + "provider_scope_test", + Duration::from_secs(1), + ) + .await + .expect("dynamic attempt loop should succeed"); + + assert!(matches!( + outcome, + LocalExecutionRequestOutcome::Responded(_) + )); + assert_eq!( + port.executed.lock().unwrap().as_slice(), + ["a-key1-retry0", "b-key1-retry0"] + ); + assert_eq!(source.skipped_providers, ["provider-a"]); + } + #[test] fn transfer_timeout_is_checked_at_candidate_boundary_and_zero_disables_limits() { let started_at = Instant::now(); @@ -2171,14 +2329,24 @@ mod tests { "claude_cli_stream", &plan, Some(&report_context), - || std::future::pending::>, GatewayError>>(), + || { + std::future::pending::< + Result>, GatewayError>, + >() + }, ) .await }); tokio::time::sleep(Duration::from_millis(40)).await; let result = task.await.expect("watchdog task should join"); - assert!(matches!(result, Ok(None))); + assert!(matches!( + result, + Ok(AiAttemptExecutionOutcome::Retry { + scope: AiAttemptRetryScope::Candidate, + fallback_response: None, + }) + )); let records = writer.records.lock().await; assert_eq!(records.len(), 1); @@ -2226,7 +2394,13 @@ mod tests { ) .await; - assert!(matches!(result, Ok(None))); + assert!(matches!( + result, + Ok(AiAttemptExecutionOutcome::Retry { + scope: AiAttemptRetryScope::Candidate, + fallback_response: None, + }) + )); let records = writer.records.lock().await; assert_eq!(records.len(), 1); let record = &records[0]; @@ -2268,7 +2442,13 @@ mod tests { ) .await; - assert!(matches!(result, Ok(None))); + assert!(matches!( + result, + Ok(AiAttemptExecutionOutcome::Retry { + scope: AiAttemptRetryScope::Candidate, + fallback_response: None, + }) + )); assert!(writer.records.lock().await.is_empty()); } } diff --git a/apps/aether-gateway/src/executor/mod.rs b/apps/aether-gateway/src/executor/mod.rs index b5c8f3ad2..fe2fc9e6a 100644 --- a/apps/aether-gateway/src/executor/mod.rs +++ b/apps/aether-gateway/src/executor/mod.rs @@ -19,7 +19,8 @@ pub(crate) use orchestration::*; pub(crate) use outcome::{ beautify_local_execution_client_error_message, build_fast_local_execution_exhaustion, build_fast_local_execution_runtime_miss_context, build_local_execution_exhaustion, - build_local_execution_runtime_miss_context, record_failed_usage_for_exhausted_request, + build_local_execution_runtime_miss_context, is_deferred_upstream_response, + mark_deferred_upstream_response, record_failed_usage_for_exhausted_request, record_failed_usage_for_runtime_miss_request, LocalExecutionExhaustion, LocalExecutionRequestOutcome, LocalExecutionRuntimeMissContext, }; diff --git a/apps/aether-gateway/src/executor/orchestration.rs b/apps/aether-gateway/src/executor/orchestration.rs index ed9883ceb..379eab37f 100644 --- a/apps/aether-gateway/src/executor/orchestration.rs +++ b/apps/aether-gateway/src/executor/orchestration.rs @@ -1091,7 +1091,7 @@ fn standard_text_sync_heartbeat_error_kind(status_code: u16) -> LocalCoreSyncErr 401 => LocalCoreSyncErrorKind::Authentication, 403 => LocalCoreSyncErrorKind::PermissionDenied, 404 => LocalCoreSyncErrorKind::NotFound, - 413 => LocalCoreSyncErrorKind::ContextLengthExceeded, + 413 => LocalCoreSyncErrorKind::RequestTooLarge, 429 => LocalCoreSyncErrorKind::RateLimit, 503 => LocalCoreSyncErrorKind::Overloaded, _ => LocalCoreSyncErrorKind::ServerError, @@ -1586,6 +1586,18 @@ mod tests { Ok(self.attempts.drain(..).collect()) } + async fn skip_credential(&mut self, key_id: &str) -> Result<(), GatewayError> { + self.attempts + .retain(|attempt| attempt.plan.key_id != key_id); + Ok(()) + } + + async fn skip_endpoint(&mut self, endpoint_id: &str) -> Result<(), GatewayError> { + self.attempts + .retain(|attempt| attempt.plan.endpoint_id != endpoint_id); + Ok(()) + } + async fn skip_provider(&mut self, provider_id: &str) -> Result<(), GatewayError> { self.attempts .retain(|attempt| attempt.plan.provider_id != provider_id); @@ -1813,7 +1825,6 @@ mod tests { test_openai_image_heartbeat_attempt(0, "endpoint-retry", "candidate-retry"), test_openai_image_heartbeat_attempt(1, "endpoint-success", "candidate-success"), ]; - let outcome = execute_openai_image_sync_heartbeat_attempts( state, "/v1/images/generations".to_string(), @@ -2104,7 +2115,6 @@ mod tests { "openai:responses:compact", ), ]; - let (parts, _) = http::Request::builder() .method(http::Method::POST) .uri("/v1/responses") diff --git a/apps/aether-gateway/src/executor/outcome.rs b/apps/aether-gateway/src/executor/outcome.rs index 818ed2e63..046e415e4 100644 --- a/apps/aether-gateway/src/executor/outcome.rs +++ b/apps/aether-gateway/src/executor/outcome.rs @@ -18,6 +18,7 @@ use base64::Engine as _; use serde_json::{json, Map, Value}; use tracing::warn; +use crate::ai_serving::{build_core_error_body_for_client_format, LocalCoreSyncErrorKind}; use crate::constants::{ EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS, LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER, }; @@ -32,6 +33,9 @@ pub(crate) enum LocalExecutionRequestOutcome { NoPath, } +#[derive(Debug, Clone, Copy)] +pub(crate) struct DeferredUpstreamResponse; + #[derive(Debug, Clone)] pub(crate) struct LocalExecutionExhaustion { request_id: String, @@ -70,6 +74,18 @@ impl LocalExecutionRequestOutcome { } } +pub(crate) fn mark_deferred_upstream_response(mut response: Response) -> Response { + response.extensions_mut().insert(DeferredUpstreamResponse); + response +} + +pub(crate) fn is_deferred_upstream_response(response: &Response) -> bool { + response + .extensions() + .get::() + .is_some() +} + impl LocalExecutionRuntimeMissContext { pub(crate) fn persisted_candidate_count(&self) -> usize { self.candidate_contexts.len() @@ -316,12 +332,12 @@ pub(crate) async fn record_failed_usage_for_exhausted_request( ); } data.client_response_headers = Some(Value::Object(client_headers)); - data.client_response_body = Some(json!({ - "error": { - "type": "http_error", - "message": beautify_local_execution_client_error_message(local_execution_runtime_miss_detail), - } - })); + let client_message = + beautify_local_execution_client_error_message(local_execution_runtime_miss_detail); + data.client_response_body = Some(runtime_miss_client_error_body( + data.api_format.as_deref(), + &client_message, + )); let mut request_metadata = match data.request_metadata.take() { Some(Value::Object(object)) => object, @@ -392,12 +408,7 @@ pub(crate) async fn record_failed_usage_for_runtime_miss_request( let status_code = http::StatusCode::SERVICE_UNAVAILABLE.as_u16(); let client_message = beautify_local_execution_client_error_message(local_execution_runtime_miss_detail); - let client_body = json!({ - "error": { - "type": "http_error", - "message": client_message, - } - }); + let client_body = runtime_miss_client_error_body(api_format.as_deref(), &client_message); let mut client_headers = Map::from_iter([( "content-type".to_string(), Value::String("application/json".to_string()), @@ -654,6 +665,30 @@ fn json_header_map() -> Value { )])) } +fn runtime_miss_client_error_body(api_format: Option<&str>, message: &str) -> Value { + let fallback = json!({ + "error": { + "type": "http_error", + "message": message, + } + }); + let is_claude = api_format.is_some_and(|format| { + crate::ai_serving::normalize_api_format_alias(format) + .eq_ignore_ascii_case("claude:messages") + }); + if !is_claude { + return fallback; + } + + build_core_error_body_for_client_format( + "claude:messages", + message, + None, + LocalCoreSyncErrorKind::Overloaded, + ) + .unwrap_or(fallback) +} + fn runtime_miss_original_headers_json(headers: &HeaderMap) -> Value { let mut headers = crate::headers::collect_control_headers(headers); for (name, value) in headers.iter_mut() { @@ -1135,7 +1170,7 @@ fn trimmed_non_empty(value: Option<&str>) -> Option { mod tests { use super::{ apply_runtime_miss_usage_routing, beautify_local_execution_client_error_message, - request_candidate_represents_provider_execution, + request_candidate_represents_provider_execution, runtime_miss_client_error_body, select_last_runtime_miss_executed_candidate, LocalExecutionRuntimeMissContext, RuntimeMissCandidateContext, }; @@ -1169,6 +1204,17 @@ mod tests { ); } + #[test] + fn runtime_miss_usage_body_matches_claude_client_envelope() { + let claude = runtime_miss_client_error_body(Some("claude:messages"), "busy"); + assert_eq!(claude["type"], "error"); + assert_eq!(claude["error"]["type"], "overloaded_error"); + + let openai = runtime_miss_client_error_body(Some("openai:chat"), "busy"); + assert_eq!(openai["error"]["type"], "http_error"); + assert!(openai.get("type").is_none()); + } + #[test] fn runtime_miss_routing_moves_to_typed_usage_fields_and_keeps_metadata_lightweight() { let mut data = UsageEventData::default(); diff --git a/apps/aether-gateway/src/executor/stream_path.rs b/apps/aether-gateway/src/executor/stream_path.rs index 7077bc19e..921122e16 100644 --- a/apps/aether-gateway/src/executor/stream_path.rs +++ b/apps/aether-gateway/src/executor/stream_path.rs @@ -1,6 +1,6 @@ use aether_ai_serving::{ run_ai_stream_execution_path, AiPlanFallbackReason, AiServingExecutionOutcome, - AiStreamExecutionPathPort, AiStreamExecutionStep, + AiStreamExecutionPathPort, AiStreamExecutionStep, OriginalRequestPayload, }; use async_trait::async_trait; use axum::body::{Body, Bytes}; @@ -59,6 +59,21 @@ pub(crate) async fn maybe_execute_via_stream_decision_path( ); return Ok(LocalExecutionRequestOutcome::NoPath); }; + let mut planning_parts = parts.clone(); + if crate::ai_serving::is_json_request(&planning_parts.headers) { + if let Ok(decoded_body) = crate::ai_serving::decoded_request_body_bytes( + &planning_parts.headers, + body_bytes.as_ref(), + ) { + planning_parts + .extensions + .insert(OriginalRequestPayload::from_parsed_json( + body_json.clone(), + decoded_body.as_ref(), + )); + } + } + let parts = &planning_parts; observe_gateway_stage_ms( "frontdoor_stream_parse", parse_started_at.elapsed().as_millis() as u64, @@ -422,7 +437,11 @@ fn to_ai_serving_outcome( ) -> AiServingExecutionOutcome, super::LocalExecutionExhaustion> { match outcome { LocalExecutionRequestOutcome::Responded(response) => { - AiServingExecutionOutcome::Responded(response) + if super::is_deferred_upstream_response(&response) { + AiServingExecutionOutcome::Deferred(response) + } else { + AiServingExecutionOutcome::Responded(response) + } } LocalExecutionRequestOutcome::Exhausted(outcome) => { AiServingExecutionOutcome::Exhausted(outcome) @@ -438,6 +457,9 @@ fn from_ai_serving_outcome( AiServingExecutionOutcome::Responded(response) => { LocalExecutionRequestOutcome::Responded(response) } + AiServingExecutionOutcome::Deferred(response) => { + LocalExecutionRequestOutcome::Responded(response) + } AiServingExecutionOutcome::Exhausted(outcome) => { LocalExecutionRequestOutcome::Exhausted(outcome) } diff --git a/apps/aether-gateway/src/executor/sync_path.rs b/apps/aether-gateway/src/executor/sync_path.rs index 3b7054ea1..3af16aa78 100644 --- a/apps/aether-gateway/src/executor/sync_path.rs +++ b/apps/aether-gateway/src/executor/sync_path.rs @@ -1,6 +1,6 @@ use aether_ai_serving::{ run_ai_sync_execution_path, AiPlanFallbackReason, AiServingExecutionOutcome, - AiSyncExecutionPathPort, AiSyncExecutionStep, + AiSyncExecutionPathPort, AiSyncExecutionStep, OriginalRequestPayload, }; use async_trait::async_trait; use axum::body::{Body, Bytes}; @@ -51,6 +51,22 @@ pub(crate) async fn maybe_execute_via_sync_decision_path( return Ok(LocalExecutionRequestOutcome::NoPath); }; + let mut planning_parts = parts.clone(); + if crate::ai_serving::is_json_request(&planning_parts.headers) { + if let Ok(decoded_body) = crate::ai_serving::decoded_request_body_bytes( + &planning_parts.headers, + body_bytes.as_ref(), + ) { + planning_parts + .extensions + .insert(OriginalRequestPayload::from_parsed_json( + body_json.clone(), + decoded_body.as_ref(), + )); + } + } + let parts = &planning_parts; + if let Some(stream_plan_kind) = resolve_execution_runtime_stream_plan_kind(parts, decision) { if is_matching_stream_request(stream_plan_kind, parts, &body_json, body_base64.as_deref()) { return Ok(LocalExecutionRequestOutcome::NoPath); @@ -256,7 +272,11 @@ fn to_ai_serving_outcome( ) -> AiServingExecutionOutcome, super::LocalExecutionExhaustion> { match outcome { LocalExecutionRequestOutcome::Responded(response) => { - AiServingExecutionOutcome::Responded(response) + if super::is_deferred_upstream_response(&response) { + AiServingExecutionOutcome::Deferred(response) + } else { + AiServingExecutionOutcome::Responded(response) + } } LocalExecutionRequestOutcome::Exhausted(outcome) => { AiServingExecutionOutcome::Exhausted(outcome) @@ -272,6 +292,9 @@ fn from_ai_serving_outcome( AiServingExecutionOutcome::Responded(response) => { LocalExecutionRequestOutcome::Responded(response) } + AiServingExecutionOutcome::Deferred(response) => { + LocalExecutionRequestOutcome::Responded(response) + } AiServingExecutionOutcome::Exhausted(outcome) => { LocalExecutionRequestOutcome::Exhausted(outcome) } diff --git a/apps/aether-gateway/src/fallback_metrics.rs b/apps/aether-gateway/src/fallback_metrics.rs index 94f614063..9920e7923 100644 --- a/apps/aether-gateway/src/fallback_metrics.rs +++ b/apps/aether-gateway/src/fallback_metrics.rs @@ -173,6 +173,9 @@ mod tests { route_class: Some("ai_public".to_string()), route_family: Some("openai".to_string()), route_kind: Some("chat".to_string()), + client_surface: None, + api_operation: None, + gateway_credential_carrier: None, request_auth_channel: None, auth_endpoint_signature: None, execution_runtime_candidate: true, diff --git a/apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/complete/key.rs b/apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/complete/key.rs index 7ef6be85f..cc440a89d 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/complete/key.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/complete/key.rs @@ -287,6 +287,7 @@ pub(super) async fn handle_admin_provider_oauth_complete_key( &ProviderCatalogKeyOAuthRuntimeStateCasUpdate { key_id: key_id.clone(), expected_encrypted_auth_config: state_data.expected_encrypted_auth_config, + expected_credential: None, encrypted_auth_config: persisted_encrypted_auth_config.clone(), encrypted_api_key_update: Some(encrypted_api_key), expires_at_unix_secs_update: Some(expires_at), 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 063bb7260..822e86c9f 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 @@ -370,6 +370,7 @@ pub(crate) async fn persist_fenced_provider_quota_refresh_state( &ProviderCatalogKeyOAuthRuntimeStateCasUpdate { key_id: key_id.to_string(), expected_encrypted_auth_config: Some(expected_encrypted_auth_config.to_string()), + expected_credential: None, encrypted_auth_config: expected_encrypted_auth_config.to_string(), encrypted_api_key_update: None, expires_at_unix_secs_update: None, diff --git a/apps/aether-gateway/src/handlers/admin/provider/oauth/runtime.rs b/apps/aether-gateway/src/handlers/admin/provider/oauth/runtime.rs index ce49f0a6c..d61c9bee0 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/oauth/runtime.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/oauth/runtime.rs @@ -100,14 +100,6 @@ fn select_provider_oauth_runtime_endpoint( .api_format .trim() .eq_ignore_ascii_case("gemini:generate_content") - }) - .or_else(|| { - matching_endpoint(endpoints, include_inactive, |endpoint| { - endpoint - .api_format - .trim() - .eq_ignore_ascii_case("claude:messages") - }) }), _ => matching_endpoint(endpoints, include_inactive, |_| true), } @@ -255,3 +247,44 @@ pub(crate) fn spawn_provider_oauth_account_state_refresh_after_update( .await; }); } + +#[cfg(test)] +mod tests { + use super::{ + provider_oauth_maintenance_endpoint_for_provider, + provider_oauth_runtime_endpoint_for_provider, + }; + use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogEndpoint; + + fn endpoint(id: &str, api_format: &str, is_active: bool) -> StoredProviderCatalogEndpoint { + StoredProviderCatalogEndpoint::new( + id.to_string(), + "provider-1".to_string(), + api_format.to_string(), + None, + None, + is_active, + ) + .expect("endpoint should build") + } + + #[test] + fn vertex_oauth_runtime_never_falls_back_to_retired_claude_endpoint() { + let endpoints = vec![endpoint("claude", "claude:messages", true)]; + + assert!(provider_oauth_runtime_endpoint_for_provider("vertex_ai", &endpoints).is_none()); + assert!( + provider_oauth_maintenance_endpoint_for_provider("vertex_ai", &endpoints).is_none() + ); + + let endpoints = vec![ + endpoint("claude", "claude:messages", true), + endpoint("gemini", "gemini:generate_content", true), + ]; + assert_eq!( + provider_oauth_runtime_endpoint_for_provider("vertex_ai", &endpoints) + .map(|endpoint| endpoint.id), + Some("gemini".to_string()) + ); + } +} diff --git a/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test.rs b/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test.rs index aa3e060d2..e96973d01 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test.rs @@ -2668,6 +2668,7 @@ async fn provider_query_execute_antigravity_test_candidate( upstream_is_stream: false, request_query: parts.uri.query(), kiro_api_region: None, + api_operation: None, }, ); let Some(request_url) = request_url else { @@ -3364,6 +3365,7 @@ async fn provider_query_execute_standard_test_candidate( upstream_is_stream, request_query: parts.uri.query(), kiro_api_region: None, + api_operation: None, }, Some(&provider_request_body), ); diff --git a/apps/aether-gateway/src/handlers/admin/provider/write/normalize.rs b/apps/aether-gateway/src/handlers/admin/provider/write/normalize.rs index 5f5f8c5f6..636f4d616 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/write/normalize.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/write/normalize.rs @@ -195,11 +195,7 @@ pub(crate) fn validate_vertex_api_formats( let allowed = match auth_type { "api_key" => &["gemini:generate_content", "gemini:embedding"][..], - "service_account" | "vertex_ai" => &[ - "claude:messages", - "gemini:generate_content", - "gemini:embedding", - ][..], + "service_account" | "vertex_ai" => &["gemini:generate_content", "gemini:embedding"][..], _ => return Ok(()), }; let invalid = api_formats @@ -410,20 +406,11 @@ mod tests { } #[test] - fn validate_vertex_api_formats_uses_canonical_message_formats() { + fn validate_vertex_api_formats_rejects_unimplemented_anthropic_transport() { assert!(validate_vertex_api_formats( "vertex_ai", "service_account", - &[ - "claude:messages".to_string(), - "gemini:generate_content".to_string() - ], - ) - .is_ok()); - assert!(validate_vertex_api_formats( - "vertex_ai", - "service_account", - &["claude:chat".to_string()], + &["claude:messages".to_string()], ) .is_err()); } @@ -443,7 +430,6 @@ mod tests { "vertex_ai", "service_account", &[ - "claude:messages".to_string(), "gemini:generate_content".to_string(), "gemini:embedding".to_string() ], diff --git a/apps/aether-gateway/src/handlers/admin/provider/write/provider/create.rs b/apps/aether-gateway/src/handlers/admin/provider/write/provider/create.rs index a791c89bf..bc3ba3d4b 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/write/provider/create.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/write/provider/create.rs @@ -158,6 +158,8 @@ pub(crate) async fn build_admin_create_provider_record( } } let config = (!config_map.is_empty()).then_some(serde_json::Value::Object(config_map)); + crate::provider_transport::validate_anthropic_compatibility_profile_config(config.as_ref()) + .map_err(|_| "无效的 Anthropic compatibility profile".to_string())?; let now_unix_secs = SystemTime::now() .duration_since(UNIX_EPOCH) diff --git a/apps/aether-gateway/src/handlers/admin/provider/write/provider/update.rs b/apps/aether-gateway/src/handlers/admin/provider/write/provider/update.rs index d45eea9b2..15b23ea05 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/write/provider/update.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/write/provider/update.rs @@ -312,6 +312,10 @@ pub(crate) async fn build_admin_update_provider_record( } updated.config = (!config_map.is_empty()).then_some(serde_json::Value::Object(config_map)); + crate::provider_transport::validate_anthropic_compatibility_profile_config( + updated.config.as_ref(), + ) + .map_err(|_| "无效的 Anthropic compatibility profile".to_string())?; updated.updated_at_unix_secs = SystemTime::now() .duration_since(UNIX_EPOCH) .ok() diff --git a/apps/aether-gateway/src/handlers/admin/request/provider/builders.rs b/apps/aether-gateway/src/handlers/admin/request/provider/builders.rs index 7bb2bcf76..fe86e0491 100644 --- a/apps/aether-gateway/src/handlers/admin/request/provider/builders.rs +++ b/apps/aether-gateway/src/handlers/admin/request/provider/builders.rs @@ -259,6 +259,10 @@ impl<'a> AdminAppState<'a> { admin_endpoint_signature_parts(&payload.api_format) .ok_or_else(|| format!("无效的 api_format: {}", payload.api_format))?; validate_admin_endpoint_stream_policy(normalized_api_format, payload.config.as_ref())?; + crate::provider_transport::validate_anthropic_compatibility_profile_config( + payload.config.as_ref(), + ) + .map_err(|_| "无效的 Anthropic compatibility profile".to_string())?; let base_url = normalize_admin_base_url(&payload.base_url)?; let existing_endpoints = self @@ -369,6 +373,10 @@ impl<'a> AdminAppState<'a> { existing_endpoint.api_format.as_str(), updated.config.as_ref(), )?; + crate::provider_transport::validate_anthropic_compatibility_profile_config( + updated.config.as_ref(), + ) + .map_err(|_| "无效的 Anthropic compatibility profile".to_string())?; } if provider_type == "codex" diff --git a/apps/aether-gateway/src/handlers/admin/request/system/import.rs b/apps/aether-gateway/src/handlers/admin/request/system/import.rs index 5f9ce6917..71540a7c2 100644 --- a/apps/aether-gateway/src/handlers/admin/request/system/import.rs +++ b/apps/aether-gateway/src/handlers/admin/request/system/import.rs @@ -192,6 +192,15 @@ fn normalize_import_endpoint_format(value: &str) -> Result { .ok_or_else(|| format!("无效的 api_format: {value}")) } +fn fixed_provider_import_endpoint_supported(provider_type: &str, api_format: &str) -> bool { + crate::provider_transport::provider_types::fixed_provider_template(provider_type).is_none() + || crate::provider_transport::provider_types::fixed_provider_endpoint_template_by_api_format( + provider_type, + api_format, + ) + .is_some() +} + fn normalize_import_key_formats( item: &ImportedProviderKey, provider_endpoint_formats: &BTreeSet, @@ -1419,6 +1428,12 @@ impl<'a> AdminAppState<'a> { for imported_provider_item in imported_providers { let (raw_provider, imported_provider) = imported_provider_item.into_parts(); let provider_name = invalid!(trim_required(&imported_provider.name, "name")); + invalid!( + crate::provider_transport::validate_anthropic_compatibility_profile_config( + imported_provider.config.as_ref(), + ) + .map_err(|_| "无效的 Anthropic compatibility profile".to_string()) + ); let existing_provider = providers_by_name.get(&provider_name).cloned(); let provider = if let Some(existing) = existing_provider { @@ -1513,6 +1528,42 @@ impl<'a> AdminAppState<'a> { let normalized_api_format = invalid!(normalize_import_endpoint_format( &imported_endpoint.api_format )); + invalid!( + crate::provider_transport::validate_anthropic_compatibility_profile_config( + imported_endpoint.config.as_ref(), + ) + .map_err(|_| "无效的 Anthropic compatibility profile".to_string()) + ); + if !fixed_provider_import_endpoint_supported( + &provider.provider_type, + &normalized_api_format, + ) { + let retired = existing_endpoints_by_format.remove(&normalized_api_format); + if let Some(mut retired) = retired { + if retired.is_active { + retired.is_active = false; + retired.updated_at_unix_secs = SystemTime::now() + .duration_since(UNIX_EPOCH) + .ok() + .map(|duration| duration.as_secs()); + let Some(_) = self.update_provider_catalog_endpoint(&retired).await? + else { + return Ok(Err(invalid_request(format!( + "停用 Provider '{provider_name}' 的已移除 Endpoint '{normalized_api_format}' 失败" + )))); + }; + stats.endpoints.updated += 1; + } else { + stats.endpoints.skipped += 1; + } + } else { + stats.endpoints.skipped += 1; + } + stats.errors.push(format!( + "固定 Provider '{provider_name}' 不再支持 Endpoint '{normalized_api_format}'，已跳过或停用" + )); + continue; + } let existing_endpoint = existing_endpoints_by_format .get(&normalized_api_format) .cloned(); diff --git a/apps/aether-gateway/src/handlers/proxy/mod.rs b/apps/aether-gateway/src/handlers/proxy/mod.rs index 19c797c8c..c9fd5dcdc 100644 --- a/apps/aether-gateway/src/handlers/proxy/mod.rs +++ b/apps/aether-gateway/src/handlers/proxy/mod.rs @@ -17,8 +17,8 @@ use crate::ai_serving::api::{ }; use crate::api::response::{ build_client_response, build_client_response_from_parts, build_local_auth_rejection_response, - build_local_http_error_response, build_local_overloaded_response, - build_local_user_rpm_limited_response, + build_local_http_error_response, build_local_http_error_response_with_request_path, + build_local_overloaded_response, build_local_user_rpm_limited_response, }; use crate::constants::{ CONTROL_CANDIDATE_ID_HEADER, DEPENDENCY_REASON_HEADER, EXECUTION_PATH_CONTROL_EXECUTE_STREAM, @@ -935,7 +935,13 @@ async fn proxy_request_inner( limit, })) => { let trace_id = extract_or_generate_trace_id(request.headers()); - let response = build_local_overloaded_response(&trace_id, None, gate, limit)?; + let response = build_local_overloaded_response( + &trace_id, + None, + Some(request.uri().path()), + gate, + limit, + )?; return Ok(finalize_gateway_response( &state, response, @@ -965,7 +971,13 @@ async fn proxy_request_inner( aether_runtime_state::RuntimeSemaphoreError::Unavailable { gate, limit, .. }, )) => { let trace_id = extract_or_generate_trace_id(request.headers()); - let response = build_local_overloaded_response(&trace_id, None, gate, limit)?; + let response = build_local_overloaded_response( + &trace_id, + None, + Some(request.uri().path()), + gate, + limit, + )?; return Ok(finalize_gateway_response( &state, response, @@ -999,9 +1011,10 @@ async fn proxy_request_inner( path = %request.uri().path(), "gateway rejected blacklisted client IP" ); - let response = build_local_http_error_response( + let response = build_local_http_error_response_with_request_path( &trace_id, None, + Some(request.uri().path()), http::StatusCode::FORBIDDEN, "当前 IP 已被禁止访问", )?; @@ -1057,9 +1070,10 @@ async fn proxy_request_inner( loop_guard_header = EXECUTION_RUNTIME_LOOP_GUARD_HEADER, "gateway rejected execution runtime request loop into frontdoor" ); - let response = build_local_http_error_response( + let response = build_local_http_error_response_with_request_path( &trace_id, None, + Some(parts.uri.path()), http::StatusCode::LOOP_DETECTED, LOCAL_EXECUTION_LOOP_DETECTED_DETAIL, )?; @@ -1534,9 +1548,10 @@ async fn proxy_request_inner( } if control_decision.is_none() { - let response = build_local_http_error_response( + let response = build_local_http_error_response_with_request_path( &trace_id, None, + Some(request_context.request_path.as_str()), http::StatusCode::NOT_FOUND, LOCAL_ROUTE_NOT_FOUND_DETAIL, )?; diff --git a/apps/aether-gateway/src/handlers/public/ai_public.rs b/apps/aether-gateway/src/handlers/public/ai_public.rs index 8ca6fc401..b689cc30a 100644 --- a/apps/aether-gateway/src/handlers/public/ai_public.rs +++ b/apps/aether-gateway/src/handlers/public/ai_public.rs @@ -1,4 +1,6 @@ -use crate::ai_serving::normalize_openai_image_quality; +use crate::ai_serving::{ + build_core_error_body_for_client_format, normalize_openai_image_quality, LocalCoreSyncErrorKind, +}; use crate::async_task::CancelVideoTaskError; use crate::control::GatewayControlDecision; use crate::control::GatewayPublicRequestContext; @@ -13,8 +15,6 @@ use axum::response::IntoResponse; use axum::Json; use serde_json::{json, Value}; -const CLAUDE_COUNT_TOKENS_INVALID_PAYLOAD_DETAIL: &str = "Invalid token count payload"; -const CLAUDE_COUNT_TOKENS_MISSING_BODY_DETAIL: &str = "请求体不能为空"; const GEMINI_VIDEO_TASK_NOT_FOUND_DETAIL: &str = "Video task not found"; const AI_PUBLIC_METHOD_NOT_ALLOWED_DETAIL: &str = "Method not allowed"; const AI_PUBLIC_UNAUTHORIZED_DETAIL: &str = "Unauthorized"; @@ -51,6 +51,10 @@ const OPENAI_RERANK_TOP_N_DETAIL: &str = "Rerank request top_n must be a positiv const OPENAI_RERANK_CHAT_PAYLOAD_DETAIL: &str = "Rerank request must use query/documents, not chat messages"; const OPENAI_RERANK_STREAM_UNSUPPORTED_DETAIL: &str = "Rerank requests do not support streaming"; +const CLAUDE_COUNT_TOKENS_BODY_REQUIRED_DETAIL: &str = "Request body is required"; +const CLAUDE_COUNT_TOKENS_INVALID_JSON_DETAIL: &str = "Invalid JSON body"; +const CLAUDE_COUNT_TOKENS_MODEL_REQUIRED_DETAIL: &str = "model: Field required"; +const CLAUDE_COUNT_TOKENS_MESSAGES_REQUIRED_DETAIL: &str = "messages: Field required"; const ANTIGRAVITY_USER_SETTINGS_MISSING_BODY_DETAIL: &str = "Antigravity setUserSettings request body is required"; const ANTIGRAVITY_USER_SETTINGS_INVALID_JSON_DETAIL: &str = @@ -135,7 +139,7 @@ pub(crate) async fn maybe_build_local_ai_public_response( } if let Some(response) = - maybe_build_local_claude_count_tokens_response(request_context, request_body) + maybe_build_local_claude_count_tokens_validation_response(request_context, request_body) { return Some(response); } @@ -863,7 +867,7 @@ fn maybe_build_local_ai_public_route_guard_response( None } -fn maybe_build_local_claude_count_tokens_response( +fn maybe_build_local_claude_count_tokens_validation_response( request_context: &GatewayPublicRequestContext, request_body: Option<&Bytes>, ) -> Option> { @@ -876,34 +880,45 @@ fn maybe_build_local_claude_count_tokens_response( return None; } - let Some(request_body) = request_body else { - return Some(build_ai_public_error_response( - http::StatusCode::BAD_REQUEST, - CLAUDE_COUNT_TOKENS_MISSING_BODY_DETAIL, - )); - }; + let validation = validate_claude_count_tokens_request(request_body); + validation.err().map(build_claude_invalid_request_response) +} - let payload = match serde_json::from_slice::(request_body) { - Ok(payload) => payload, - Err(_) => { - return Some(build_ai_public_error_response( - http::StatusCode::BAD_REQUEST, - CLAUDE_COUNT_TOKENS_INVALID_PAYLOAD_DETAIL, - )); - } - }; +fn validate_claude_count_tokens_request(request_body: Option<&Bytes>) -> Result<(), &'static str> { + let request_body = request_body + .filter(|body| !body.is_empty()) + .ok_or(CLAUDE_COUNT_TOKENS_BODY_REQUIRED_DETAIL)?; + let payload = serde_json::from_slice::(request_body) + .map_err(|_| CLAUDE_COUNT_TOKENS_INVALID_JSON_DETAIL)?; + let object = payload + .as_object() + .ok_or(CLAUDE_COUNT_TOKENS_INVALID_JSON_DETAIL)?; - let input_tokens = match estimate_claude_count_tokens(&payload) { - Ok(tokens) => tokens, - Err(_) => { - return Some(build_ai_public_error_response( - http::StatusCode::BAD_REQUEST, - CLAUDE_COUNT_TOKENS_INVALID_PAYLOAD_DETAIL, - )); - } - }; + if object + .get("model") + .and_then(Value::as_str) + .map(str::trim) + .filter(|model| !model.is_empty()) + .is_none() + { + return Err(CLAUDE_COUNT_TOKENS_MODEL_REQUIRED_DETAIL); + } + if object.get("messages").and_then(Value::as_array).is_none() { + return Err(CLAUDE_COUNT_TOKENS_MESSAGES_REQUIRED_DETAIL); + } - Some(Json(json!({ "input_tokens": input_tokens })).into_response()) + Ok(()) +} + +fn build_claude_invalid_request_response(detail: &'static str) -> Response { + let body = build_core_error_body_for_client_format( + "claude:messages", + detail, + None, + LocalCoreSyncErrorKind::InvalidRequest, + ) + .expect("Claude core error format should be available"); + (http::StatusCode::BAD_REQUEST, Json(body)).into_response() } fn maybe_build_local_antigravity_v1internal_response( @@ -1583,132 +1598,43 @@ fn build_ai_public_error_response( (status, Json(json!({ "detail": detail.into() }))).into_response() } -fn estimate_claude_count_tokens(payload: &serde_json::Value) -> Result { - let object = payload.as_object().ok_or(())?; - let model = object - .get("model") - .and_then(serde_json::Value::as_str) - .ok_or(())?; - if model.trim().is_empty() { - return Err(()); - } - - let messages = object - .get("messages") - .and_then(serde_json::Value::as_array) - .ok_or(())?; - - let system_tokens = estimate_claude_system_tokens(object.get("system"))?; - let message_tokens = estimate_claude_message_tokens(messages)?; - Ok(system_tokens.saturating_add(message_tokens)) -} - -fn estimate_claude_system_tokens(system: Option<&serde_json::Value>) -> Result { - let Some(system) = system else { - return Ok(0); - }; - - match system { - serde_json::Value::Null => Ok(0), - serde_json::Value::String(text) => Ok(estimate_text_tokens(text)), - serde_json::Value::Array(blocks) => { - let mut total = 0_u64; - for block in blocks { - let block = block.as_object().ok_or(())?; - if let Some(text) = block.get("text").and_then(serde_json::Value::as_str) { - total = total.saturating_add(estimate_text_tokens(text)); - } - } - Ok(total) - } - serde_json::Value::Object(_) => Ok(0), - _ => Err(()), - } -} - -fn estimate_claude_message_tokens(messages: &[serde_json::Value]) -> Result { - let mut total = 0_u64; - - for message in messages { - let message = message.as_object().ok_or(())?; - let role = message - .get("role") - .and_then(serde_json::Value::as_str) - .ok_or(())?; - if !matches!(role, "user" | "assistant") { - return Err(()); - } - - total = total.saturating_add(4); - let content = message.get("content").ok_or(())?; - match content { - serde_json::Value::String(text) => { - total = total.saturating_add(estimate_text_tokens(text)); - } - serde_json::Value::Array(items) => { - for item in items { - let item = item.as_object().ok_or(())?; - if let Some(text) = item.get("text").and_then(serde_json::Value::as_str) { - total = total.saturating_add(estimate_text_tokens(text)); - } - } - } - _ => return Err(()), - } - } - - Ok(total) -} - -fn estimate_text_tokens(text: &str) -> u64 { - if text.is_empty() { - return 0; - } - - let char_count = text.chars().count() as u64; - std::cmp::max(1, char_count / 4) -} - #[cfg(test)] mod tests { use super::{ - estimate_claude_count_tokens, parse_openai_image_validation_input, validate_openai_image_n, - OpenAiImageOperation, + parse_openai_image_validation_input, validate_claude_count_tokens_request, + validate_openai_image_n, OpenAiImageOperation, CLAUDE_COUNT_TOKENS_BODY_REQUIRED_DETAIL, + CLAUDE_COUNT_TOKENS_INVALID_JSON_DETAIL, CLAUDE_COUNT_TOKENS_MESSAGES_REQUIRED_DETAIL, + CLAUDE_COUNT_TOKENS_MODEL_REQUIRED_DETAIL, }; use axum::body::Bytes; use serde_json::json; #[test] - fn estimates_claude_count_tokens_from_system_and_messages() { - let payload = json!({ - "model": "claude-sonnet-4-5", - "system": [{"type": "text", "text": "abcdefghijklmnop"}], - "messages": [ - { - "role": "user", - "content": "abcdefghijkl" - }, - { - "role": "assistant", - "content": [ - {"type": "text", "text": "abcdefgh"}, - {"type": "tool_use", "name": "ignored", "input": {"city": "SF"}} - ] - } - ] - }); - - assert_eq!(estimate_claude_count_tokens(&payload), Ok(17)); - } - - #[test] - fn rejects_invalid_claude_count_tokens_payload() { - let payload = json!({ - "model": "claude-sonnet-4-5", - "messages": [{"role": "system", "content": "bad"}] - }); - - assert_eq!(estimate_claude_count_tokens(&payload), Err(())); + fn count_tokens_validation_rejects_only_structurally_invalid_requests() { + assert_eq!( + validate_claude_count_tokens_request(None), + Err(CLAUDE_COUNT_TOKENS_BODY_REQUIRED_DETAIL) + ); + assert_eq!( + validate_claude_count_tokens_request(Some(&Bytes::from_static(b"{"))), + Err(CLAUDE_COUNT_TOKENS_INVALID_JSON_DETAIL) + ); + assert_eq!( + validate_claude_count_tokens_request(Some(&Bytes::from_static(br#"{"messages":[]}"#,))), + Err(CLAUDE_COUNT_TOKENS_MODEL_REQUIRED_DETAIL) + ); + assert_eq!( + validate_claude_count_tokens_request(Some(&Bytes::from_static( + br#"{"model":"claude-sonnet-4-5"}"#, + ))), + Err(CLAUDE_COUNT_TOKENS_MESSAGES_REQUIRED_DETAIL) + ); + assert_eq!( + validate_claude_count_tokens_request(Some(&Bytes::from_static( + br#"{"model":"claude-sonnet-4-5","messages":[],"tools":[{"name":"x"}]}"#, + ))), + Ok(()) + ); } #[test] diff --git a/apps/aether-gateway/src/handlers/public/support/test_connection/route.rs b/apps/aether-gateway/src/handlers/public/support/test_connection/route.rs index 0e68bff02..b879f62b0 100644 --- a/apps/aether-gateway/src/handlers/public/support/test_connection/route.rs +++ b/apps/aether-gateway/src/handlers/public/support/test_connection/route.rs @@ -253,6 +253,7 @@ pub(super) async fn maybe_build_local_test_connection_route_response( upstream_is_stream: false, request_query: None, kiro_api_region: None, + api_operation: None, }, ); let Some(upstream_url) = upstream_url else { diff --git a/apps/aether-gateway/src/maintenance/runtime/fixed_provider_reconciliation.rs b/apps/aether-gateway/src/maintenance/runtime/fixed_provider_reconciliation.rs index 9ebf13544..b1beefbde 100644 --- a/apps/aether-gateway/src/maintenance/runtime/fixed_provider_reconciliation.rs +++ b/apps/aether-gateway/src/maintenance/runtime/fixed_provider_reconciliation.rs @@ -12,7 +12,6 @@ const FIXED_PROVIDER_RECONCILIATION_LOCK_KEY: &str = "task_runtime:lock:maintenance.provider.fixed_template.reconcile"; const FIXED_PROVIDER_RECONCILIATION_LOCK_TTL: Duration = Duration::from_secs(10 * 60); const FIXED_PROVIDER_RECONCILIATION_RETRY_DELAY: Duration = Duration::from_secs(2); -const RECONCILED_PROVIDER_TYPE: &str = "codex"; pub(crate) async fn perform_fixed_provider_reconciliation_once( state: &AppState, @@ -51,13 +50,9 @@ async fn reconcile_fixed_provider_templates(state: &AppState) -> Result<(), Gate let admin_state = AdminAppState::new(state); let mut failures = Vec::new(); for provider in &providers { - if !provider - .provider_type - .trim() - .eq_ignore_ascii_case(RECONCILED_PROVIDER_TYPE) - || admin_state - .fixed_provider_template(&provider.provider_type) - .is_none() + if admin_state + .fixed_provider_template(&provider.provider_type) + .is_none() { continue; } @@ -226,16 +221,16 @@ mod tests { .expect("key should build"); key.api_formats = Some(json!(["openai:responses"])); - let unrelated_fixed_provider = StoredProviderCatalogProvider::new( - "provider-claude-code".to_string(), - "Claude Code".to_string(), + let unrelated_provider = StoredProviderCatalogProvider::new( + "provider-custom".to_string(), + "Custom".to_string(), None, - "claude_code".to_string(), + "custom".to_string(), ) - .expect("unrelated fixed provider should build"); + .expect("unrelated provider should build"); let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( - vec![provider, unrelated_fixed_provider], + vec![provider, unrelated_provider], vec![responses], vec![key], )); @@ -284,7 +279,7 @@ mod tests { assert_eq!(keys.len(), 1); assert_eq!(keys[0].api_formats, Some(json!(["openai:responses"]))); assert!(repository - .list_endpoints_by_provider_ids(&["provider-claude-code".to_string()]) + .list_endpoints_by_provider_ids(&["provider-custom".to_string()]) .await .expect("unrelated endpoints should list") .is_empty()); @@ -298,4 +293,106 @@ mod tests { .expect("endpoints should list again"); assert_eq!(second_endpoints, first_endpoints); } + + #[tokio::test] + async fn fixed_provider_reconciliation_retires_removed_vertex_claude_endpoint() { + let provider = StoredProviderCatalogProvider::new( + "provider-vertex".to_string(), + "Vertex AI".to_string(), + None, + "vertex_ai".to_string(), + ) + .expect("provider should build"); + + let gemini = StoredProviderCatalogEndpoint::new( + "endpoint-vertex-gemini".to_string(), + provider.id.clone(), + "gemini:generate_content".to_string(), + Some("gemini".to_string()), + Some("generate_content".to_string()), + true, + ) + .expect("gemini endpoint should build") + .with_transport_fields( + "https://aiplatform.googleapis.com".to_string(), + None, + None, + Some(2), + None, + None, + None, + None, + ) + .expect("gemini endpoint transport should build"); + + let claude = StoredProviderCatalogEndpoint::new( + "endpoint-vertex-claude".to_string(), + provider.id.clone(), + "claude:messages".to_string(), + Some("claude".to_string()), + Some("messages".to_string()), + true, + ) + .expect("claude endpoint should build") + .with_transport_fields( + "https://aiplatform.googleapis.com".to_string(), + None, + None, + Some(2), + None, + Some(json!({ + "_aether_fixed_provider_template": { + "managed": true, + "provider_type": "vertex_ai", + "item_key": "claude:messages", + "version": 1, + "retired": false, + "overrides": [], + "config_keys": [] + } + })), + None, + None, + ) + .expect("claude endpoint transport should build"); + + let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( + vec![provider], + vec![gemini, claude], + vec![], + )); + let state = AppState::new() + .expect("gateway state should build") + .with_data_state_for_tests( + GatewayDataState::with_provider_catalog_repository_for_tests(repository.clone()), + ); + + assert!(perform_fixed_provider_reconciliation_once(&state) + .await + .expect("reconciliation should run")); + + let endpoints = repository + .list_endpoints_by_provider_ids(&["provider-vertex".to_string()]) + .await + .expect("vertex endpoints should list"); + let gemini = endpoints + .iter() + .find(|endpoint| endpoint.id == "endpoint-vertex-gemini") + .expect("gemini endpoint should remain"); + assert!(gemini.is_active); + + let claude = endpoints + .iter() + .find(|endpoint| endpoint.id == "endpoint-vertex-claude") + .expect("legacy claude endpoint should remain as retired history"); + assert!(!claude.is_active); + assert_eq!( + claude + .config + .as_ref() + .and_then(|value| value.get("_aether_fixed_provider_template")) + .and_then(|value| value.get("retired")), + Some(&json!(true)) + ); + } } diff --git a/apps/aether-gateway/src/orchestration/classifier.rs b/apps/aether-gateway/src/orchestration/classifier.rs index 16eec5ac1..0245f511a 100644 --- a/apps/aether-gateway/src/orchestration/classifier.rs +++ b/apps/aether-gateway/src/orchestration/classifier.rs @@ -54,6 +54,223 @@ impl LocalFailoverClassification { } } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum FailureRetryAction { + Stop, + SameCredential, + NextCandidate, + NextCredential, + NextEndpoint, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum FailureScope { + None, + Credential, + CredentialModel, + Endpoint, + Provider, +} + +impl FailureScope { + pub(crate) const fn affects_credential(self) -> bool { + matches!(self, Self::Credential | Self::CredentialModel) + } + + pub(crate) const fn allows_key_wide_effects(self) -> bool { + matches!(self, Self::None | Self::Credential) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum FailureTokenAction { + None, + ForceRefresh, + #[allow(dead_code)] + Quarantine, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct FailureDisposition { + pub(crate) retry_action: FailureRetryAction, + pub(crate) failure_scope: FailureScope, + pub(crate) token_action: FailureTokenAction, + pub(crate) preserve_upstream_error: bool, +} + +impl FailureDisposition { + const fn new( + retry_action: FailureRetryAction, + failure_scope: FailureScope, + token_action: FailureTokenAction, + preserve_upstream_error: bool, + ) -> Self { + Self { + retry_action, + failure_scope, + token_action, + preserve_upstream_error, + } + } +} + +pub(crate) const fn failure_disposition_from_local_classification( + classification: LocalFailoverClassification, + status_code: u16, +) -> FailureDisposition { + match classification { + LocalFailoverClassification::StopStatusCode + | LocalFailoverClassification::StopErrorPattern + | LocalFailoverClassification::StopExecutionError + | LocalFailoverClassification::StopCyberPolicy => FailureDisposition::new( + FailureRetryAction::Stop, + FailureScope::None, + FailureTokenAction::None, + true, + ), + LocalFailoverClassification::UseDefault => FailureDisposition::new( + FailureRetryAction::Stop, + FailureScope::None, + FailureTokenAction::None, + status_code >= 400, + ), + LocalFailoverClassification::RetrySuccessPattern => FailureDisposition::new( + FailureRetryAction::NextCandidate, + FailureScope::None, + FailureTokenAction::None, + false, + ), + LocalFailoverClassification::RetryStatusCode + | LocalFailoverClassification::RetryUpstreamFailure => FailureDisposition::new( + FailureRetryAction::NextCandidate, + FailureScope::None, + FailureTokenAction::None, + false, + ), + } +} + +pub(crate) const fn classify_anthropic_failure_disposition( + classification: LocalFailoverClassification, + status_code: u16, +) -> FailureDisposition { + if matches!( + classification, + LocalFailoverClassification::StopStatusCode + | LocalFailoverClassification::StopErrorPattern + | LocalFailoverClassification::StopExecutionError + | LocalFailoverClassification::StopCyberPolicy + ) { + let generic = failure_disposition_from_local_classification(classification, status_code); + return match status_code { + 401 => FailureDisposition::new( + generic.retry_action, + FailureScope::Credential, + FailureTokenAction::ForceRefresh, + generic.preserve_upstream_error, + ), + 403 => FailureDisposition::new( + generic.retry_action, + FailureScope::Credential, + FailureTokenAction::None, + generic.preserve_upstream_error, + ), + 404 => FailureDisposition::new( + generic.retry_action, + FailureScope::Endpoint, + FailureTokenAction::None, + generic.preserve_upstream_error, + ), + 429 => FailureDisposition::new( + generic.retry_action, + FailureScope::CredentialModel, + FailureTokenAction::None, + generic.preserve_upstream_error, + ), + 529 => FailureDisposition::new( + generic.retry_action, + FailureScope::Provider, + FailureTokenAction::None, + generic.preserve_upstream_error, + ), + 500..=599 => FailureDisposition::new( + generic.retry_action, + FailureScope::Endpoint, + FailureTokenAction::None, + generic.preserve_upstream_error, + ), + _ => generic, + }; + } + + match status_code { + 400 => FailureDisposition::new( + FailureRetryAction::Stop, + FailureScope::None, + FailureTokenAction::None, + true, + ), + 401 => FailureDisposition::new( + FailureRetryAction::NextCredential, + FailureScope::Credential, + FailureTokenAction::ForceRefresh, + true, + ), + 403 => FailureDisposition::new( + FailureRetryAction::NextCredential, + FailureScope::Credential, + FailureTokenAction::None, + true, + ), + 404 => FailureDisposition::new( + FailureRetryAction::NextEndpoint, + FailureScope::Endpoint, + FailureTokenAction::None, + true, + ), + 413 => FailureDisposition::new( + FailureRetryAction::Stop, + FailureScope::None, + FailureTokenAction::None, + true, + ), + 429 => FailureDisposition::new( + FailureRetryAction::NextCredential, + FailureScope::CredentialModel, + FailureTokenAction::None, + true, + ), + 529 => FailureDisposition::new( + FailureRetryAction::NextEndpoint, + FailureScope::Provider, + FailureTokenAction::None, + true, + ), + 500..=599 => FailureDisposition::new( + FailureRetryAction::NextEndpoint, + FailureScope::Endpoint, + FailureTokenAction::None, + true, + ), + _ => failure_disposition_from_local_classification(classification, status_code), + } +} + +pub(crate) fn classify_failure_disposition( + provider_api_format: &str, + classification: LocalFailoverClassification, + status_code: u16, +) -> FailureDisposition { + if provider_api_format + .trim() + .eq_ignore_ascii_case("claude:messages") + { + classify_anthropic_failure_disposition(classification, status_code) + } else { + failure_disposition_from_local_classification(classification, status_code) + } +} + pub(crate) fn classify_local_failover( policy: &LocalFailoverPolicy, input: LocalFailoverInput<'_>, @@ -267,7 +484,11 @@ fn local_failover_regex_rule_matches( mod tests { use std::collections::BTreeSet; - use super::{classify_local_failover, LocalFailoverClassification, LocalFailoverInput}; + use super::{ + classify_anthropic_failure_disposition, classify_local_failover, + failure_disposition_from_local_classification, FailureDisposition, FailureRetryAction, + FailureScope, FailureTokenAction, LocalFailoverClassification, LocalFailoverInput, + }; use crate::orchestration::{LocalFailoverPolicy, LocalFailoverRegexRule}; #[test] @@ -544,4 +765,138 @@ mod tests { LocalFailoverClassification::UseDefault ); } + + #[test] + fn legacy_classification_preserves_candidate_by_candidate_retry() { + assert_eq!( + failure_disposition_from_local_classification( + LocalFailoverClassification::RetryUpstreamFailure, + 429, + ), + FailureDisposition { + retry_action: FailureRetryAction::NextCandidate, + failure_scope: FailureScope::None, + token_action: FailureTokenAction::None, + preserve_upstream_error: false, + } + ); + assert_eq!( + failure_disposition_from_local_classification( + LocalFailoverClassification::StopErrorPattern, + 400, + ) + .retry_action, + FailureRetryAction::Stop + ); + } + + #[test] + fn anthropic_bad_request_stops_and_preserves_upstream_error() { + let disposition = classify_anthropic_failure_disposition( + LocalFailoverClassification::RetryUpstreamFailure, + 400, + ); + + assert_eq!(disposition.retry_action, FailureRetryAction::Stop); + assert_eq!(disposition.failure_scope, FailureScope::None); + assert_eq!(disposition.token_action, FailureTokenAction::None); + assert!(disposition.preserve_upstream_error); + } + + #[test] + fn anthropic_auth_failures_refresh_then_rotate_only_when_needed() { + let unauthorized = classify_anthropic_failure_disposition( + LocalFailoverClassification::RetryUpstreamFailure, + 401, + ); + assert_eq!( + unauthorized.retry_action, + FailureRetryAction::NextCredential + ); + assert_eq!(unauthorized.failure_scope, FailureScope::Credential); + assert_eq!(unauthorized.token_action, FailureTokenAction::ForceRefresh); + + let forbidden = classify_anthropic_failure_disposition( + LocalFailoverClassification::RetryUpstreamFailure, + 403, + ); + assert_eq!(forbidden.retry_action, FailureRetryAction::NextCredential); + assert_eq!(forbidden.failure_scope, FailureScope::Credential); + assert_eq!(forbidden.token_action, FailureTokenAction::None); + } + + #[test] + fn anthropic_rate_limit_rotates_with_credential_model_scope() { + let disposition = classify_anthropic_failure_disposition( + LocalFailoverClassification::RetryUpstreamFailure, + 429, + ); + + assert_eq!(disposition.retry_action, FailureRetryAction::NextCredential); + assert_eq!(disposition.failure_scope, FailureScope::CredentialModel); + assert!(disposition.failure_scope.affects_credential()); + assert!(!disposition.failure_scope.allows_key_wide_effects()); + assert!(disposition.preserve_upstream_error); + } + + #[test] + fn anthropic_overload_moves_endpoint_without_credential_penalty() { + let disposition = classify_anthropic_failure_disposition( + LocalFailoverClassification::RetryUpstreamFailure, + 529, + ); + + assert_eq!(disposition.retry_action, FailureRetryAction::NextEndpoint); + assert_eq!(disposition.failure_scope, FailureScope::Provider); + assert!(!disposition.failure_scope.affects_credential()); + assert!(!disposition.failure_scope.allows_key_wide_effects()); + assert_eq!(disposition.token_action, FailureTokenAction::None); + assert!(disposition.preserve_upstream_error); + } + + #[test] + fn anthropic_not_found_moves_endpoint_and_oversize_stops() { + let not_found = classify_anthropic_failure_disposition( + LocalFailoverClassification::RetryUpstreamFailure, + 404, + ); + assert_eq!(not_found.retry_action, FailureRetryAction::NextEndpoint); + assert_eq!(not_found.failure_scope, FailureScope::Endpoint); + assert!(not_found.preserve_upstream_error); + + let oversized = classify_anthropic_failure_disposition( + LocalFailoverClassification::RetryUpstreamFailure, + 413, + ); + assert_eq!(oversized.retry_action, FailureRetryAction::Stop); + assert_eq!(oversized.failure_scope, FailureScope::None); + assert!(oversized.preserve_upstream_error); + } + + #[test] + fn only_unscoped_and_credential_failures_allow_key_wide_effects() { + assert!(FailureScope::None.allows_key_wide_effects()); + assert!(FailureScope::Credential.allows_key_wide_effects()); + assert!(!FailureScope::CredentialModel.allows_key_wide_effects()); + assert!(!FailureScope::Endpoint.allows_key_wide_effects()); + assert!(!FailureScope::Provider.allows_key_wide_effects()); + } + + #[test] + fn anthropic_explicit_stop_keeps_failure_resource_scope() { + let auth = classify_anthropic_failure_disposition( + LocalFailoverClassification::StopStatusCode, + 401, + ); + assert_eq!(auth.retry_action, FailureRetryAction::Stop); + assert_eq!(auth.failure_scope, FailureScope::Credential); + assert_eq!(auth.token_action, FailureTokenAction::ForceRefresh); + + let overloaded = classify_anthropic_failure_disposition( + LocalFailoverClassification::StopStatusCode, + 529, + ); + assert_eq!(overloaded.retry_action, FailureRetryAction::Stop); + assert_eq!(overloaded.failure_scope, FailureScope::Provider); + } } diff --git a/apps/aether-gateway/src/orchestration/effects.rs b/apps/aether-gateway/src/orchestration/effects.rs index e5376a037..83c0113c2 100644 --- a/apps/aether-gateway/src/orchestration/effects.rs +++ b/apps/aether-gateway/src/orchestration/effects.rs @@ -26,9 +26,10 @@ use tokio::sync::Mutex as TokioMutex; use tracing::warn; use super::{ - local_failover_error_message, project_local_adaptive_rate_limit, + classify_failure_disposition, local_failover_error_message, project_local_adaptive_rate_limit, project_local_adaptive_success, project_local_failure_health, project_local_key_circuit_closed, - project_local_key_circuit_failure, project_local_success_health, LocalFailoverClassification, + project_local_key_circuit_failure, project_local_success_health, FailureScope, + LocalFailoverClassification, }; use crate::ai_serving::extract_pool_sticky_session_token; use crate::client_session_affinity::{ @@ -613,7 +614,8 @@ async fn record_attempt_failure_effect( context: LocalExecutionEffectContext<'_>, effect: LocalAttemptFailureEffect, ) { - if !local_candidate_failure_should_invalidate_affinity( + if !local_candidate_failure_should_invalidate_affinity_for_provider( + &context.plan.provider_api_format, effect.classification, effect.status_code, ) { @@ -683,6 +685,13 @@ async fn record_adaptive_rate_limit_effect( context: LocalExecutionEffectContext<'_>, effect: LocalAdaptiveRateLimitEffect<'_>, ) { + if !local_candidate_failure_should_apply_key_effects( + &context.plan.provider_api_format, + effect.classification, + effect.status_code, + ) { + return; + } let Some(auth_config_fence) = capture_local_execution_auth_config_fence(state, context.plan).await else { @@ -964,6 +973,13 @@ async fn record_health_failure_effect( context: LocalExecutionEffectContext<'_>, effect: LocalHealthFailureEffect, ) { + if !local_candidate_failure_should_apply_key_effects( + &context.plan.provider_api_format, + effect.classification, + effect.status_code, + ) { + return; + } let api_format = context.plan.provider_api_format.trim(); if api_format.is_empty() { return; @@ -1231,6 +1247,13 @@ async fn record_pool_error_effect( context: LocalExecutionEffectContext<'_>, effect: LocalPoolErrorEffect<'_>, ) { + if !local_candidate_failure_should_apply_key_effects( + &context.plan.provider_api_format, + effect.classification, + effect.status_code, + ) { + return; + } let terminal_error_reason = admin_provider_pool_key_terminal_error_reason(effect.status_code, effect.error_body); if terminal_error_reason.is_none() @@ -1379,13 +1402,7 @@ async fn record_oauth_invalidation_effect( if !transport.key.auth_type.trim().eq_ignore_ascii_case("oauth") { return; } - if transport - .provider - .provider_type - .trim() - .eq_ignore_ascii_case("codex") - && !execution_plan_bearer_matches_transport(plan, &transport) - { + if !execution_plan_bearer_matches_transport(plan, &transport) { return; } @@ -1397,28 +1414,8 @@ async fn record_oauth_invalidation_effect( return; }; - let expected_auth_config = match state - .capture_provider_transport_auth_config_fence(&transport) - .await - { - Ok(Some(ciphertext)) => ciphertext, - Ok(None) => return, - Err(err) => { - warn!( - "gateway orchestration effects: failed to capture oauth invalidation fence for provider {} endpoint {} key {}: {:?}", - plan.provider_id, plan.endpoint_id, plan.key_id, err - ); - return; - } - }; - if let Err(err) = state - .mark_provider_catalog_key_oauth_invalid_fenced( - &plan.key_id, - transport.provider.provider_type.as_str(), - invalid_reason.as_str(), - expected_auth_config.as_str(), - ) + .mark_provider_transport_oauth_invalid_fenced(&transport, invalid_reason.as_str()) .await { warn!( @@ -1444,16 +1441,20 @@ fn execution_plan_bearer_matches_transport( plan: &ExecutionPlan, transport: &crate::provider_transport::GatewayProviderTransportSnapshot, ) -> bool { - let current_token = transport.key.decrypted_api_key.trim(); - !current_token.is_empty() - && plan.headers.iter().any(|(name, value)| { - name.eq_ignore_ascii_case("authorization") - && value - .trim() - .strip_prefix("Bearer ") - .map(str::trim) - .is_some_and(|token| token == current_token) - }) + let Some(plan_token) = execution_plan_authorization(plan).and_then(bearer_access_token) else { + return false; + }; + crate::provider_transport::resolve_local_generic_oauth_transport_authorization(transport) + .as_deref() + .and_then(bearer_access_token) + .is_some_and(|current_token| current_token == plan_token) +} + +fn bearer_access_token(authorization: &str) -> Option<&str> { + let mut parts = authorization.split_ascii_whitespace(); + let scheme = parts.next()?; + let token = parts.next()?; + (scheme.eq_ignore_ascii_case("bearer") && parts.next().is_none()).then_some(token) } fn resolve_local_oauth_invalid_reason( @@ -1467,6 +1468,12 @@ fn resolve_local_oauth_invalid_reason( status_code, upstream_message.as_deref(), ), + _ if super::oauth_status_may_be_invalid(status_code, response_text) => Some(format!( + "[OAUTH_EXPIRED] {}", + upstream_message + .as_deref() + .unwrap_or("OAuth access token was rejected") + )), _ => None, } } @@ -1492,6 +1499,46 @@ fn local_candidate_failure_should_invalidate_affinity( } } +fn local_candidate_failure_should_invalidate_affinity_for_provider( + provider_api_format: &str, + classification: LocalFailoverClassification, + status_code: u16, +) -> bool { + if !local_candidate_failure_should_invalidate_affinity(classification, status_code) { + return false; + } + if !provider_api_format + .trim() + .eq_ignore_ascii_case("claude:messages") + { + return true; + } + + let disposition = + classify_failure_disposition(provider_api_format, classification, status_code); + !(disposition.retry_action == crate::orchestration::FailureRetryAction::Stop + && disposition.failure_scope == FailureScope::None) +} + +fn local_candidate_failure_should_apply_key_effects( + provider_api_format: &str, + classification: LocalFailoverClassification, + status_code: u16, +) -> bool { + if !provider_api_format + .trim() + .eq_ignore_ascii_case("claude:messages") + { + return true; + } + + matches!( + classify_failure_disposition(provider_api_format, classification, status_code) + .failure_scope, + FailureScope::Credential + ) +} + fn local_candidate_failure_should_record_pool_error( classification: LocalFailoverClassification, status_code: u16, @@ -1708,12 +1755,13 @@ mod tests { use serde_json::{json, Value}; use super::{ - apply_local_execution_effect, local_candidate_failure_should_record_pool_error, - pool_score_feedback_gate_allows, pool_score_hard_state_for_status, - LocalAdaptiveRateLimitEffect, LocalAdaptiveSuccessEffect, LocalAttemptFailureEffect, - LocalExecutionEffect, LocalExecutionEffectContext, LocalHealthFailureEffect, - LocalHealthSuccessEffect, LocalOAuthInvalidationEffect, LocalPoolErrorEffect, - ProviderKeyEffectLockPool, + apply_local_execution_effect, execution_plan_bearer_matches_transport, + local_candidate_failure_should_apply_key_effects, + local_candidate_failure_should_record_pool_error, pool_score_feedback_gate_allows, + pool_score_hard_state_for_status, LocalAdaptiveRateLimitEffect, LocalAdaptiveSuccessEffect, + LocalAttemptFailureEffect, LocalExecutionEffect, LocalExecutionEffectContext, + LocalHealthFailureEffect, LocalHealthSuccessEffect, LocalOAuthInvalidationEffect, + LocalPoolErrorEffect, ProviderKeyEffectLockPool, }; use crate::data::{GatewayDataConfig, GatewayDataState}; use crate::orchestration::LocalFailoverClassification; @@ -1760,6 +1808,13 @@ mod tests { } } + fn sample_claude_plan() -> ExecutionPlan { + let mut plan = sample_plan(); + plan.provider_name = Some("anthropic".to_string()); + plan.provider_api_format = "claude:messages".to_string(); + plan + } + #[test] fn pool_score_feedback_gate_suppresses_repeated_success_writes() { super::POOL_SCORE_FEEDBACK_GATE.clear(); @@ -1864,7 +1919,7 @@ mod tests { url: "https://chatgpt.com/backend-api/codex".to_string(), headers: BTreeMap::from([( "authorization".to_string(), - "Bearer __placeholder__".to_string(), + "Bearer codex-access-token".to_string(), )]), content_type: Some("application/json".to_string()), content_encoding: None, @@ -1959,8 +2014,8 @@ mod tests { .expect("key should build") .with_transport_fields( Some(serde_json::json!(["openai:responses"])), - encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "__placeholder__") - .expect("placeholder api key should encrypt"), + encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "codex-access-token") + .expect("access token should encrypt"), Some(encrypted_auth_config), None, Some(serde_json::json!({"openai:responses": 1})), @@ -1975,6 +2030,10 @@ mod tests { fn sample_codex_agent_identity_key() -> StoredProviderCatalogKey { let mut key = sample_codex_key(); key.name = "Agent Identity".to_string(); + key.encrypted_api_key = Some( + encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "__placeholder__") + .expect("placeholder api key should encrypt"), + ); key.encrypted_auth_config = Some( encrypt_python_fernet_plaintext( DEVELOPMENT_ENCRYPTION_KEY, @@ -2014,6 +2073,36 @@ mod tests { ) } + fn claude_code_oauth_state() -> AppState { + let mut provider = sample_codex_provider(); + provider.name = "claude_code".to_string(); + provider.provider_type = "claude_code".to_string(); + let mut endpoint = sample_codex_endpoint(); + endpoint.api_format = "claude:messages".to_string(); + endpoint.api_family = Some("claude".to_string()); + endpoint.base_url = "https://api.anthropic.com".to_string(); + let mut key = sample_codex_key(); + key.api_formats = Some(json!(["claude:messages"])); + key.encrypted_auth_config = Some( + encrypt_python_fernet_plaintext( + DEVELOPMENT_ENCRYPTION_KEY, + r#"{"provider_type":"claude_code","refresh_token":"rt-claude-local-123"}"#, + ) + .expect("Claude Code auth config should encrypt"), + ); + let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( + vec![provider], + vec![endpoint], + vec![key], + )); + AppState::new() + .expect("gateway state should build") + .with_data_state_for_tests( + GatewayDataState::with_provider_catalog_repository_for_tests(repository) + .with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY), + ) + } + fn codex_state_with_redis(redis_url: &str, redis_key_prefix: &str) -> AppState { let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( vec![sample_codex_provider()], @@ -2724,6 +2813,179 @@ mod tests { )); } + #[test] + fn anthropic_non_credential_failures_do_not_apply_key_wide_effects() { + assert!(!local_candidate_failure_should_apply_key_effects( + "claude:messages", + LocalFailoverClassification::RetryUpstreamFailure, + 529, + )); + assert!(!local_candidate_failure_should_apply_key_effects( + "claude:messages", + LocalFailoverClassification::RetryUpstreamFailure, + 429, + )); + assert!(!local_candidate_failure_should_apply_key_effects( + "claude:messages", + LocalFailoverClassification::RetryUpstreamFailure, + 503, + )); + assert!(local_candidate_failure_should_apply_key_effects( + "claude:messages", + LocalFailoverClassification::RetryUpstreamFailure, + 401, + )); + assert!(local_candidate_failure_should_apply_key_effects( + "claude:messages", + LocalFailoverClassification::RetryUpstreamFailure, + 403, + )); + assert!(!local_candidate_failure_should_apply_key_effects( + "claude:messages", + LocalFailoverClassification::RetryUpstreamFailure, + 400, + )); + assert!(local_candidate_failure_should_apply_key_effects( + "openai:chat", + LocalFailoverClassification::RetryUpstreamFailure, + 529, + )); + assert!(local_candidate_failure_should_apply_key_effects( + "openai:chat", + LocalFailoverClassification::RetryUpstreamFailure, + 429, + )); + assert!(local_candidate_failure_should_apply_key_effects( + "openai:chat", + LocalFailoverClassification::RetryUpstreamFailure, + 503, + )); + } + + #[tokio::test] + async fn anthropic_non_credential_failures_preserve_key_wide_state() { + for status_code in [400, 429, 503, 529] { + let mut key = sample_adaptive_key(); + let circuit = json!({ + "openai:chat": { + "open": true, + "reason": "existing-state" + } + }); + key.circuit_breaker_by_format = Some(circuit.clone()); + let expected_adaptive_state = ProviderCatalogKeyAdaptiveState::from(&key); + let expected_health = key.health_by_format.clone(); + let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( + vec![sample_pool_health_provider()], + vec![sample_health_endpoint()], + vec![key], + )); + let state = AppState::new() + .expect("gateway state should build") + .with_data_state_for_tests( + GatewayDataState::with_provider_catalog_repository_for_tests(repository) + .with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY), + ); + let plan = sample_claude_plan(); + let report_context = json!({ + "api_key_id": "api-key-1", + "client_api_format": "claude:messages", + "model": "claude-sonnet-4-5", + }); + let cache_key = build_scheduler_affinity_cache_key_for_api_key_id( + "api-key-1", + "claude:messages", + "claude-sonnet-4-5", + ) + .expect("scheduler affinity cache key should build"); + let target = SchedulerAffinityTarget { + provider_id: plan.provider_id.clone(), + endpoint_id: plan.endpoint_id.clone(), + key_id: plan.key_id.clone(), + }; + state.remember_scheduler_affinity_target( + &cache_key, + target.clone(), + SCHEDULER_AFFINITY_TTL, + 16, + ); + let headers = BTreeMap::from([("Retry-After".to_string(), "120".to_string())]); + let context = LocalExecutionEffectContext { + plan: &plan, + report_context: Some(&report_context), + }; + let classification = LocalFailoverClassification::RetryUpstreamFailure; + + apply_local_execution_effect( + &state, + context, + LocalExecutionEffect::AttemptFailure(LocalAttemptFailureEffect { + status_code, + classification, + }), + ) + .await; + apply_local_execution_effect( + &state, + context, + LocalExecutionEffect::AdaptiveRateLimit(LocalAdaptiveRateLimitEffect { + status_code, + classification, + headers: Some(&headers), + }), + ) + .await; + apply_local_execution_effect( + &state, + context, + LocalExecutionEffect::HealthFailure(LocalHealthFailureEffect { + status_code, + classification, + }), + ) + .await; + apply_local_execution_effect( + &state, + context, + LocalExecutionEffect::PoolError(LocalPoolErrorEffect { + status_code, + classification, + headers: &headers, + error_body: Some(r#"{"error":{"message":"temporarily unavailable"}}"#), + }), + ) + .await; + + let stored_key = state + .read_provider_catalog_keys_by_ids(std::slice::from_ref(&plan.key_id)) + .await + .expect("provider catalog keys should load") + .into_iter() + .next() + .expect("stored key should exist"); + assert_eq!( + ProviderCatalogKeyAdaptiveState::from(&stored_key), + expected_adaptive_state, + "Anthropic status {status_code} must not update key-wide adaptive state" + ); + assert_eq!( + stored_key.health_by_format, expected_health, + "Anthropic status {status_code} must not update key-wide health" + ); + assert_eq!( + stored_key.circuit_breaker_by_format, + Some(circuit), + "Anthropic status {status_code} must not clear pool key state" + ); + let expected_affinity = (status_code == 400).then_some(target); + assert_eq!( + state.read_scheduler_affinity_target(&cache_key, SCHEDULER_AFFINITY_TTL), + expected_affinity, + "Anthropic status {status_code} must invalidate only retryable target affinity" + ); + } + } + #[test] fn terminal_pool_account_errors_project_pool_hard_state() { assert_eq!( @@ -2790,6 +3052,162 @@ mod tests { assert_eq!(stored_key.circuit_breaker_by_format, None); } + #[tokio::test] + async fn oauth_bearer_generation_match_supports_generic_auth_config_token() { + let state = codex_state(); + let mut transport = state + .read_provider_transport_snapshot( + "provider-codex-cli-local-1", + "endpoint-codex-cli-local-1", + "key-codex-cli-local-1", + ) + .await + .expect("transport should load") + .expect("transport should exist"); + transport.provider.provider_type = "claude_code".to_string(); + transport.key.decrypted_api_key = "__placeholder__".to_string(); + transport.key.decrypted_auth_config = + Some(json!({"accessToken": "current-access-token"}).to_string()); + let mut plan = sample_codex_plan(); + plan.headers.insert( + "authorization".to_string(), + "Bearer current-access-token".to_string(), + ); + assert!(execution_plan_bearer_matches_transport(&plan, &transport)); + + plan.headers.insert( + "authorization".to_string(), + "Bearer stale-access-token".to_string(), + ); + assert!(!execution_plan_bearer_matches_transport(&plan, &transport)); + + transport.key.decrypted_api_key = "replacement-access-token".to_string(); + plan.headers.insert( + "authorization".to_string(), + "Bearer current-access-token".to_string(), + ); + assert!(!execution_plan_bearer_matches_transport(&plan, &transport)); + plan.headers.insert( + "authorization".to_string(), + "Bearer replacement-access-token".to_string(), + ); + assert!(execution_plan_bearer_matches_transport(&plan, &transport)); + + transport.key.decrypted_auth_config = Some( + json!({ + "accessToken": "current-access-token", + "request": { + "extraHeaders": { + "Authorization": "Bearer nested-override-token" + } + } + }) + .to_string(), + ); + plan.headers.insert( + "authorization".to_string(), + "Bearer replacement-access-token".to_string(), + ); + assert!(!execution_plan_bearer_matches_transport(&plan, &transport)); + plan.headers.insert( + "authorization".to_string(), + "Bearer nested-override-token".to_string(), + ); + assert!(execution_plan_bearer_matches_transport(&plan, &transport)); + } + + #[tokio::test] + async fn oauth_invalidation_marks_claude_code_authentication_failures_only() { + let state = claude_code_oauth_state(); + let mut plan = sample_codex_plan(); + plan.provider_name = Some("claude_code".to_string()); + plan.provider_api_format = "claude:messages".to_string(); + + apply_local_execution_effect( + &state, + LocalExecutionEffectContext { + plan: &plan, + report_context: None, + }, + LocalExecutionEffect::OauthInvalidation(LocalOAuthInvalidationEffect { + status_code: 403, + response_text: Some( + r#"{"type":"error","error":{"type":"permission_error","message":"insufficient scope"}}"#, + ), + }), + ) + .await; + let unmarked = state + .read_provider_catalog_keys_by_ids(std::slice::from_ref(&plan.key_id)) + .await + .expect("provider catalog keys should load") + .into_iter() + .next() + .expect("stored key should exist"); + assert!(unmarked.oauth_invalid_at_unix_secs.is_none()); + + apply_local_execution_effect( + &state, + LocalExecutionEffectContext { + plan: &plan, + report_context: None, + }, + LocalExecutionEffect::OauthInvalidation(LocalOAuthInvalidationEffect { + status_code: 403, + response_text: Some( + r#"{"type":"error","error":{"type":"authentication_error","message":"invalid access token"}}"#, + ), + }), + ) + .await; + let marked = state + .read_provider_catalog_keys_by_ids(std::slice::from_ref(&plan.key_id)) + .await + .expect("provider catalog keys should load") + .into_iter() + .next() + .expect("stored key should exist"); + assert!(marked.oauth_invalid_at_unix_secs.is_some()); + assert_eq!( + marked.oauth_invalid_reason.as_deref(), + Some("[OAUTH_EXPIRED] invalid access token") + ); + } + + #[tokio::test] + async fn oauth_invalidation_marks_claude_code_unauthorized_without_body() { + let state = claude_code_oauth_state(); + let mut plan = sample_codex_plan(); + plan.provider_name = Some("claude_code".to_string()); + plan.provider_api_format = "claude:messages".to_string(); + + apply_local_execution_effect( + &state, + LocalExecutionEffectContext { + plan: &plan, + report_context: None, + }, + LocalExecutionEffect::OauthInvalidation(LocalOAuthInvalidationEffect { + status_code: 401, + response_text: None, + }), + ) + .await; + + let stored_key = state + .read_provider_catalog_keys_by_ids(std::slice::from_ref(&plan.key_id)) + .await + .expect("provider catalog keys should load") + .into_iter() + .next() + .expect("stored key should exist"); + assert!(stored_key.oauth_invalid_at_unix_secs.is_some()); + assert_eq!( + stored_key.oauth_invalid_reason.as_deref(), + Some("[OAUTH_EXPIRED] OAuth access token was rejected") + ); + } + #[tokio::test] async fn oauth_invalidation_marks_codex_key_expired() { let state = codex_state(); diff --git a/apps/aether-gateway/src/orchestration/mod.rs b/apps/aether-gateway/src/orchestration/mod.rs index cd4c838d9..b9782b319 100644 --- a/apps/aether-gateway/src/orchestration/mod.rs +++ b/apps/aether-gateway/src/orchestration/mod.rs @@ -9,6 +9,7 @@ mod attempt; mod classifier; mod effects; mod health; +mod oauth_error; mod policy; mod recovery; mod report_effects; @@ -24,8 +25,10 @@ pub(crate) use self::attempt::{ LocalExecutionCandidateMetadata, SCHEDULER_AFFINITY_EPOCH_REPORT_FIELD, }; pub(crate) use self::classifier::{ - classify_local_failover, local_failover_error_message, LocalFailoverClassification, - LocalFailoverInput, + classify_anthropic_failure_disposition, classify_failure_disposition, classify_local_failover, + failure_disposition_from_local_classification, local_failover_error_message, + FailureDisposition, FailureRetryAction, FailureScope, FailureTokenAction, + LocalFailoverClassification, LocalFailoverInput, }; pub(crate) use self::effects::{ apply_local_execution_effect, LocalAdaptiveRateLimitEffect, LocalAdaptiveSuccessEffect, @@ -37,6 +40,9 @@ pub(crate) use self::health::{ project_local_failure_health, project_local_key_circuit_closed, project_local_key_circuit_failure, project_local_success_health, }; +pub(crate) use self::oauth_error::{ + oauth_status_may_be_invalid, oauth_status_proves_access_token_invalid, +}; pub(crate) use self::policy::{ append_local_failover_policy_to_value, codex_cyber_flag_passthrough_enabled, cyber_continue_failover_enabled, local_failover_policy_from_report_context, @@ -44,8 +50,8 @@ pub(crate) use self::policy::{ LocalFailoverRegexRule, CYBER_CONTINUE_FAILOVER_CONFIG_KEY, }; pub(crate) use self::recovery::{ - analyze_local_failover, recover_local_failover_decision, LocalFailoverAnalysis, - LocalFailoverDecision, + analyze_local_failover, apply_provider_failure_disposition, recover_local_failover_decision, + LocalFailoverAnalysis, LocalFailoverDecision, }; #[cfg(test)] pub(crate) use self::report_effects::clear_local_report_effect_caches_for_tests; @@ -65,7 +71,9 @@ pub(crate) async fn resolve_local_failover_analysis_for_attempt( } let policy = resolve_local_failover_policy(state, plan, report_context).await; - analyze_local_failover(&policy, LocalFailoverInput::new(status_code, response_text)) + let analysis = + analyze_local_failover(&policy, LocalFailoverInput::new(status_code, response_text)); + apply_provider_failure_disposition(&plan.provider_api_format, status_code, analysis) } pub(crate) async fn resolve_local_failover_decision_for_attempt( diff --git a/apps/aether-gateway/src/orchestration/oauth_error.rs b/apps/aether-gateway/src/orchestration/oauth_error.rs new file mode 100644 index 000000000..f7a2618e0 --- /dev/null +++ b/apps/aether-gateway/src/orchestration/oauth_error.rs @@ -0,0 +1,113 @@ +pub(crate) fn oauth_status_may_be_invalid(status_code: u16, response_text: Option<&str>) -> bool { + if status_code == 401 { + return true; + } + if status_code != 403 { + return false; + } + + let Some(response_text) = response_text else { + return false; + }; + if let Ok(body) = serde_json::from_str::(response_text) { + let error_type = body + .get("error") + .and_then(|error| error.get("type")) + .and_then(serde_json::Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .or_else(|| { + body.get("type") + .and_then(serde_json::Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty() && !value.eq_ignore_ascii_case("error")) + }); + if let Some(error_type) = error_type { + return is_oauth_invalid_error_taxonomy(error_type); + } + + let error_code = body + .get("error") + .and_then(|error| error.get("code")) + .or_else(|| body.get("code")) + .or_else(|| body.get("error").filter(|error| error.is_string())) + .and_then(serde_json::Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()); + if error_code.is_some_and(is_oauth_invalid_error_taxonomy) { + return true; + } + + return response_has_oauth_invalid_phrase(response_text); + } + + response_has_oauth_invalid_phrase(response_text) +} + +pub(crate) fn oauth_status_proves_access_token_invalid( + status_code: u16, + response_text: Option<&str>, +) -> bool { + if status_code == 401 { + return true; + } + if status_code != 403 { + return false; + } + + response_text.is_some_and(response_has_oauth_invalid_phrase) +} + +fn is_oauth_invalid_error_taxonomy(value: &str) -> bool { + matches!( + value.trim().to_ascii_lowercase().as_str(), + "authentication_error" + | "invalid_authentication_token" + | "invalid_token" + | "oauth_token_invalid" + | "token_invalid" + | "token_expired" + | "unauthenticated" + | "biscuit_baker_service_auth_credential_error_status" + ) +} + +fn response_has_oauth_invalid_phrase(response_text: &str) -> bool { + let response_text = response_text.to_ascii_lowercase(); + if [ + "oauth_token_invalid", + "invalid_token", + "biscuit_baker_service_auth_credential_error_status", + ] + .iter() + .any(|taxonomy| contains_ascii_taxonomy_token(&response_text, taxonomy)) + { + return true; + } + + [ + "oauth token is invalid", + "oauth token is expired", + "oauth token has expired", + "invalid access token", + "access token invalid", + "access token expired", + "expired access token", + "authentication token has been invalidated", + "token has been invalidated", + "personal access token owner is inactive", + "security token included in the request is expired", + ] + .iter() + .any(|needle| response_text.contains(needle)) +} + +fn contains_ascii_taxonomy_token(text: &str, taxonomy: &str) -> bool { + text.match_indices(taxonomy).any(|(start, matched)| { + let end = start + matched.len(); + let is_identifier_byte = |byte: u8| byte.is_ascii_alphanumeric() || byte == b'_'; + let has_left_boundary = start == 0 || !is_identifier_byte(text.as_bytes()[start - 1]); + let has_right_boundary = end == text.len() || !is_identifier_byte(text.as_bytes()[end]); + has_left_boundary && has_right_boundary + }) +} diff --git a/apps/aether-gateway/src/orchestration/recovery.rs b/apps/aether-gateway/src/orchestration/recovery.rs index eecbf3ed8..d6266eede 100644 --- a/apps/aether-gateway/src/orchestration/recovery.rs +++ b/apps/aether-gateway/src/orchestration/recovery.rs @@ -1,4 +1,7 @@ -use super::classifier::{classify_local_failover, LocalFailoverClassification, LocalFailoverInput}; +use super::classifier::{ + classify_failure_disposition, classify_local_failover, FailureRetryAction, + LocalFailoverClassification, LocalFailoverInput, +}; use super::LocalFailoverPolicy; #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -44,6 +47,37 @@ pub(crate) fn analyze_local_failover( } } +pub(crate) fn apply_provider_failure_disposition( + provider_api_format: &str, + status_code: u16, + analysis: LocalFailoverAnalysis, +) -> LocalFailoverAnalysis { + if status_code < 400 + && matches!( + analysis.classification, + LocalFailoverClassification::UseDefault + ) + { + return analysis; + } + + let disposition = + classify_failure_disposition(provider_api_format, analysis.classification, status_code); + let decision = match disposition.retry_action { + FailureRetryAction::Stop | FailureRetryAction::SameCredential => { + LocalFailoverDecision::StopLocalFailover + } + FailureRetryAction::NextCandidate + | FailureRetryAction::NextCredential + | FailureRetryAction::NextEndpoint => LocalFailoverDecision::RetryNextCandidate, + }; + + LocalFailoverAnalysis { + classification: analysis.classification, + decision, + } +} + pub(crate) fn recover_local_failover_decision( policy: &LocalFailoverPolicy, input: LocalFailoverInput<'_>, @@ -70,7 +104,10 @@ const fn decision_from_classification( #[cfg(test)] mod tests { - use super::{analyze_local_failover, recover_local_failover_decision, LocalFailoverDecision}; + use super::{ + analyze_local_failover, apply_provider_failure_disposition, + recover_local_failover_decision, LocalFailoverAnalysis, LocalFailoverDecision, + }; use crate::orchestration::{ LocalFailoverClassification, LocalFailoverInput, LocalFailoverPolicy, }; @@ -161,4 +198,44 @@ mod tests { LocalFailoverClassification::StopCyberPolicy ); } + + #[test] + fn anthropic_failure_disposition_controls_candidate_retry() { + let policy = LocalFailoverPolicy::default(); + + for status_code in [400, 413] { + let analysis = analyze_local_failover( + &policy, + LocalFailoverInput::new(status_code, Some(r#"{"error":{"message":"failed"}}"#)), + ); + assert_eq!( + apply_provider_failure_disposition("claude:messages", status_code, analysis,) + .decision, + LocalFailoverDecision::StopLocalFailover, + "Anthropic status {status_code} must not blindly rotate credentials" + ); + } + + for status_code in [401, 403, 404, 429, 529] { + let analysis = analyze_local_failover( + &policy, + LocalFailoverInput::new(status_code, Some(r#"{"error":{"message":"failed"}}"#)), + ); + assert_eq!( + apply_provider_failure_disposition("claude:messages", status_code, analysis,) + .decision, + LocalFailoverDecision::RetryNextCandidate, + "Anthropic status {status_code} should continue candidate failover" + ); + } + } + + #[test] + fn provider_failure_disposition_preserves_non_failure_default() { + let analysis = LocalFailoverAnalysis::use_default(); + assert_eq!( + apply_provider_failure_disposition("claude:messages", 200, analysis).decision, + LocalFailoverDecision::UseDefault + ); + } } diff --git a/apps/aether-gateway/src/rate_limit.rs b/apps/aether-gateway/src/rate_limit.rs index 39f7541d7..411513dfd 100644 --- a/apps/aether-gateway/src/rate_limit.rs +++ b/apps/aether-gateway/src/rate_limit.rs @@ -502,6 +502,9 @@ mod tests { route_class: Some("ai_public".to_string()), route_family: Some("openai".to_string()), route_kind: Some("chat".to_string()), + client_surface: None, + api_operation: None, + gateway_credential_carrier: None, request_auth_channel: None, auth_endpoint_signature: Some("openai:chat".to_string()), execution_runtime_candidate: true, diff --git a/apps/aether-gateway/src/state/oauth.rs b/apps/aether-gateway/src/state/oauth.rs index 56df17434..5b7baac80 100644 --- a/apps/aether-gateway/src/state/oauth.rs +++ b/apps/aether-gateway/src/state/oauth.rs @@ -17,8 +17,8 @@ use aether_contracts::{ EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER, EXECUTION_REQUEST_HTTP1_ONLY_HEADER, }; use aether_data_contracts::repository::provider_catalog::{ - ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyStatusSnapshotUpdate, - StoredProviderCatalogKey, + ProviderCatalogKeyOAuthCredentialFence, ProviderCatalogKeyOAuthRuntimeStateCasUpdate, + ProviderCatalogKeyStatusSnapshotUpdate, StoredProviderCatalogKey, }; use aether_runtime_state::RuntimeLockLease; use base64::{engine::general_purpose::STANDARD, Engine as _}; @@ -42,6 +42,12 @@ const OAUTH_EXPIRED_PREFIX: &str = "[OAUTH_EXPIRED] "; const OAUTH_REFRESH_FAILED_PREFIX: &str = "[REFRESH_FAILED] "; const OAUTH_REQUEST_FAILED_PREFIX: &str = "[REQUEST_FAILED] "; +#[derive(Debug, Clone, PartialEq)] +struct ProviderTransportCredentialFence { + encrypted_auth_config: String, + credential: ProviderCatalogKeyOAuthCredentialFence, +} + struct GatewayLocalOAuthHttpExecutor<'a> { state: &'a AppState, } @@ -173,32 +179,6 @@ fn local_oauth_transport_context_allows_reload( initial, current, ); } - if initial - .provider - .provider_type - .trim() - .eq_ignore_ascii_case("codex") - && initial.key.auth_type.trim().eq_ignore_ascii_case("oauth") - { - let initial_config = initial - .key - .decrypted_auth_config - .as_deref() - .and_then(|value| serde_json::from_str::(value).ok()); - let current_config = current - .key - .decrypted_auth_config - .as_deref() - .and_then(|value| serde_json::from_str::(value).ok()); - return current - .provider - .provider_type - .trim() - .eq_ignore_ascii_case("codex") - && current.key.auth_type.trim().eq_ignore_ascii_case("oauth") - && initial_config == current_config - && initial.key.decrypted_api_key == current.key.decrypted_api_key; - } true } @@ -1380,7 +1360,7 @@ impl AppState { { return Ok(None); } - let expected_auth_config = if current_transport + let expected_credential_fence = if current_transport .key .decrypted_auth_config .as_deref() @@ -1388,10 +1368,10 @@ impl AppState { .is_some_and(|value| !value.is_empty()) { match self - .capture_provider_transport_auth_config_fence(¤t_transport) + .capture_provider_transport_credential_fence(¤t_transport) .await? { - Some(ciphertext) => Some(ciphertext), + Some(fence) => Some(fence), None => { let Some(reloaded) = self .read_provider_transport_snapshot_uncached( @@ -1485,7 +1465,7 @@ impl AppState { .persist_local_oauth_refresh_entry( ¤t_transport, &refreshed_entry, - expected_auth_config.as_deref(), + expected_credential_fence.as_ref(), ) .await { @@ -1533,9 +1513,9 @@ impl AppState { let lock_owner = format!("aether-gateway-admin-{}", std::process::id()); let initial_transport = transport.clone(); let mut current_transport = transport.clone(); - current_transport.key.decrypted_api_key = "__placeholder__".to_string(); - let expected_refresh_fingerprint = - provider_transport::codex_agent_identity_refresh_fingerprint(¤t_transport, None); + let expected_refresh_fingerprint = self + .oauth_refresh + .refresh_fingerprint_for_transport(&initial_transport); let executor = GatewayLocalOAuthHttpExecutor { state: self }; let transport_refresh_token_fingerprint = oauth_auth_config_refresh_token_fingerprint( current_transport.key.decrypted_auth_config.as_deref(), @@ -1560,8 +1540,8 @@ impl AppState { { return Ok(None); } - let expected_auth_config = match self - .capture_provider_transport_auth_config_fence(¤t_transport) + let expected_credential_fence = match self + .capture_provider_transport_credential_fence(¤t_transport) .await .map_err( |err| provider_transport::LocalOAuthRefreshError::InvalidResponse { @@ -1569,7 +1549,7 @@ impl AppState { message: format!("{err:?}"), }, )? { - Some(ciphertext) => Some(ciphertext), + Some(fence) => Some(fence), None if current_transport.key.decrypted_auth_config.is_some() => { let Some(reloaded) = self .read_provider_transport_snapshot_uncached( @@ -1588,7 +1568,6 @@ impl AppState { return Ok(None); }; current_transport = reloaded; - current_transport.key.decrypted_api_key = "__placeholder__".to_string(); continue; } None => None, @@ -1621,7 +1600,6 @@ impl AppState { continue; }; current_transport = reloaded_transport; - current_transport.key.decrypted_api_key = "__placeholder__".to_string(); continue; } @@ -1684,7 +1662,7 @@ impl AppState { .persist_local_oauth_refresh_entry( ¤t_transport, &refreshed_entry, - expected_auth_config.as_deref(), + expected_credential_fence.as_ref(), ) .await { @@ -1772,6 +1750,16 @@ impl AppState { &self, transport: &provider_transport::GatewayProviderTransportSnapshot, ) -> Result, GatewayError> { + Ok(self + .capture_provider_transport_credential_fence(transport) + .await? + .map(|fence| fence.encrypted_auth_config)) + } + + async fn capture_provider_transport_credential_fence( + &self, + transport: &provider_transport::GatewayProviderTransportSnapshot, + ) -> Result, GatewayError> { let key_id = transport.key.id.trim(); let stored = self .data @@ -1780,11 +1768,51 @@ impl AppState { .map_err(|err| GatewayError::Internal(err.to_string()))? .into_iter() .next(); - let Some(ciphertext) = stored.and_then(|key| key.encrypted_auth_config) else { + let Some(stored) = stored else { + return Ok(None); + }; + if stored.provider_id != transport.provider.id + || stored.auth_type != transport.key.auth_type + { + return Ok(None); + } + let provider = self + .data + .list_provider_catalog_providers_by_ids(std::slice::from_ref(&stored.provider_id)) + .await + .map_err(|err| GatewayError::Internal(err.to_string()))? + .into_iter() + .next(); + let Some(provider) = provider else { + return Ok(None); + }; + if provider.provider_type != transport.provider.provider_type { + return Ok(None); + } + + let stored_api_key = match stored.encrypted_api_key.as_deref() { + Some(ciphertext) => Some( + decrypt_catalog_secret_with_fallbacks(self.data.encryption_key(), ciphertext) + .ok_or_else(|| { + GatewayError::Internal( + "provider api_key could not be verified for runtime fencing" + .to_string(), + ) + })?, + ), + None => None, + }; + let transport_api_key = (!transport.key.decrypted_api_key.is_empty()) + .then_some(transport.key.decrypted_api_key.as_str()); + if stored_api_key.as_deref() != transport_api_key { + return Ok(None); + } + + let Some(ciphertext) = stored.encrypted_auth_config.as_deref() else { return Ok(None); }; let plaintext = - decrypt_catalog_secret_with_fallbacks(self.data.encryption_key(), ciphertext.as_str()) + decrypt_catalog_secret_with_fallbacks(self.data.encryption_key(), ciphertext) .ok_or_else(|| { GatewayError::Internal( "provider auth_config could not be verified for runtime fencing" @@ -1810,7 +1838,15 @@ impl AppState { if config != transport_config { return Ok(None); } - Ok(Some(ciphertext)) + Ok(Some(ProviderTransportCredentialFence { + encrypted_auth_config: ciphertext.to_string(), + credential: ProviderCatalogKeyOAuthCredentialFence { + encrypted_api_key: stored.encrypted_api_key, + auth_type: stored.auth_type, + provider_id: stored.provider_id, + provider_type: provider.provider_type, + }, + })) } pub(crate) async fn mark_provider_catalog_key_oauth_invalid( @@ -1883,18 +1919,23 @@ impl AppState { Ok(updated) } - pub(crate) async fn mark_provider_catalog_key_oauth_invalid_fenced( + pub(crate) async fn mark_provider_transport_oauth_invalid_fenced( &self, - key_id: &str, - provider_type: &str, + transport: &provider_transport::GatewayProviderTransportSnapshot, invalid_reason: &str, - expected_encrypted_auth_config: &str, ) -> Result { let invalid_reason = invalid_reason.trim(); - let expected_encrypted_auth_config = expected_encrypted_auth_config.trim(); - if invalid_reason.is_empty() || expected_encrypted_auth_config.is_empty() { + let key_id = transport.key.id.trim(); + let provider_type = transport.provider.provider_type.as_str(); + if invalid_reason.is_empty() || key_id.is_empty() { return Ok(false); } + let Some(expected_credential_fence) = self + .capture_provider_transport_credential_fence(transport) + .await? + else { + return Ok(false); + }; let Some(mut latest_key) = self .data @@ -1906,7 +1947,12 @@ impl AppState { else { return Ok(false); }; - if latest_key.encrypted_auth_config.as_deref() != Some(expected_encrypted_auth_config) + if latest_key.encrypted_auth_config.as_deref() + != Some(expected_credential_fence.encrypted_auth_config.as_str()) + || latest_key.encrypted_api_key + != expected_credential_fence.credential.encrypted_api_key + || latest_key.auth_type != expected_credential_fence.credential.auth_type + || latest_key.provider_id != expected_credential_fence.credential.provider_id || !provider_key_is_oauth_managed(&latest_key, provider_type) { return Ok(false); @@ -1940,9 +1986,10 @@ impl AppState { &ProviderCatalogKeyOAuthRuntimeStateCasUpdate { key_id: key_id.to_string(), expected_encrypted_auth_config: Some( - expected_encrypted_auth_config.to_string(), + expected_credential_fence.encrypted_auth_config.clone(), ), - encrypted_auth_config: expected_encrypted_auth_config.to_string(), + expected_credential: Some(expected_credential_fence.credential), + encrypted_auth_config: expected_credential_fence.encrypted_auth_config, encrypted_api_key_update: None, expires_at_unix_secs_update: None, oauth_invalid_at_unix_secs: latest_key.oauth_invalid_at_unix_secs, @@ -1965,7 +2012,7 @@ impl AppState { &self, transport: &provider_transport::GatewayProviderTransportSnapshot, entry: &provider_transport::CachedOAuthEntry, - expected_auth_config: Option<&str>, + expected_credential_fence: Option<&ProviderTransportCredentialFence>, ) -> Result<(), GatewayError> { let key_id = transport.key.id.trim(); if key_id.is_empty() { @@ -1973,6 +2020,19 @@ impl AppState { } if local_oauth_refresh_entry_should_stay_memory_only(transport, entry) { + let expected_credential_fence = expected_credential_fence.ok_or_else(|| { + GatewayError::Internal( + "memory-only OAuth refresh has no starting credential fence".to_string(), + ) + })?; + let current_credential_fence = self + .capture_provider_transport_credential_fence(transport) + .await?; + if current_credential_fence.as_ref() != Some(expected_credential_fence) { + return Err(GatewayError::Internal( + "OAuth credential changed while memory-only refresh was in flight".to_string(), + )); + } tracing::info!( key_id = %key_id, provider_id = %transport.provider.id, @@ -2034,18 +2094,22 @@ impl AppState { else { return Ok(()); }; + let expected_credential_fence = expected_credential_fence.ok_or_else(|| { + GatewayError::Internal( + "Agent Identity task registration has no starting credential fence".to_string(), + ) + })?; let expected_encrypted_auth_config = - expected_auth_config.map(str::to_string).ok_or_else(|| { - GatewayError::Internal( - "Agent Identity task registration has no starting auth_config fence" - .to_string(), - ) - })?; + expected_credential_fence.encrypted_auth_config.clone(); if latest_key.encrypted_auth_config.as_deref() != Some(expected_encrypted_auth_config.as_str()) + || latest_key.encrypted_api_key + != expected_credential_fence.credential.encrypted_api_key + || latest_key.auth_type != expected_credential_fence.credential.auth_type + || latest_key.provider_id != expected_credential_fence.credential.provider_id { return Err(GatewayError::Internal( - "Agent Identity auth_config changed while task registration was in flight" + "Agent Identity credential changed while task registration was in flight" .to_string(), )); } @@ -2093,6 +2157,7 @@ impl AppState { &ProviderCatalogKeyOAuthRuntimeStateCasUpdate { key_id: key_id.to_string(), expected_encrypted_auth_config: Some(expected_encrypted_auth_config), + expected_credential: Some(expected_credential_fence.credential.clone()), encrypted_auth_config, encrypted_api_key_update: None, expires_at_unix_secs_update: None, @@ -2147,17 +2212,13 @@ impl AppState { .map(|value| encrypt_python_fernet_plaintext(encryption_key, value.as_str())) .transpose() .map_err(|err| GatewayError::Internal(err.to_string()))?; - let requires_fenced_persistence = transport - .provider - .provider_type - .trim() - .eq_ignore_ascii_case("codex") - && transport.key.auth_type.trim().eq_ignore_ascii_case("oauth"); + let requires_fenced_persistence = + provider_transport::supports_local_oauth_request_auth_resolution(transport); if requires_fenced_persistence - && (expected_auth_config.is_none() || encrypted_auth_config.is_none()) + && (expected_credential_fence.is_none() || encrypted_auth_config.is_none()) { return Err(GatewayError::Internal( - "Codex OAuth refresh persistence is missing its auth_config fence".to_string(), + "OAuth refresh persistence is missing its credential fence".to_string(), )); } @@ -2172,7 +2233,13 @@ impl AppState { return Ok(()); }; - let observed_encrypted_auth_config = latest_key.encrypted_auth_config.clone(); + let observed_credential_matches = expected_credential_fence.is_none_or(|expected| { + latest_key.encrypted_auth_config.as_deref() + == Some(expected.encrypted_auth_config.as_str()) + && latest_key.encrypted_api_key == expected.credential.encrypted_api_key + && latest_key.auth_type == expected.credential.auth_type + && latest_key.provider_id == expected.credential.provider_id + }); latest_key.encrypted_api_key = Some(encrypted_api_key.clone()); latest_key.encrypted_auth_config = encrypted_auth_config.clone(); latest_key.expires_at_unix_secs = entry.expires_at_unix_secs; @@ -2191,17 +2258,20 @@ impl AppState { latest_key.status_snapshot = sync_provider_key_oauth_status_snapshot(current_status_snapshot, &latest_key); let used_fenced_persistence = - expected_auth_config.is_some() && encrypted_auth_config.is_some(); - let updated = if let (Some(expected_auth_config), Some(encrypted_auth_config)) = - (expected_auth_config, encrypted_auth_config.as_deref()) + expected_credential_fence.is_some() && encrypted_auth_config.is_some(); + let updated = if let (Some(expected_credential_fence), Some(encrypted_auth_config)) = + (expected_credential_fence, encrypted_auth_config.as_deref()) { - if observed_encrypted_auth_config.as_deref() != Some(expected_auth_config) { + if !observed_credential_matches { false } else { self.compare_and_update_provider_catalog_key_oauth_runtime_state( &ProviderCatalogKeyOAuthRuntimeStateCasUpdate { key_id: key_id.to_string(), - expected_encrypted_auth_config: Some(expected_auth_config.to_string()), + expected_encrypted_auth_config: Some( + expected_credential_fence.encrypted_auth_config.clone(), + ), + expected_credential: Some(expected_credential_fence.credential.clone()), encrypted_auth_config: encrypted_auth_config.to_string(), encrypted_api_key_update: Some(encrypted_api_key.clone()), expires_at_unix_secs_update: Some(entry.expires_at_unix_secs), @@ -2250,7 +2320,7 @@ impl AppState { }; if !updated && (requires_fenced_persistence || used_fenced_persistence) { return Err(GatewayError::Internal( - "Codex OAuth credential changed during refresh persistence".to_string(), + "OAuth credential changed during refresh persistence".to_string(), )); } let metadata_refresh_token_fingerprint = @@ -2295,13 +2365,13 @@ impl AppState { .as_deref() .map(str::trim) .is_some_and(|value| !value.is_empty()); - let expected_auth_config = if transport_has_auth_config { - self.capture_provider_transport_auth_config_fence(transport) + let expected_credential_fence = if transport_has_auth_config { + self.capture_provider_transport_credential_fence(transport) .await? } else { None }; - if transport_has_auth_config && expected_auth_config.is_none() { + if transport_has_auth_config && expected_credential_fence.is_none() { return Ok(false); } @@ -2320,10 +2390,13 @@ impl AppState { return Ok(false); } - if expected_auth_config - .as_deref() - .is_some_and(|expected| latest_key.encrypted_auth_config.as_deref() != Some(expected)) - { + if expected_credential_fence.as_ref().is_some_and(|expected| { + latest_key.encrypted_auth_config.as_deref() + != Some(expected.encrypted_auth_config.as_str()) + || latest_key.encrypted_api_key != expected.credential.encrypted_api_key + || latest_key.auth_type != expected.credential.auth_type + || latest_key.provider_id != expected.credential.provider_id + }) { return Ok(false); } @@ -2356,13 +2429,18 @@ impl AppState { latest_key.status_snapshot = sync_provider_key_oauth_status_snapshot(current_status_snapshot, &latest_key); - if let Some(expected_auth_config) = expected_auth_config.as_ref() { + if let Some(expected_credential_fence) = expected_credential_fence.as_ref() { updated = self .compare_and_update_provider_catalog_key_oauth_runtime_state( &ProviderCatalogKeyOAuthRuntimeStateCasUpdate { key_id: key_id.to_string(), - expected_encrypted_auth_config: Some(expected_auth_config.clone()), - encrypted_auth_config: expected_auth_config.clone(), + expected_encrypted_auth_config: Some( + expected_credential_fence.encrypted_auth_config.clone(), + ), + expected_credential: Some(expected_credential_fence.credential.clone()), + encrypted_auth_config: expected_credential_fence + .encrypted_auth_config + .clone(), encrypted_api_key_update: None, expires_at_unix_secs_update: None, oauth_invalid_at_unix_secs: latest_key.oauth_invalid_at_unix_secs, @@ -2407,11 +2485,12 @@ impl AppState { // Codex credentials are replaceable under a stable key id. Without a // conditional delete, refresh failure handling must retain them after // writing the generation-fenced marker. - let auto_removed = if !transport - .provider - .provider_type - .trim() - .eq_ignore_ascii_case("codex") + let auto_removed = if expected_credential_fence.is_none() + && !transport + .provider + .provider_type + .trim() + .eq_ignore_ascii_case("codex") && admin_provider_quota_pure::provider_auto_remove_banned_keys( transport.provider.config.as_ref(), ) @@ -2909,6 +2988,60 @@ mod tests { (state, repository, encrypted_auth_config) } + fn vertex_service_account_state( + ) -> (AppState, Arc, String) { + let mut provider = sample_provider(); + provider.provider_type = "vertex_ai".to_string(); + let mut endpoint = sample_endpoint(); + endpoint.api_format = "gemini:generate_content".to_string(); + endpoint.api_family = Some("gemini".to_string()); + endpoint.base_url = "https://aiplatform.googleapis.com".to_string(); + let auth_config = json!({ + "client_email": "svc@example.iam.gserviceaccount.com", + "private_key": "TEST-PRIVATE-KEY", + "project_id": "demo-project" + }); + let encrypted_auth_config = + encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, &auth_config.to_string()) + .expect("Vertex auth config should encrypt"); + let encrypted_api_key = + encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "__placeholder__") + .expect("Vertex placeholder should encrypt"); + let key = StoredProviderCatalogKey::new( + "key-1".to_string(), + "provider-1".to_string(), + "Vertex service account".to_string(), + "service_account".to_string(), + None, + true, + ) + .expect("Vertex key should build") + .with_transport_fields( + Some(json!(["gemini:generate_content"])), + encrypted_api_key, + Some(encrypted_auth_config.clone()), + None, + Some(json!({"gemini:generate_content": 1})), + None, + None, + None, + None, + ) + .expect("Vertex key transport should build"); + let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( + vec![provider], + vec![endpoint], + vec![key], + )); + let state = AppState::new() + .expect("state should build") + .with_data_state_for_tests( + GatewayDataState::with_provider_catalog_repository_for_tests(repository.clone()) + .with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY), + ); + (state, repository, encrypted_auth_config) + } + fn state_with_global_format_conversion(enabled: bool) -> AppState { let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( vec![sample_provider()], @@ -3769,6 +3902,65 @@ mod tests { )); } + #[tokio::test] + async fn memory_only_vertex_refresh_rejects_replaced_credential_generation() { + let (state, repository, initial_encrypted_auth_config) = vertex_service_account_state(); + let transport = state + .read_provider_transport_snapshot("provider-1", "endpoint-1", "key-1") + .await + .expect("Vertex transport should load") + .expect("Vertex transport should exist"); + let expected_credential_fence = state + .capture_provider_transport_credential_fence(&transport) + .await + .expect("Vertex fence should load") + .expect("Vertex fence should match"); + let entry = crate::provider_transport::CachedOAuthEntry { + provider_type: "vertex_ai".to_string(), + auth_header_name: "authorization".to_string(), + auth_header_value: "Bearer memory-only-token".to_string(), + expires_at_unix_secs: Some(4_102_444_800), + metadata: None, + source_fingerprint: Some("service-account-generation".to_string()), + }; + + state + .persist_local_oauth_refresh_entry(&transport, &entry, Some(&expected_credential_fence)) + .await + .expect("unchanged Vertex credential should accept memory-only token"); + let unchanged = repository + .list_keys_by_ids(&["key-1".to_string()]) + .await + .expect("Vertex key should load") + .pop() + .expect("Vertex key should exist"); + assert_eq!( + unchanged.encrypted_auth_config.as_deref(), + Some(initial_encrypted_auth_config.as_str()) + ); + + let mut replacement = unchanged; + replacement.encrypted_api_key = Some( + encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "admin-replacement") + .expect("replacement credential should encrypt"), + ); + repository + .update_key(&replacement) + .await + .expect("replacement credential should persist"); + + assert!( + state + .persist_local_oauth_refresh_entry( + &transport, + &entry, + Some(&expected_credential_fence), + ) + .await + .is_err() + ); + } + #[test] fn failed_refresh_persistence_discards_provisional_auth_and_cache_entry() { let mut resolution = Some(crate::provider_transport::LocalOAuthResolution { @@ -3789,6 +3981,7 @@ mod tests { refresh_in_flight: false, reused_refresh: false, distributed_lease: None, + local_refresh_guard: None, }); super::discard_failed_local_oauth_refresh_resolution(&mut resolution); @@ -3860,13 +4053,18 @@ mod tests { "email": "before@example.com", "expires_at": 4102444800_u64 }); - let (state, repository, expected_auth_config) = + let (state, repository, _expected_auth_config) = codex_oauth_state(&initial_config, "access-old"); let transport = state .read_provider_transport_snapshot("provider-1", "endpoint-1", "key-1") .await .expect("transport should load") .expect("transport should exist"); + let expected_credential_fence = state + .capture_provider_transport_credential_fence(&transport) + .await + .expect("credential fence should load") + .expect("credential fence should match the initial transport"); let replacement_config = json!({ "provider_type": "codex", @@ -3914,7 +4112,7 @@ mod tests { .persist_local_oauth_refresh_entry( &transport, &refreshed_entry, - Some(expected_auth_config.as_str()), + Some(&expected_credential_fence), ) .await .is_err()); @@ -3935,4 +4133,108 @@ mod tests { ); assert_eq!(stored.expires_at_unix_secs, None); } + + #[tokio::test] + async fn stale_refresh_failure_does_not_mark_access_token_only_replacement() { + let initial_config = json!({ + "provider_type": "codex", + "refresh_token": "refresh-stable", + "expires_at": 4102444800_u64 + }); + let (state, repository, _) = codex_oauth_state(&initial_config, "access-old"); + let stale_transport = state + .read_provider_transport_snapshot("provider-1", "endpoint-1", "key-1") + .await + .expect("transport should load") + .expect("transport should exist"); + + let replacement_api_key = + encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "access-admin") + .expect("replacement api key should encrypt"); + let mut replaced = repository + .list_keys_by_ids(&["key-1".to_string()]) + .await + .expect("key should load") + .pop() + .expect("key should exist"); + replaced.encrypted_api_key = Some(replacement_api_key.clone()); + repository + .update_key(&replaced) + .await + .expect("access token replacement should persist"); + + assert!(!state + .persist_local_oauth_refresh_failure_state( + &stale_transport, + 401, + r#"{"error":"invalid_token"}"#, + true, + ) + .await + .expect("stale failure should be ignored")); + + let stored = repository + .list_keys_by_ids(&["key-1".to_string()]) + .await + .expect("key should reload") + .pop() + .expect("replacement should remain"); + assert_eq!( + stored.encrypted_api_key.as_deref(), + Some(replacement_api_key.as_str()) + ); + assert!(stored.oauth_invalid_at_unix_secs.is_none()); + assert!(stored.oauth_invalid_reason.is_none()); + } + + #[tokio::test] + async fn stale_request_invalidation_does_not_mark_access_token_only_replacement() { + let initial_config = json!({ + "provider_type": "codex", + "refresh_token": "refresh-stable", + "expires_at": 4102444800_u64 + }); + let (state, repository, _) = codex_oauth_state(&initial_config, "access-old"); + let stale_transport = state + .read_provider_transport_snapshot("provider-1", "endpoint-1", "key-1") + .await + .expect("transport should load") + .expect("transport should exist"); + + let replacement_api_key = + encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "access-admin") + .expect("replacement api key should encrypt"); + let mut replaced = repository + .list_keys_by_ids(&["key-1".to_string()]) + .await + .expect("key should load") + .pop() + .expect("key should exist"); + replaced.encrypted_api_key = Some(replacement_api_key.clone()); + repository + .update_key(&replaced) + .await + .expect("access token replacement should persist"); + + assert!(!state + .mark_provider_transport_oauth_invalid_fenced( + &stale_transport, + "[OAUTH_EXPIRED] stale request", + ) + .await + .expect("stale invalidation should be ignored")); + + let stored = repository + .list_keys_by_ids(&["key-1".to_string()]) + .await + .expect("key should reload") + .pop() + .expect("replacement should remain"); + assert_eq!( + stored.encrypted_api_key.as_deref(), + Some(replacement_api_key.as_str()) + ); + assert!(stored.oauth_invalid_at_unix_secs.is_none()); + assert!(stored.oauth_invalid_reason.is_none()); + } } diff --git a/apps/aether-gateway/src/tests/ai_execute/fallback.rs b/apps/aether-gateway/src/tests/ai_execute/fallback.rs index 6eb79d820..d66ec63ae 100644 --- a/apps/aether-gateway/src/tests/ai_execute/fallback.rs +++ b/apps/aether-gateway/src/tests/ai_execute/fallback.rs @@ -511,7 +511,13 @@ async fn assert_ai_route_locally_denied_after_execution_runtime_miss_with_reques None ); let payload: serde_json::Value = response.json().await.expect("body should parse"); - assert_eq!(payload["error"]["type"], "http_error"); + if request_path.trim_end_matches('/') == "/v1/messages" { + assert_eq!(payload["type"], "error"); + assert_eq!(payload["error"]["type"], "overloaded_error"); + } else { + assert!(payload.get("type").is_none()); + assert_eq!(payload["error"]["type"], "http_error"); + } assert_eq!(payload["error"]["message"], expected_message); assert_eq!(*control_execute_hits.lock().expect("mutex should lock"), 0); assert_eq!(*public_hits.lock().expect("mutex should lock"), 0); diff --git a/apps/aether-gateway/src/tests/ai_execute/finalize_local_cli/direct.rs b/apps/aether-gateway/src/tests/ai_execute/finalize_local_cli/direct.rs index 429e52774..21db78caf 100644 --- a/apps/aether-gateway/src/tests/ai_execute/finalize_local_cli/direct.rs +++ b/apps/aether-gateway/src/tests/ai_execute/finalize_local_cli/direct.rs @@ -992,6 +992,7 @@ async fn gateway_executes_kiro_claude_cli_sync_upstream_stream_via_local_finaliz let response = reqwest::Client::new() .post(format!("{gateway_url}/v1/messages")) .header(http::header::CONTENT_TYPE, "application/json") + .header(http::header::USER_AGENT, "Claude-Code/2.1.0") .header( http::header::AUTHORIZATION, "Bearer sk-client-kiro-cli-finalize-local", diff --git a/apps/aether-gateway/src/tests/ai_execute/finalize_local_provider/claude.rs b/apps/aether-gateway/src/tests/ai_execute/finalize_local_provider/claude.rs index 5991495b0..8121aab2e 100644 --- a/apps/aether-gateway/src/tests/ai_execute/finalize_local_provider/claude.rs +++ b/apps/aether-gateway/src/tests/ai_execute/finalize_local_provider/claude.rs @@ -1326,6 +1326,7 @@ async fn gateway_executes_claude_cli_sync_upstream_stream_via_local_finalize_res let response = reqwest::Client::new() .post(format!("{gateway_url}/v1/messages")) .header(http::header::CONTENT_TYPE, "application/json") + .header(http::header::USER_AGENT, "Claude-Code/2.1.0") .header( http::header::AUTHORIZATION, "Bearer sk-client-claude-cli-stream-sync-local", diff --git a/apps/aether-gateway/src/tests/ai_execute/stream_provider.rs b/apps/aether-gateway/src/tests/ai_execute/stream_provider.rs index 68ec79da7..8143f5af9 100644 --- a/apps/aether-gateway/src/tests/ai_execute/stream_provider.rs +++ b/apps/aether-gateway/src/tests/ai_execute/stream_provider.rs @@ -516,6 +516,7 @@ async fn gateway_executes_kiro_claude_cli_stream_via_local_provider_catalog_cand let response = reqwest::Client::new() .post(format!("{gateway_url}/v1/messages")) .header(http::header::CONTENT_TYPE, "application/json") + .header(http::header::USER_AGENT, "Claude-Code/2.1.0") .header( http::header::AUTHORIZATION, "Bearer sk-client-kiro-cli-local-stream", @@ -927,6 +928,9 @@ async fn gateway_executes_claude_cli_stream_via_local_decision_gate_without_wait b"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: message_start\\ndata: {\\\"type\\\":\\\"message_start\\\"}\\n\\n\"}}\n" )); tokio::time::sleep(std::time::Duration::from_millis(250)).await; + yield Ok::(Bytes::from_static( + b"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: message_stop\\ndata: {\\\"type\\\":\\\"message_stop\\\"}\\n\\n\"}}\n" + )); yield Ok::(Bytes::from_static( b"{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":31,\"ttfb_ms\":11,\"upstream_bytes\":37}}}\n" )); @@ -980,6 +984,7 @@ async fn gateway_executes_claude_cli_stream_via_local_decision_gate_without_wait let mut response = reqwest::Client::new() .post(format!("{gateway_url}/v1/messages")) .header(http::header::CONTENT_TYPE, "application/json") + .header(http::header::USER_AGENT, "Claude-Code/2.1.0") .header( http::header::AUTHORIZATION, "Bearer sk-client-claude-cli-local", @@ -1004,7 +1009,7 @@ async fn gateway_executes_claude_cli_stream_via_local_decision_gate_without_wait ); assert_eq!( response.text().await.expect("remaining body should read"), - "" + "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n" ); let seen_execution_runtime_request = seen_execution_runtime @@ -1438,6 +1443,7 @@ async fn gateway_executes_claude_code_cli_stream_via_local_decision_gate_with_lo let frames = concat!( "{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n", "{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: message_start\\ndata: {\\\"type\\\":\\\"message_start\\\"}\\n\\n\"}}\n", + "{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: message_stop\\ndata: {\\\"type\\\":\\\"message_stop\\\"}\\n\\n\"}}\n", "{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":31,\"ttfb_ms\":11,\"upstream_bytes\":37}}}\n", "{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n" ); @@ -1490,6 +1496,7 @@ async fn gateway_executes_claude_code_cli_stream_via_local_decision_gate_with_lo let response = reqwest::Client::new() .post(format!("{gateway_url}/v1/messages")) .header(http::header::CONTENT_TYPE, "application/json") + .header(http::header::USER_AGENT, "Claude-Code/2.1.0") .header( http::header::AUTHORIZATION, "Bearer sk-client-claude-code-cli-local", @@ -1522,7 +1529,10 @@ async fn gateway_executes_claude_code_cli_stream_via_local_decision_gate_with_lo assert_eq!(response.status(), StatusCode::OK); assert_eq!( strip_sse_keepalive_comments(&response.text().await.expect("body should read")), - "event: message_start\ndata: {\"type\":\"message_start\"}\n\n" + concat!( + "event: message_start\ndata: {\"type\":\"message_start\"}\n\n", + "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n", + ) ); let seen_execution_runtime_request = seen_execution_runtime @@ -1551,14 +1561,17 @@ async fn gateway_executes_claude_code_cli_stream_via_local_decision_gate_with_lo ); assert_eq!( seen_execution_runtime_request.anthropic_beta, - "claude-code-20250219,oauth-2025-04-20,interleaved-thinking-2025-05-14,custom-beta" + "claude-code-20250219,oauth-2025-04-20,interleaved-thinking-2025-05-14,prompt-caching-scope-2026-01-05,effort-2025-11-24,context-management-2025-06-27,extended-cache-ttl-2025-04-11,context-1m-2025-08-07,custom-beta" ); assert_eq!(seen_execution_runtime_request.x_app, "cli"); assert_eq!( seen_execution_runtime_request.x_stainless_helper_method, "stream" ); - assert_eq!(seen_execution_runtime_request.user_agent, "Claude-Code/9.9"); + assert_eq!( + seen_execution_runtime_request.user_agent, + "claude-cli/2.1.161 (external, cli)" + ); assert_eq!( seen_execution_runtime_request.endpoint_tag, "claude-code-cli-local" @@ -1921,6 +1934,7 @@ async fn gateway_executes_claude_chat_stream_via_local_decision_gate_with_local_ let frames = concat!( "{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n", "{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: message_start\\ndata: {\\\"type\\\":\\\"message_start\\\"}\\n\\n\"}}\n", + "{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: message_stop\\ndata: {\\\"type\\\":\\\"message_stop\\\"}\\n\\n\"}}\n", "{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":31,\"ttfb_ms\":11,\"upstream_bytes\":37}}}\n", "{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n" ); @@ -1985,7 +1999,10 @@ async fn gateway_executes_claude_chat_stream_via_local_decision_gate_with_local_ assert_eq!(response.status(), StatusCode::OK); assert_eq!( strip_sse_keepalive_comments(&response.text().await.expect("body should read")), - "event: message_start\ndata: {\"type\":\"message_start\"}\n\n" + concat!( + "event: message_start\ndata: {\"type\":\"message_start\"}\n\n", + "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n", + ) ); let seen_execution_runtime_request = seen_execution_runtime diff --git a/apps/aether-gateway/src/tests/ai_execute/sync/claude/claude_code.rs b/apps/aether-gateway/src/tests/ai_execute/sync/claude/claude_code.rs index e5d38946a..b6d290690 100644 --- a/apps/aether-gateway/src/tests/ai_execute/sync/claude/claude_code.rs +++ b/apps/aether-gateway/src/tests/ai_execute/sync/claude/claude_code.rs @@ -469,6 +469,7 @@ async fn gateway_executes_claude_code_cli_sync_via_local_decision_gate_with_loca let response = reqwest::Client::new() .post(format!("{gateway_url}/v1/messages")) .header(http::header::CONTENT_TYPE, "application/json") + .header(http::header::USER_AGENT, "Claude-Code/2.1.0") .header( http::header::AUTHORIZATION, "Bearer sk-client-claude-code-cli-local", @@ -535,15 +536,18 @@ async fn gateway_executes_claude_code_cli_sync_via_local_decision_gate_with_loca ); assert_eq!( seen_execution_runtime_request.anthropic_beta, - "claude-code-20250219,oauth-2025-04-20,interleaved-thinking-2025-05-14,custom-beta" + "claude-code-20250219,oauth-2025-04-20,interleaved-thinking-2025-05-14,prompt-caching-scope-2026-01-05,effort-2025-11-24,context-management-2025-06-27,extended-cache-ttl-2025-04-11,context-1m-2025-08-07,custom-beta" ); assert_eq!(seen_execution_runtime_request.x_app, "cli"); assert_eq!(seen_execution_runtime_request.x_stainless_helper_method, ""); assert_eq!( seen_execution_runtime_request.x_stainless_package_version, - "1.0.5" + "0.94.0" + ); + assert_eq!( + seen_execution_runtime_request.user_agent, + "claude-cli/2.1.161 (external, cli)" ); - assert_eq!(seen_execution_runtime_request.user_agent, "Claude-Code/9.9"); assert_eq!( seen_execution_runtime_request.endpoint_tag, "claude-code-cli-local" diff --git a/apps/aether-gateway/src/tests/ai_execute/sync/claude/kiro.rs b/apps/aether-gateway/src/tests/ai_execute/sync/claude/kiro.rs index 16a3db78a..04a775038 100644 --- a/apps/aether-gateway/src/tests/ai_execute/sync/claude/kiro.rs +++ b/apps/aether-gateway/src/tests/ai_execute/sync/claude/kiro.rs @@ -545,6 +545,7 @@ async fn gateway_executes_kiro_claude_cli_sync_via_local_provider_catalog_candid let response = reqwest::Client::new() .post(format!("{gateway_url}/v1/messages")) .header(http::header::CONTENT_TYPE, "application/json") + .header(http::header::USER_AGENT, "Claude-Code/2.1.0") .header( http::header::AUTHORIZATION, "Bearer sk-client-kiro-cli-local-sync", @@ -1153,6 +1154,7 @@ async fn gateway_executes_kiro_claude_cli_sync_via_local_provider_catalog_candid let response = reqwest::Client::new() .post(format!("{gateway_url}/v1/messages")) .header(http::header::CONTENT_TYPE, "application/json") + .header(http::header::USER_AGENT, "Claude-Code/2.1.0") .header( http::header::AUTHORIZATION, "Bearer sk-client-kiro-cli-local-refresh", diff --git a/apps/aether-gateway/src/tests/ai_execute/sync/claude/local_chat.rs b/apps/aether-gateway/src/tests/ai_execute/sync/claude/local_chat.rs index 733c60669..ad3952b00 100644 --- a/apps/aether-gateway/src/tests/ai_execute/sync/claude/local_chat.rs +++ b/apps/aether-gateway/src/tests/ai_execute/sync/claude/local_chat.rs @@ -49,7 +49,9 @@ async fn gateway_executes_claude_chat_sync_via_local_decision_gate_with_local_sy trace_id: String, url: String, model: String, + stream: Option, auth_header_value: String, + accept: String, anthropic_version: String, anthropic_beta: String, endpoint_tag: String, @@ -270,6 +272,12 @@ async fn gateway_executes_claude_chat_sync_via_local_decision_gate_with_local_sy let raw_body = to_bytes(body, usize::MAX).await.expect("body should read"); let payload: serde_json::Value = serde_json::from_slice(&raw_body) .expect("execution runtime payload should parse"); + let upstream_url = payload + .get("url") + .and_then(|value| value.as_str()) + .unwrap_or_default() + .to_string(); + let is_count_tokens = upstream_url.ends_with("/v1/messages/count_tokens"); *seen_execution_runtime_inner .lock() .expect("mutex should lock") = Some(SeenExecutionRuntimeSyncRequest { @@ -279,11 +287,7 @@ async fn gateway_executes_claude_chat_sync_via_local_decision_gate_with_local_sy .and_then(|value| value.to_str().ok()) .unwrap_or_default() .to_string(), - url: payload - .get("url") - .and_then(|value| value.as_str()) - .unwrap_or_default() - .to_string(), + url: upstream_url, model: payload .get("body") .and_then(|value| value.get("json_body")) @@ -291,12 +295,23 @@ async fn gateway_executes_claude_chat_sync_via_local_decision_gate_with_local_sy .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), + stream: payload + .get("body") + .and_then(|value| value.get("json_body")) + .and_then(|value| value.get("stream")) + .and_then(|value| value.as_bool()), auth_header_value: payload .get("headers") .and_then(|value| value.get("x-api-key")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), + accept: payload + .get("headers") + .and_then(|value| value.get("accept")) + .and_then(|value| value.as_str()) + .unwrap_or_default() + .to_string(), anthropic_version: payload .get("headers") .and_then(|value| value.get("anthropic-version")) @@ -344,29 +359,39 @@ async fn gateway_executes_claude_chat_sync_via_local_decision_gate_with_local_sy .unwrap_or_default() .to_string(), }); - Json(json!({ - "request_id": "trace-claude-chat-local-123", - "status_code": 200, - "headers": { - "content-type": "application/json" - }, - "body": { - "json_body": { - "id": "msg-local-claude-123", - "type": "message", - "model": "claude-sonnet-4-5-upstream", - "role": "assistant", - "content": [], - "usage": { - "input_tokens": 2, - "output_tokens": 3 + if is_count_tokens { + Json(json!({ + "request_id": "trace-claude-count-tokens-local-123", + "status_code": 200, + "headers": {"content-type": "application/json"}, + "body": {"json_body": {"input_tokens": 17}}, + "telemetry": {"elapsed_ms": 11} + })) + } else { + Json(json!({ + "request_id": "trace-claude-chat-local-123", + "status_code": 200, + "headers": { + "content-type": "application/json" + }, + "body": { + "json_body": { + "id": "msg-local-claude-123", + "type": "message", + "model": "claude-sonnet-4-5-upstream", + "role": "assistant", + "content": [], + "usage": { + "input_tokens": 2, + "output_tokens": 3 + } } + }, + "telemetry": { + "elapsed_ms": 29 } - }, - "telemetry": { - "elapsed_ms": 29 - } - })) + })) + } } }), ); @@ -484,6 +509,125 @@ async fn gateway_executes_claude_chat_sync_via_local_decision_gate_with_local_sy assert_eq!(*plan_hits.lock().expect("mutex should lock"), 0); assert_eq!(*public_hits.lock().expect("mutex should lock"), 0); + for (case, body, expected_message) in [ + ("missing-body", None, "Request body is required"), + ("invalid-json", Some("{"), "Invalid JSON body"), + ( + "missing-model", + Some(r#"{"messages":[]}"#), + "model: Field required", + ), + ( + "missing-messages", + Some(r#"{"model":"claude-sonnet-4-5"}"#), + "messages: Field required", + ), + ] { + let mut request = reqwest::Client::new() + .post(format!("{gateway_url}/v1/messages/count_tokens")) + .header(http::header::CONTENT_TYPE, "application/json") + .header("x-api-key", "sk-client-claude-chat-local") + .header("anthropic-version", "2023-06-01") + .header(TRACE_ID_HEADER, format!("trace-claude-count-tokens-{case}")); + if let Some(body) = body { + request = request.body(body); + } + let invalid_response = request + .send() + .await + .expect("invalid count_tokens request should complete locally"); + assert_eq!(invalid_response.status(), StatusCode::BAD_REQUEST); + let invalid_json: serde_json::Value = invalid_response + .json() + .await + .expect("Anthropic error body should parse"); + assert_eq!(invalid_json["type"], "error"); + assert_eq!(invalid_json["error"]["type"], "invalid_request_error"); + assert_eq!(invalid_json["error"]["message"], expected_message); + assert_eq!( + seen_execution_runtime + .lock() + .expect("mutex should lock") + .as_ref() + .map(|request| request.url.as_str()), + Some("https://api.anthropic.example/custom/v1/messages"), + "invalid count_tokens request must not reach the execution runtime" + ); + } + + let count_tokens_response = reqwest::Client::new() + .post(format!("{gateway_url}/v1/messages/count_tokens")) + .header(http::header::CONTENT_TYPE, "application/json") + .header("x-api-key", "sk-client-claude-chat-local") + .header("anthropic-version", "2023-06-01") + .header(TRACE_ID_HEADER, "trace-claude-count-tokens-local-123") + .body( + "{\"model\":\"claude-sonnet-4-5\",\"messages\":[{\"role\":\"user\",\"content\":\"hello\"}],\"stream\":true}", + ) + .send() + .await + .expect("count_tokens request should succeed"); + + assert_eq!(count_tokens_response.status(), StatusCode::OK); + assert_eq!( + count_tokens_response + .headers() + .get(EXECUTION_PATH_HEADER) + .and_then(|value| value.to_str().ok()), + Some(EXECUTION_PATH_EXECUTION_RUNTIME_SYNC) + ); + let seen_count_tokens = seen_execution_runtime + .lock() + .expect("mutex should lock") + .clone() + .expect("count_tokens execution request should be captured"); + assert_eq!( + seen_count_tokens.url, + "https://api.anthropic.example/custom/v1/messages/count_tokens" + ); + assert_eq!(seen_count_tokens.model, "claude-sonnet-4-5-upstream"); + assert_eq!(seen_count_tokens.stream, None); + assert_eq!(seen_count_tokens.accept, "application/json"); + assert_eq!( + seen_count_tokens.auth_header_value, + "sk-upstream-claude-chat" + ); + + let count_tokens_json: serde_json::Value = count_tokens_response + .json() + .await + .expect("count_tokens response should parse"); + assert_eq!( + count_tokens_json["input_tokens"], 17, + "unexpected count_tokens response: {count_tokens_json}" + ); + + use std::io::Write as _; + let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default()); + encoder + .write_all( + br#"{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"hello"}]}"#, + ) + .expect("gzip request body should encode"); + let gzip_body = encoder.finish().expect("gzip request body should finish"); + let gzip_count_tokens_response = reqwest::Client::new() + .post(format!("{gateway_url}/v1/messages/count_tokens")) + .header(http::header::CONTENT_TYPE, "application/json") + .header(http::header::CONTENT_ENCODING, "gzip") + .header("x-api-key", "sk-client-claude-chat-local") + .header("anthropic-version", "2023-06-01") + .header(TRACE_ID_HEADER, "trace-claude-count-tokens-gzip-123") + .body(gzip_body) + .send() + .await + .expect("gzip count_tokens request should succeed"); + assert_eq!(gzip_count_tokens_response.status(), StatusCode::OK); + let gzip_count_tokens_json: serde_json::Value = gzip_count_tokens_response + .json() + .await + .expect("gzip count_tokens response should parse"); + assert_eq!(gzip_count_tokens_json["input_tokens"], 17); + gateway_handle.abort(); execution_runtime_handle.abort(); upstream_handle.abort(); @@ -603,7 +747,8 @@ async fn gateway_surfaces_candidate_list_empty_reason_for_claude_chat_runtime_mi Some("candidate_list_empty") ); let payload: serde_json::Value = response.json().await.expect("body should parse"); - assert_eq!(payload["error"]["type"], "http_error"); + assert_eq!(payload["type"], "error"); + assert_eq!(payload["error"]["type"], "overloaded_error"); assert_eq!( payload["error"]["message"], "没有可用提供商支持模型 claude-sonnet-4-5 的同步请求" diff --git a/apps/aether-gateway/src/tests/ai_execute/sync/claude/local_cli.rs b/apps/aether-gateway/src/tests/ai_execute/sync/claude/local_cli.rs index 4687c1aa4..fcc657b3e 100644 --- a/apps/aether-gateway/src/tests/ai_execute/sync/claude/local_cli.rs +++ b/apps/aether-gateway/src/tests/ai_execute/sync/claude/local_cli.rs @@ -404,6 +404,7 @@ async fn gateway_executes_claude_cli_sync_via_local_decision_gate_with_local_syn let response = reqwest::Client::new() .post(format!("{gateway_url}/v1/messages")) .header(http::header::CONTENT_TYPE, "application/json") + .header(http::header::USER_AGENT, "Claude-Code/2.1.0") .header( http::header::AUTHORIZATION, "Bearer sk-client-claude-cli-local", @@ -730,6 +731,7 @@ async fn gateway_returns_claude_cli_error_for_local_sync_failure_impl() { let response = reqwest::Client::new() .post(format!("{gateway_url}/v1/messages")) .header(http::header::CONTENT_TYPE, "application/json") + .header(http::header::USER_AGENT, "Claude-Code/2.1.0") .header( http::header::AUTHORIZATION, "Bearer sk-client-claude-cli-local-error", @@ -985,6 +987,7 @@ async fn gateway_marks_claude_cli_cross_format_runtime_miss_when_format_conversi let response = reqwest::Client::new() .post(format!("{gateway_url}/v1/messages?beta=true")) .header(http::header::CONTENT_TYPE, "application/json") + .header(http::header::USER_AGENT, "Claude-Code/2.1.0") .header( http::header::AUTHORIZATION, "Bearer sk-client-claude-cli-openai-local-miss", @@ -1011,7 +1014,8 @@ async fn gateway_marks_claude_cli_cross_format_runtime_miss_when_format_conversi Some("all_candidates_skipped") ); let response_json: serde_json::Value = response.json().await.expect("body should parse"); - assert_eq!(response_json["error"]["type"], "http_error"); + assert_eq!(response_json["type"], "error"); + assert_eq!(response_json["error"]["type"], "overloaded_error"); assert_eq!( response_json["error"]["message"], "没有可用提供商支持模型 gpt-5.4 的同步请求" diff --git a/apps/aether-gateway/src/tests/concurrency.rs b/apps/aether-gateway/src/tests/concurrency.rs index 62f701a44..02fa37fa1 100644 --- a/apps/aether-gateway/src/tests/concurrency.rs +++ b/apps/aether-gateway/src/tests/concurrency.rs @@ -60,6 +60,9 @@ fn sample_decision() -> crate::control::GatewayControlDecision { route_class: Some("ai_public".to_string()), route_family: Some("openai".to_string()), route_kind: Some("chat".to_string()), + client_surface: None, + api_operation: None, + gateway_credential_carrier: None, request_auth_channel: None, auth_endpoint_signature: None, execution_runtime_candidate: true, 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 33b27787c..09ceb46a6 100644 --- a/apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs +++ b/apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs @@ -1604,7 +1604,7 @@ async fn gateway_refresh_quota_reconciles_unsupported_fixed_provider_endpoints_b ( "provider-vertex-ai-reconcile", "vertex_ai", - 3usize, + 2usize, "gemini:generate_content", "https://aiplatform.googleapis.com", "Vertex AI 暂不支持自动刷新额度", diff --git a/apps/aether-gateway/src/tests/control/admin/security.rs b/apps/aether-gateway/src/tests/control/admin/security.rs index c1a7f6eb7..2e8d1cded 100644 --- a/apps/aether-gateway/src/tests/control/admin/security.rs +++ b/apps/aether-gateway/src/tests/control/admin/security.rs @@ -48,6 +48,45 @@ async fn gateway_blocks_blacklisted_ip_before_routing() { assert_eq!(payload["error"]["message"], "当前 IP 已被禁止访问"); } +#[tokio::test] +async fn gateway_shapes_blacklist_rejections_for_claude_routes_before_routing() { + let gateway = build_router_with_state( + AppState::new() + .expect("gateway should build") + .with_admin_security_blacklist_for_tests([( + "127.0.0.1".to_string(), + "blocked".to_string(), + )]), + ); + + for path in ["/v1/messages", "/v1/messages/count_tokens"] { + let request = Request::builder() + .method("POST") + .uri(path) + .header(http::header::CONTENT_TYPE, "application/json") + .body(Body::from(r#"{"model":"claude-sonnet-4","messages":[]}"#)) + .expect("request should build"); + + let response = send_request(gateway.clone(), request).await; + + assert_eq!(response.status(), StatusCode::FORBIDDEN, "path: {path}"); + let payload = response + .into_body() + .collect() + .await + .expect("body should collect") + .to_bytes(); + let payload: serde_json::Value = + serde_json::from_slice(&payload).expect("response should be json"); + assert_eq!(payload["type"], "error", "path: {path}"); + assert_eq!(payload["error"]["type"], "permission_error", "path: {path}"); + assert_eq!( + payload["error"]["message"], "当前 IP 已被禁止访问", + "path: {path}" + ); + } +} + #[tokio::test] async fn gateway_blocks_forwarded_ip_from_trusted_proxy() { let gateway = build_router_with_state( diff --git a/apps/aether-gateway/src/tests/control/admin/system_import.rs b/apps/aether-gateway/src/tests/control/admin/system_import.rs index ad05042d1..d03cb3cac 100644 --- a/apps/aether-gateway/src/tests/control/admin/system_import.rs +++ b/apps/aether-gateway/src/tests/control/admin/system_import.rs @@ -1758,6 +1758,144 @@ async fn gateway_reports_field_path_for_invalid_admin_system_config_import_shape gateway_handle.abort(); } +#[test] +fn gateway_rejects_invalid_anthropic_profiles_during_admin_system_config_import() { + run_admin_system_import_test( + "gateway_rejects_invalid_anthropic_profiles_during_admin_system_config_import", + gateway_rejects_invalid_anthropic_profiles_during_admin_system_config_import_impl, + ); +} + +async fn gateway_rejects_invalid_anthropic_profiles_during_admin_system_config_import_impl() { + let gateway = build_router_with_state( + AppState::new() + .expect("gateway should build") + .with_data_state_for_tests(build_empty_admin_system_data_state()), + ); + let (gateway_url, gateway_handle) = start_server(gateway).await; + let client = reqwest::Client::new(); + + for config_scope in ["provider", "endpoint"] { + let mut payload = sample_system_import_payload(); + let invalid_config = json!({ + "anthropic": {"compatibility_profile": "claude_cod_typo"} + }); + if config_scope == "provider" { + payload["providers"][0]["config"] = invalid_config; + } else { + payload["providers"][0]["endpoints"][0]["config"] = invalid_config; + } + + let response = client + .post(format!("{gateway_url}/api/admin/system/config/import")) + .header(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(&payload) + .send() + .await + .expect("invalid Anthropic profile import should complete locally"); + + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + let body: Value = response.json().await.expect("json body should parse"); + assert_eq!( + body["detail"], "无效的 Anthropic compatibility profile", + "unexpected {config_scope} validation response: {body}" + ); + } + + gateway_handle.abort(); +} + +#[test] +fn gateway_does_not_restore_retired_vertex_claude_endpoint_from_system_import() { + run_admin_system_import_test( + "gateway_does_not_restore_retired_vertex_claude_endpoint_from_system_import", + gateway_does_not_restore_retired_vertex_claude_endpoint_from_system_import_impl, + ); +} + +async fn gateway_does_not_restore_retired_vertex_claude_endpoint_from_system_import_impl() { + let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( + Vec::new(), + Vec::new(), + Vec::new(), + )); + let global_model_repository = Arc::new(InMemoryGlobalModelReadRepository::seed(Vec::< + StoredPublicGlobalModel, + >::new())); + let data_state = build_admin_system_data_state_with_repositories( + Arc::clone(&provider_catalog_repository), + Arc::clone(&global_model_repository), + ); + let gateway = build_router_with_state( + AppState::new() + .expect("gateway should build") + .with_data_state_for_tests(data_state), + ); + let (gateway_url, gateway_handle) = start_server(gateway).await; + + let mut payload = sample_system_import_payload(); + payload["providers"][0]["name"] = json!("legacy-vertex-backup"); + payload["providers"][0]["provider_type"] = json!("vertex_ai"); + payload["providers"][0]["endpoints"] = json!([ + { + "api_format": "gemini:generate_content", + "base_url": "https://aiplatform.googleapis.com", + "max_retries": 2, + "is_active": true + }, + { + "api_format": "claude:messages", + "base_url": "https://aiplatform.googleapis.com", + "max_retries": 2, + "is_active": true + } + ]); + payload["providers"][0]["api_keys"] = json!([]); + payload["providers"][0]["models"] = json!([]); + + let response = reqwest::Client::new() + .post(format!("{gateway_url}/api/admin/system/config/import")) + .header(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(&payload) + .send() + .await + .expect("legacy Vertex import should complete"); + + let status = response.status(); + let response_body: Value = response.json().await.expect("json body should parse"); + assert_eq!(status, StatusCode::OK, "payload={response_body}"); + assert_eq!(response_body["stats"]["endpoints"]["created"], json!(1)); + assert_eq!(response_body["stats"]["endpoints"]["skipped"], json!(1)); + assert!(response_body["stats"]["errors"] + .as_array() + .is_some_and(|errors| errors.iter().any(|error| { + error + .as_str() + .is_some_and(|error| error.contains("claude:messages")) + }))); + + let providers = provider_catalog_repository + .list_providers(false) + .await + .expect("providers should load"); + assert_eq!(providers.len(), 1); + let endpoints = provider_catalog_repository + .list_endpoints_by_provider_ids(std::slice::from_ref(&providers[0].id)) + .await + .expect("endpoints should load"); + assert_eq!(endpoints.len(), 1, "unexpected endpoints: {endpoints:?}"); + assert_eq!(endpoints[0].api_format, "gemini:generate_content"); + assert!(endpoints[0].is_active); + + gateway_handle.abort(); +} + #[test] fn gateway_imports_admin_system_config_with_numeric_string_prices() { run_admin_system_import_test( diff --git a/apps/aether-gateway/src/tests/control/proxy/local_denials.rs b/apps/aether-gateway/src/tests/control/proxy/local_denials.rs index ef3c06fdb..fe65d1d76 100644 --- a/apps/aether-gateway/src/tests/control/proxy/local_denials.rs +++ b/apps/aether-gateway/src/tests/control/proxy/local_denials.rs @@ -322,6 +322,7 @@ async fn gateway_locally_denies_invalid_bearer_api_key_without_hitting_control_o Some(EXECUTION_PATH_LOCAL_AUTH_DENIED) ); let payload: serde_json::Value = response.json().await.expect("response json should parse"); + assert!(payload.get("type").is_none()); assert_eq!(payload["error"]["type"], "http_error"); assert_eq!(payload["error"]["message"], "无效的API密钥"); assert_eq!(*auth_context_hits.lock().expect("mutex should lock"), 0); @@ -331,6 +332,57 @@ async fn gateway_locally_denies_invalid_bearer_api_key_without_hitting_control_o upstream_handle.abort(); } +#[tokio::test] +async fn gateway_claude_routes_use_anthropic_authentication_error_for_invalid_api_key() { + let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( + Some(hash_api_key("sk-other-claude-key")), + sample_currently_usable_auth_snapshot("key-claude-other", "user-claude-other"), + )])); + let gateway = build_router_with_state( + AppState::new() + .expect("gateway state should build") + .with_auth_api_key_data_reader_for_tests(repository), + ); + let (gateway_url, gateway_handle) = start_server(gateway).await; + let client = reqwest::Client::new(); + + for (path, trace_id) in [ + ("/v1/messages", "trace-control-claude-invalid-key-messages"), + ( + "/v1/messages/count_tokens", + "trace-control-claude-invalid-key-count-tokens", + ), + ] { + let response = client + .post(format!("{gateway_url}{path}")) + .header(http::header::CONTENT_TYPE, "application/json") + .header("x-api-key", "sk-missing-claude-key") + .header(TRACE_ID_HEADER, trace_id) + .body("{\"model\":\"claude-sonnet-4-5\",\"messages\":[]}") + .send() + .await + .expect("request should complete locally"); + + assert_eq!(response.status(), StatusCode::UNAUTHORIZED, "path: {path}"); + assert_eq!( + response + .headers() + .get(EXECUTION_PATH_HEADER) + .and_then(|value| value.to_str().ok()), + Some(EXECUTION_PATH_LOCAL_AUTH_DENIED), + "path: {path}" + ); + let payload: serde_json::Value = response.json().await.expect("response json should parse"); + assert_eq!(payload["type"], "error", "path: {path}"); + assert_eq!( + payload["error"]["type"], "authentication_error", + "path: {path}" + ); + } + + gateway_handle.abort(); +} + #[tokio::test] async fn gateway_locally_denies_admin_proxy_without_admin_principal_and_without_hitting_upstream() { let upstream_hits = Arc::new(Mutex::new(0usize)); @@ -498,7 +550,8 @@ async fn gateway_locally_denies_disallowed_claude_api_format_without_hitting_con assert_eq!(response.status(), StatusCode::FORBIDDEN); let payload: serde_json::Value = response.json().await.expect("response json should parse"); - assert_eq!(payload["error"]["type"], "http_error"); + assert_eq!(payload["type"], "error"); + assert_eq!(payload["error"]["type"], "permission_error"); assert_eq!( payload["error"]["message"], "当前用户、用户组或密钥的访问控制策略不允许访问 claude:messages 格式" @@ -581,7 +634,8 @@ async fn gateway_locally_denies_disallowed_provider_without_hitting_control_or_u assert_eq!(response.status(), StatusCode::FORBIDDEN); let payload: serde_json::Value = response.json().await.expect("response json should parse"); - assert_eq!(payload["error"]["type"], "http_error"); + assert_eq!(payload["type"], "error"); + assert_eq!(payload["error"]["type"], "permission_error"); assert_eq!( payload["error"]["message"], "当前用户、用户组或密钥的访问控制策略不允许访问 claude 提供商" @@ -817,3 +871,52 @@ async fn gateway_locally_denies_disallowed_openai_model_without_hitting_control_ gateway_handle.abort(); upstream_handle.abort(); } + +#[tokio::test] +async fn gateway_locally_denies_disallowed_claude_model_with_anthropic_permission_error() { + let mut snapshot = + sample_currently_usable_auth_snapshot("key-claude-model-123", "user-claude-model-123"); + snapshot.api_key_allowed_providers = Some(vec!["claude".to_string()]); + snapshot.user_allowed_providers = Some(vec!["claude".to_string()]); + snapshot.api_key_allowed_api_formats = Some(vec!["claude:messages".to_string()]); + snapshot.user_allowed_api_formats = Some(vec!["claude:messages".to_string()]); + snapshot.api_key_allowed_models = Some(vec!["claude-haiku-4-5".to_string()]); + let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( + Some(hash_api_key("sk-claude-model-guard-123")), + snapshot, + )])); + let gateway = build_router_with_state( + AppState::new() + .expect("gateway state should build") + .with_auth_api_key_data_reader_for_tests(repository), + ); + let (gateway_url, gateway_handle) = start_server(gateway).await; + + let response = reqwest::Client::new() + .post(format!("{gateway_url}/v1/messages")) + .header(http::header::CONTENT_TYPE, "application/json") + .header("x-api-key", "sk-claude-model-guard-123") + .header(TRACE_ID_HEADER, "trace-control-claude-model-guard-1") + .body("{\"model\":\"claude-sonnet-4-5\",\"messages\":[]}") + .send() + .await + .expect("request should complete locally"); + + assert_eq!(response.status(), StatusCode::FORBIDDEN); + assert_eq!( + response + .headers() + .get(EXECUTION_PATH_HEADER) + .and_then(|value| value.to_str().ok()), + Some(EXECUTION_PATH_LOCAL_AUTH_DENIED) + ); + let payload: serde_json::Value = response.json().await.expect("response json should parse"); + assert_eq!(payload["type"], "error"); + assert_eq!(payload["error"]["type"], "permission_error"); + assert_eq!( + payload["error"]["message"], + "当前用户、用户组或密钥的访问控制策略不允许访问模型 claude-sonnet-4-5" + ); + + gateway_handle.abort(); +} diff --git a/apps/aether-gateway/src/tests/frontdoor/ai.rs b/apps/aether-gateway/src/tests/frontdoor/ai.rs index 2ddf5b6ab..5f19d6589 100644 --- a/apps/aether-gateway/src/tests/frontdoor/ai.rs +++ b/apps/aether-gateway/src/tests/frontdoor/ai.rs @@ -990,136 +990,6 @@ async fn gateway_handles_public_gemini_models_without_hitting_fallback_probe() { fallback_probe_handle.abort(); } -#[tokio::test] -async fn gateway_handles_claude_count_tokens_without_hitting_fallback_probe() { - let fallback_probe_hits = Arc::new(Mutex::new(0usize)); - let fallback_probe_hits_clone = Arc::clone(&fallback_probe_hits); - let fallback_probe = Router::new().route( - "/{*path}", - any(move |_request: Request| { - let fallback_probe_hits_inner = Arc::clone(&fallback_probe_hits_clone); - async move { - *fallback_probe_hits_inner.lock().expect("mutex should lock") += 1; - (StatusCode::OK, Json(json!({"proxied": true}))).into_response() - } - }), - ); - - let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( - Some(hash_api_key("sk-claude-count")), - unrestricted_models_snapshot("key-claude-count", "user-claude-count"), - )])); - - let (_unused_fallback_probe_url, fallback_probe_handle) = start_server(fallback_probe).await; - let gateway = build_router_with_state( - AppState::new() - .expect("gateway should build") - .with_auth_api_key_data_reader_for_tests(auth_repository), - ); - let (gateway_url, gateway_handle) = start_server(gateway).await; - - let response = reqwest::Client::new() - .post(format!("{gateway_url}/v1/messages/count_tokens")) - .header("x-api-key", "sk-claude-count") - .header("anthropic-version", "2023-06-01") - .body( - serde_json::to_vec(&json!({ - "model": "claude-sonnet-4-5", - "system": [{"type": "text", "text": "abcdefghijklmnop"}], - "messages": [ - { - "role": "user", - "content": "abcdefghijkl" - }, - { - "role": "assistant", - "content": [ - {"type": "text", "text": "abcdefgh"}, - {"type": "tool_use", "name": "ignored", "input": {"city": "SF"}} - ] - } - ] - })) - .expect("request body should encode"), - ) - .send() - .await - .expect("request should succeed"); - - assert_eq!(response.status(), StatusCode::OK); - assert_eq!( - response - .headers() - .get(EXECUTION_PATH_HEADER) - .and_then(|value| value.to_str().ok()), - Some(EXECUTION_PATH_LOCAL_AI_PUBLIC) - ); - let payload: serde_json::Value = response.json().await.expect("json body should parse"); - assert_eq!(payload["input_tokens"], 17); - assert_eq!(*fallback_probe_hits.lock().expect("mutex should lock"), 0); - - gateway_handle.abort(); - fallback_probe_handle.abort(); -} - -#[tokio::test] -async fn gateway_rejects_invalid_claude_count_tokens_payload_without_hitting_fallback_probe() { - let fallback_probe_hits = Arc::new(Mutex::new(0usize)); - let fallback_probe_hits_clone = Arc::clone(&fallback_probe_hits); - let fallback_probe = Router::new().route( - "/{*path}", - any(move |_request: Request| { - let fallback_probe_hits_inner = Arc::clone(&fallback_probe_hits_clone); - async move { - *fallback_probe_hits_inner.lock().expect("mutex should lock") += 1; - (StatusCode::OK, Json(json!({"proxied": true}))).into_response() - } - }), - ); - - let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( - Some(hash_api_key("sk-claude-count-invalid")), - unrestricted_models_snapshot("key-claude-count-invalid", "user-claude-count-invalid"), - )])); - - let (_unused_fallback_probe_url, fallback_probe_handle) = start_server(fallback_probe).await; - let gateway = build_router_with_state( - AppState::new() - .expect("gateway should build") - .with_auth_api_key_data_reader_for_tests(auth_repository), - ); - let (gateway_url, gateway_handle) = start_server(gateway).await; - - let response = reqwest::Client::new() - .post(format!("{gateway_url}/v1/messages/count_tokens")) - .header("x-api-key", "sk-claude-count-invalid") - .body( - serde_json::to_vec(&json!({ - "model": "claude-sonnet-4-5", - "messages": [{"role": "system", "content": "bad"}] - })) - .expect("request body should encode"), - ) - .send() - .await - .expect("request should succeed"); - - assert_eq!(response.status(), StatusCode::BAD_REQUEST); - assert_eq!( - response - .headers() - .get(EXECUTION_PATH_HEADER) - .and_then(|value| value.to_str().ok()), - Some(EXECUTION_PATH_LOCAL_AI_PUBLIC) - ); - let payload: serde_json::Value = response.json().await.expect("json body should parse"); - assert_eq!(payload["detail"], "Invalid token count payload"); - assert_eq!(*fallback_probe_hits.lock().expect("mutex should lock"), 0); - - gateway_handle.abort(); - fallback_probe_handle.abort(); -} - #[tokio::test] async fn gateway_handles_antigravity_v1internal_control_plane_without_proxying() { let fallback_probe_hits = Arc::new(Mutex::new(0usize)); diff --git a/apps/aether-gateway/src/tests/proxy.rs b/apps/aether-gateway/src/tests/proxy.rs index e4c646bb0..81c9e6dda 100644 --- a/apps/aether-gateway/src/tests/proxy.rs +++ b/apps/aether-gateway/src/tests/proxy.rs @@ -383,6 +383,86 @@ async fn gateway_rejects_execution_runtime_loop_guarded_ai_request() { gateway_handle.abort(); } +#[tokio::test] +async fn gateway_shapes_execution_loop_rejections_for_claude_routes() { + let gateway = build_router().expect("gateway should build"); + let (gateway_url, gateway_handle) = start_server(gateway).await; + + for path in ["/v1/messages", "/v1/messages/count_tokens"] { + let response = reqwest::Client::new() + .post(format!("{gateway_url}{path}")) + .header( + EXECUTION_RUNTIME_LOOP_GUARD_HEADER, + EXECUTION_RUNTIME_LOOP_GUARD_VALUE, + ) + .header(http::header::CONTENT_TYPE, "application/json") + .body(r#"{"model":"claude-sonnet-4","messages":[]}"#) + .send() + .await + .expect("request should succeed"); + + assert_eq!(response.status(), StatusCode::LOOP_DETECTED, "path: {path}"); + assert_eq!( + response + .headers() + .get(EXECUTION_PATH_HEADER) + .and_then(|value| value.to_str().ok()), + Some(EXECUTION_PATH_LOCAL_EXECUTION_LOOP_DETECTED), + "path: {path}" + ); + let payload: serde_json::Value = response.json().await.expect("body should parse"); + assert_eq!(payload["type"], "error", "path: {path}"); + assert_eq!(payload["error"]["type"], "api_error", "path: {path}"); + assert_eq!( + payload["error"]["message"], + "Gateway detected an execution runtime request loop back into the local frontdoor", + "path: {path}" + ); + } + + gateway_handle.abort(); +} + +#[tokio::test] +async fn gateway_shapes_wrong_method_rejections_for_claude_routes() { + let gateway = build_router().expect("gateway should build"); + let (gateway_url, gateway_handle) = start_server(gateway).await; + + for path in ["/v1/messages", "/v1/messages/count_tokens"] { + let response = reqwest::Client::new() + .get(format!("{gateway_url}{path}")) + .send() + .await + .expect("request should succeed"); + + assert_eq!( + response.status(), + StatusCode::METHOD_NOT_ALLOWED, + "path: {path}" + ); + assert_eq!( + response + .headers() + .get(http::header::ALLOW) + .and_then(|value| value.to_str().ok()), + Some("POST"), + "path: {path}" + ); + let payload: serde_json::Value = response.json().await.expect("body should parse"); + assert_eq!(payload["type"], "error", "path: {path}"); + assert_eq!( + payload["error"]["type"], "invalid_request_error", + "path: {path}" + ); + assert_eq!( + payload["error"]["message"], "Method not allowed", + "path: {path}" + ); + } + + gateway_handle.abort(); +} + #[tokio::test] async fn gateway_rejects_execution_runtime_via_guarded_ai_request() { let gateway = build_router().expect("gateway should build"); diff --git a/apps/aether-gateway/src/tests/usage/local.rs b/apps/aether-gateway/src/tests/usage/local.rs index ea8a63ec4..a64a3080d 100644 --- a/apps/aether-gateway/src/tests/usage/local.rs +++ b/apps/aether-gateway/src/tests/usage/local.rs @@ -1121,7 +1121,8 @@ async fn gateway_records_failed_usage_for_claude_runtime_miss_without_execution_ Some("candidate_list_empty") ); let body_json: serde_json::Value = response.json().await.expect("body should parse"); - assert_eq!(body_json["error"]["type"], "http_error"); + assert_eq!(body_json["type"], "error"); + assert_eq!(body_json["error"]["type"], "overloaded_error"); assert_eq!( body_json["error"]["message"], "没有可用提供商支持模型 claude-sonnet-4-5 的同步请求" @@ -1169,7 +1170,7 @@ async fn gateway_records_failed_usage_for_claude_runtime_miss_without_execution_ .and_then(|value| value.get("error")) .and_then(|value| value.get("type")) .and_then(|value| value.as_str()), - Some("http_error") + Some("overloaded_error") ); let stored_candidates = request_candidate_repository @@ -1749,6 +1750,7 @@ async fn gateway_records_failed_usage_when_all_local_claude_cli_candidates_are_s let response = reqwest::Client::new() .post(format!("{gateway_url}/v1/messages?beta=true")) .header(http::header::CONTENT_TYPE, "application/json") + .header(http::header::USER_AGENT, "Claude-Code/2.1.0") .header( http::header::AUTHORIZATION, "Bearer sk-client-claude-cli-usage-local-miss", @@ -1768,7 +1770,8 @@ async fn gateway_records_failed_usage_when_all_local_claude_cli_candidates_are_s Some("all_candidates_skipped") ); let body_json: serde_json::Value = response.json().await.expect("body should parse"); - assert_eq!(body_json["error"]["type"], "http_error"); + assert_eq!(body_json["type"], "error"); + assert_eq!(body_json["error"]["type"], "overloaded_error"); assert_eq!( body_json["error"]["message"], "没有可用提供商支持模型 gpt-5.4 的同步请求" @@ -2062,6 +2065,7 @@ fn gateway_keeps_failed_usage_request_capture_lightweight_for_large_local_claude let response = reqwest::Client::new() .post(format!("{gateway_url}/v1/messages?beta=true")) .header(http::header::CONTENT_TYPE, "application/json") + .header(http::header::USER_AGENT, "Claude-Code/2.1.0") .header( http::header::AUTHORIZATION, "Bearer sk-client-claude-cli-usage-local-miss-large", diff --git a/crates/aether-ai/formats/src/api.rs b/crates/aether-ai/formats/src/api.rs index 01a687ac6..d38bc0cdf 100644 --- a/crates/aether-ai/formats/src/api.rs +++ b/crates/aether-ai/formats/src/api.rs @@ -2,13 +2,14 @@ pub use crate::contracts::{ core_error_background_report_kind, core_error_default_client_api_format, core_success_background_report_kind, implicit_sync_finalize_report_kind, is_openai_responses_stream_plan_kind, is_openai_responses_sync_plan_kind, AiControlPlanRequest, - ExecutionRuntimeAuthContext, CLAUDE_CHAT_STREAM_PLAN_KIND, + ApiOperation, ClientSurface, ExecutionRuntimeAuthContext, CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_STREAM_SUCCESS_REPORT_KIND, CLAUDE_CHAT_SYNC_ERROR_REPORT_KIND, CLAUDE_CHAT_SYNC_FINALIZE_REPORT_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CHAT_SYNC_SUCCESS_REPORT_KIND, CLAUDE_CLI_STREAM_PLAN_KIND, CLAUDE_CLI_STREAM_SUCCESS_REPORT_KIND, CLAUDE_CLI_SYNC_ERROR_REPORT_KIND, CLAUDE_CLI_SYNC_FINALIZE_REPORT_KIND, CLAUDE_CLI_SYNC_PLAN_KIND, - CLAUDE_CLI_SYNC_SUCCESS_REPORT_KIND, EXECUTION_RUNTIME_STREAM_ACTION, + CLAUDE_CLI_SYNC_SUCCESS_REPORT_KIND, CLAUDE_COUNT_TOKENS_SYNC_PLAN_KIND, + CLAUDE_COUNT_TOKENS_SYNC_SUCCESS_REPORT_KIND, EXECUTION_RUNTIME_STREAM_ACTION, EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_ACTION, EXECUTION_RUNTIME_SYNC_DECISION_ACTION, GEMINI_CHAT_STREAM_PLAN_KIND, GEMINI_CHAT_STREAM_SUCCESS_REPORT_KIND, GEMINI_CHAT_SYNC_ERROR_REPORT_KIND, @@ -127,7 +128,9 @@ pub use crate::formats::shared::response::{ pub use crate::formats::shared::routing::{ is_matching_stream_http_request, is_matching_stream_request, request_path_implies_stream_request, resolve_execution_runtime_stream_plan_kind, - resolve_execution_runtime_sync_plan_kind, sanitize_request_path, + resolve_execution_runtime_stream_plan_kind_with_client_surface, + resolve_execution_runtime_sync_plan_kind, + resolve_execution_runtime_sync_plan_kind_with_client_surface, sanitize_request_path, sanitize_request_path_and_query, sanitize_request_query_string, supports_stream_execution_decision_kind, supports_sync_execution_decision_kind, }; diff --git a/crates/aether-ai/formats/src/contracts/mod.rs b/crates/aether-ai/formats/src/contracts/mod.rs index 62f11b066..2c26e0f93 100644 --- a/crates/aether-ai/formats/src/contracts/mod.rs +++ b/crates/aether-ai/formats/src/contracts/mod.rs @@ -3,6 +3,7 @@ mod auth_context; mod control_payloads; mod plan_kinds; mod report_kinds; +mod request_dimensions; pub use actions::{ EXECUTION_RUNTIME_STREAM_ACTION, EXECUTION_RUNTIME_STREAM_DECISION_ACTION, @@ -13,19 +14,20 @@ pub use control_payloads::{build_ai_control_plan_request, AiControlPlanRequest}; pub use plan_kinds::{ is_openai_responses_stream_plan_kind, is_openai_responses_sync_plan_kind, CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND, - CLAUDE_CLI_SYNC_PLAN_KIND, GEMINI_CHAT_STREAM_PLAN_KIND, GEMINI_CHAT_SYNC_PLAN_KIND, - GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_SYNC_PLAN_KIND, GEMINI_EMBEDDING_SYNC_PLAN_KIND, - GEMINI_FILES_DELETE_PLAN_KIND, GEMINI_FILES_DOWNLOAD_PLAN_KIND, GEMINI_FILES_GET_PLAN_KIND, - GEMINI_FILES_LIST_PLAN_KIND, GEMINI_FILES_UPLOAD_PLAN_KIND, - GEMINI_INTERACTIONS_STREAM_PLAN_KIND, GEMINI_INTERACTIONS_SYNC_PLAN_KIND, - GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND, GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND, - OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND, OPENAI_EMBEDDING_SYNC_PLAN_KIND, - OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND, OPENAI_RERANK_SYNC_PLAN_KIND, - OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND, - OPENAI_RESPONSES_STREAM_PLAN_KIND, OPENAI_RESPONSES_SYNC_PLAN_KIND, - OPENAI_SEARCH_SYNC_PLAN_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, - OPENAI_VIDEO_CONTENT_PLAN_KIND, OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND, - OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND, OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND, + CLAUDE_CLI_SYNC_PLAN_KIND, CLAUDE_COUNT_TOKENS_SYNC_PLAN_KIND, GEMINI_CHAT_STREAM_PLAN_KIND, + GEMINI_CHAT_SYNC_PLAN_KIND, GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_SYNC_PLAN_KIND, + GEMINI_EMBEDDING_SYNC_PLAN_KIND, GEMINI_FILES_DELETE_PLAN_KIND, + GEMINI_FILES_DOWNLOAD_PLAN_KIND, GEMINI_FILES_GET_PLAN_KIND, GEMINI_FILES_LIST_PLAN_KIND, + GEMINI_FILES_UPLOAD_PLAN_KIND, GEMINI_INTERACTIONS_STREAM_PLAN_KIND, + GEMINI_INTERACTIONS_SYNC_PLAN_KIND, GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND, + GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND, + OPENAI_EMBEDDING_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND, + OPENAI_RERANK_SYNC_PLAN_KIND, OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, + OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND, OPENAI_RESPONSES_STREAM_PLAN_KIND, + OPENAI_RESPONSES_SYNC_PLAN_KIND, OPENAI_SEARCH_SYNC_PLAN_KIND, + OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_VIDEO_CONTENT_PLAN_KIND, + OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND, + OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND, }; pub use report_kinds::{ core_error_background_report_kind, core_error_default_client_api_format, @@ -34,20 +36,20 @@ pub use report_kinds::{ CLAUDE_CHAT_SYNC_ERROR_REPORT_KIND, CLAUDE_CHAT_SYNC_FINALIZE_REPORT_KIND, CLAUDE_CHAT_SYNC_SUCCESS_REPORT_KIND, CLAUDE_CLI_STREAM_SUCCESS_REPORT_KIND, CLAUDE_CLI_SYNC_ERROR_REPORT_KIND, CLAUDE_CLI_SYNC_FINALIZE_REPORT_KIND, - CLAUDE_CLI_SYNC_SUCCESS_REPORT_KIND, GEMINI_CHAT_STREAM_SUCCESS_REPORT_KIND, - GEMINI_CHAT_SYNC_ERROR_REPORT_KIND, GEMINI_CHAT_SYNC_FINALIZE_REPORT_KIND, - GEMINI_CHAT_SYNC_SUCCESS_REPORT_KIND, GEMINI_CLI_STREAM_SUCCESS_REPORT_KIND, - GEMINI_CLI_SYNC_ERROR_REPORT_KIND, GEMINI_CLI_SYNC_FINALIZE_REPORT_KIND, - GEMINI_CLI_SYNC_SUCCESS_REPORT_KIND, GEMINI_EMBEDDING_SYNC_SUCCESS_REPORT_KIND, - GEMINI_INTERACTIONS_STREAM_SUCCESS_REPORT_KIND, GEMINI_INTERACTIONS_SYNC_ERROR_REPORT_KIND, - GEMINI_INTERACTIONS_SYNC_FINALIZE_REPORT_KIND, GEMINI_INTERACTIONS_SYNC_SUCCESS_REPORT_KIND, - GEMINI_VIDEO_CREATE_SYNC_FINALIZE_REPORT_KIND, OPENAI_CHAT_STREAM_SUCCESS_REPORT_KIND, - OPENAI_CHAT_SYNC_ERROR_REPORT_KIND, OPENAI_CHAT_SYNC_FINALIZE_REPORT_KIND, - OPENAI_CHAT_SYNC_SUCCESS_REPORT_KIND, OPENAI_EMBEDDING_SYNC_ERROR_REPORT_KIND, - OPENAI_EMBEDDING_SYNC_FINALIZE_REPORT_KIND, OPENAI_EMBEDDING_SYNC_SUCCESS_REPORT_KIND, - OPENAI_IMAGE_STREAM_SUCCESS_REPORT_KIND, OPENAI_IMAGE_SYNC_ERROR_REPORT_KIND, - OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND, OPENAI_IMAGE_SYNC_SUCCESS_REPORT_KIND, - OPENAI_RESPONSES_COMPACT_STREAM_SUCCESS_REPORT_KIND, + CLAUDE_CLI_SYNC_SUCCESS_REPORT_KIND, CLAUDE_COUNT_TOKENS_SYNC_SUCCESS_REPORT_KIND, + GEMINI_CHAT_STREAM_SUCCESS_REPORT_KIND, GEMINI_CHAT_SYNC_ERROR_REPORT_KIND, + GEMINI_CHAT_SYNC_FINALIZE_REPORT_KIND, GEMINI_CHAT_SYNC_SUCCESS_REPORT_KIND, + GEMINI_CLI_STREAM_SUCCESS_REPORT_KIND, GEMINI_CLI_SYNC_ERROR_REPORT_KIND, + GEMINI_CLI_SYNC_FINALIZE_REPORT_KIND, GEMINI_CLI_SYNC_SUCCESS_REPORT_KIND, + GEMINI_EMBEDDING_SYNC_SUCCESS_REPORT_KIND, GEMINI_INTERACTIONS_STREAM_SUCCESS_REPORT_KIND, + GEMINI_INTERACTIONS_SYNC_ERROR_REPORT_KIND, GEMINI_INTERACTIONS_SYNC_FINALIZE_REPORT_KIND, + GEMINI_INTERACTIONS_SYNC_SUCCESS_REPORT_KIND, GEMINI_VIDEO_CREATE_SYNC_FINALIZE_REPORT_KIND, + OPENAI_CHAT_STREAM_SUCCESS_REPORT_KIND, OPENAI_CHAT_SYNC_ERROR_REPORT_KIND, + OPENAI_CHAT_SYNC_FINALIZE_REPORT_KIND, OPENAI_CHAT_SYNC_SUCCESS_REPORT_KIND, + OPENAI_EMBEDDING_SYNC_ERROR_REPORT_KIND, OPENAI_EMBEDDING_SYNC_FINALIZE_REPORT_KIND, + OPENAI_EMBEDDING_SYNC_SUCCESS_REPORT_KIND, OPENAI_IMAGE_STREAM_SUCCESS_REPORT_KIND, + OPENAI_IMAGE_SYNC_ERROR_REPORT_KIND, OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND, + OPENAI_IMAGE_SYNC_SUCCESS_REPORT_KIND, OPENAI_RESPONSES_COMPACT_STREAM_SUCCESS_REPORT_KIND, OPENAI_RESPONSES_COMPACT_SYNC_ERROR_REPORT_KIND, OPENAI_RESPONSES_COMPACT_SYNC_FINALIZE_REPORT_KIND, OPENAI_RESPONSES_COMPACT_SYNC_SUCCESS_REPORT_KIND, OPENAI_RESPONSES_STREAM_SUCCESS_REPORT_KIND, @@ -55,3 +57,4 @@ pub use report_kinds::{ OPENAI_RESPONSES_SYNC_SUCCESS_REPORT_KIND, OPENAI_SEARCH_SYNC_SUCCESS_REPORT_KIND, OPENAI_VIDEO_CREATE_SYNC_FINALIZE_REPORT_KIND, }; +pub use request_dimensions::{ApiOperation, ClientSurface}; diff --git a/crates/aether-ai/formats/src/contracts/plan_kinds.rs b/crates/aether-ai/formats/src/contracts/plan_kinds.rs index 2a9e61329..8684eff10 100644 --- a/crates/aether-ai/formats/src/contracts/plan_kinds.rs +++ b/crates/aether-ai/formats/src/contracts/plan_kinds.rs @@ -13,6 +13,7 @@ pub const GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND: &str = "gemini_video_create_sync"; pub const GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND: &str = "gemini_video_cancel_sync"; pub const OPENAI_CHAT_STREAM_PLAN_KIND: &str = "openai_chat_stream"; pub const CLAUDE_CHAT_STREAM_PLAN_KIND: &str = "claude_chat_stream"; +pub const CLAUDE_COUNT_TOKENS_SYNC_PLAN_KIND: &str = "claude_count_tokens_sync"; pub const GEMINI_CHAT_STREAM_PLAN_KIND: &str = "gemini_chat_stream"; pub const GEMINI_INTERACTIONS_STREAM_PLAN_KIND: &str = "gemini_interactions_stream"; pub const OPENAI_RESPONSES_STREAM_PLAN_KIND: &str = "openai_responses_stream"; diff --git a/crates/aether-ai/formats/src/contracts/report_kinds.rs b/crates/aether-ai/formats/src/contracts/report_kinds.rs index ce19b1b4f..80d23275a 100644 --- a/crates/aether-ai/formats/src/contracts/report_kinds.rs +++ b/crates/aether-ai/formats/src/contracts/report_kinds.rs @@ -23,6 +23,7 @@ const LEGACY_OPENAI_COMPACT_SYNC_FINALIZE_REPORT_KIND: &str = "openai_compact_sy pub const OPENAI_CHAT_SYNC_SUCCESS_REPORT_KIND: &str = "openai_chat_sync_success"; pub const CLAUDE_CHAT_SYNC_SUCCESS_REPORT_KIND: &str = "claude_chat_sync_success"; +pub const CLAUDE_COUNT_TOKENS_SYNC_SUCCESS_REPORT_KIND: &str = "claude_count_tokens_sync_success"; pub const GEMINI_CHAT_SYNC_SUCCESS_REPORT_KIND: &str = "gemini_chat_sync_success"; pub const GEMINI_INTERACTIONS_SYNC_SUCCESS_REPORT_KIND: &str = "gemini_interactions_sync_success"; pub const OPENAI_RESPONSES_SYNC_SUCCESS_REPORT_KIND: &str = "openai_responses_sync_success"; diff --git a/crates/aether-ai/formats/src/contracts/request_dimensions.rs b/crates/aether-ai/formats/src/contracts/request_dimensions.rs new file mode 100644 index 000000000..2ac98aa3c --- /dev/null +++ b/crates/aether-ai/formats/src/contracts/request_dimensions.rs @@ -0,0 +1,59 @@ +use serde::{Deserialize, Serialize}; + +/// Client behavior profile detected at the ingress boundary. +/// +/// This is deliberately independent from the credential carrier: an +/// Anthropic SDK may use bearer auth, and Claude Code may be authenticated by +/// an Aether API key. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ClientSurface { + ClaudeCode, + AnthropicSdk, + GenericCompatible, +} + +impl ClientSurface { + pub const fn as_str(self) -> &'static str { + match self { + Self::ClaudeCode => "claude_code", + Self::AnthropicSdk => "anthropic_sdk", + Self::GenericCompatible => "generic_compatible", + } + } +} + +/// Semantic operation carried over an API wire format. +/// +/// Operations must not be represented as additional API formats: both +/// Anthropic message creation and token counting use the `claude:messages` +/// request contract. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ApiOperation { + ClaudeMessagesCreate, + ClaudeCountTokens, + OpenAiResponsesCompact, +} + +impl ApiOperation { + pub const fn as_str(self) -> &'static str { + match self { + Self::ClaudeMessagesCreate => "messages", + Self::ClaudeCountTokens => "count_tokens", + Self::OpenAiResponsesCompact => "compact", + } + } +} + +#[cfg(test)] +mod tests { + use super::{ApiOperation, ClientSurface}; + + #[test] + fn request_dimensions_have_stable_external_names() { + assert_eq!(ClientSurface::ClaudeCode.as_str(), "claude_code"); + assert_eq!(ApiOperation::ClaudeMessagesCreate.as_str(), "messages"); + assert_eq!(ApiOperation::ClaudeCountTokens.as_str(), "count_tokens"); + } +} diff --git a/crates/aether-ai/formats/src/formats/shared/error_body.rs b/crates/aether-ai/formats/src/formats/shared/error_body.rs index a00f053af..bc4172d49 100644 --- a/crates/aether-ai/formats/src/formats/shared/error_body.rs +++ b/crates/aether-ai/formats/src/formats/shared/error_body.rs @@ -8,6 +8,7 @@ pub enum LocalCoreSyncErrorKind { NotFound, RateLimit, ContextLengthExceeded, + RequestTooLarge, Overloaded, ServerError, } @@ -93,7 +94,9 @@ fn map_local_sync_error_kind_to_openai_type(kind: LocalCoreSyncErrorKind) -> &'s LocalCoreSyncErrorKind::PermissionDenied => "permission_error", LocalCoreSyncErrorKind::NotFound => "not_found_error", LocalCoreSyncErrorKind::RateLimit => "rate_limit_error", - LocalCoreSyncErrorKind::ContextLengthExceeded => "context_length_exceeded", + LocalCoreSyncErrorKind::ContextLengthExceeded | LocalCoreSyncErrorKind::RequestTooLarge => { + "context_length_exceeded" + } LocalCoreSyncErrorKind::Overloaded | LocalCoreSyncErrorKind::ServerError => "server_error", } } @@ -103,19 +106,21 @@ fn map_local_sync_error_kind_to_claude_type(kind: LocalCoreSyncErrorKind) -> &'s LocalCoreSyncErrorKind::InvalidRequest | LocalCoreSyncErrorKind::ContextLengthExceeded => { "invalid_request_error" } + LocalCoreSyncErrorKind::RequestTooLarge => "request_too_large", LocalCoreSyncErrorKind::Authentication => "authentication_error", LocalCoreSyncErrorKind::PermissionDenied => "permission_error", LocalCoreSyncErrorKind::NotFound => "not_found_error", LocalCoreSyncErrorKind::RateLimit => "rate_limit_error", - LocalCoreSyncErrorKind::Overloaded | LocalCoreSyncErrorKind::ServerError => "api_error", + LocalCoreSyncErrorKind::Overloaded => "overloaded_error", + LocalCoreSyncErrorKind::ServerError => "api_error", } } fn map_local_sync_error_kind_to_gemini_code(kind: LocalCoreSyncErrorKind) -> u16 { match kind { - LocalCoreSyncErrorKind::InvalidRequest | LocalCoreSyncErrorKind::ContextLengthExceeded => { - 400 - } + LocalCoreSyncErrorKind::InvalidRequest + | LocalCoreSyncErrorKind::ContextLengthExceeded + | LocalCoreSyncErrorKind::RequestTooLarge => 400, LocalCoreSyncErrorKind::Authentication => 401, LocalCoreSyncErrorKind::PermissionDenied => 403, LocalCoreSyncErrorKind::NotFound => 404, @@ -127,9 +132,9 @@ fn map_local_sync_error_kind_to_gemini_code(kind: LocalCoreSyncErrorKind) -> u16 fn map_local_sync_error_kind_to_gemini_status(kind: LocalCoreSyncErrorKind) -> &'static str { match kind { - LocalCoreSyncErrorKind::InvalidRequest | LocalCoreSyncErrorKind::ContextLengthExceeded => { - "INVALID_ARGUMENT" - } + LocalCoreSyncErrorKind::InvalidRequest + | LocalCoreSyncErrorKind::ContextLengthExceeded + | LocalCoreSyncErrorKind::RequestTooLarge => "INVALID_ARGUMENT", LocalCoreSyncErrorKind::Authentication => "UNAUTHENTICATED", LocalCoreSyncErrorKind::PermissionDenied => "PERMISSION_DENIED", LocalCoreSyncErrorKind::NotFound => "NOT_FOUND", @@ -176,6 +181,51 @@ mod tests { assert_eq!(body["error"]["code"], "upstream_unavailable"); } + #[test] + fn builds_claude_overloaded_error_body() { + let body = build_core_error_body_for_client_format( + "claude:messages", + "busy", + None, + LocalCoreSyncErrorKind::Overloaded, + ) + .expect("body should build"); + + assert_eq!(body["type"], "error"); + assert_eq!(body["error"]["type"], "overloaded_error"); + } + + #[test] + fn maps_request_too_large_for_each_client_format() { + let claude = build_core_error_body_for_client_format( + "claude:messages", + "too large", + None, + LocalCoreSyncErrorKind::RequestTooLarge, + ) + .expect("Claude body should build"); + assert_eq!(claude["error"]["type"], "request_too_large"); + + let openai = build_core_error_body_for_client_format( + "openai:chat", + "too large", + None, + LocalCoreSyncErrorKind::RequestTooLarge, + ) + .expect("OpenAI body should build"); + assert_eq!(openai["error"]["type"], "context_length_exceeded"); + + let gemini = build_core_error_body_for_client_format( + "gemini:generate_content", + "too large", + None, + LocalCoreSyncErrorKind::RequestTooLarge, + ) + .expect("Gemini body should build"); + assert_eq!(gemini["error"]["code"], 400); + assert_eq!(gemini["error"]["status"], "INVALID_ARGUMENT"); + } + #[test] fn recognizes_finalize_kind_and_success_mapping() { assert!(is_core_error_finalize_kind("openai_chat_sync_finalize")); diff --git a/crates/aether-ai/formats/src/formats/shared/passthrough.rs b/crates/aether-ai/formats/src/formats/shared/passthrough.rs index 4c1fd03ce..f1b7ccbad 100644 --- a/crates/aether-ai/formats/src/formats/shared/passthrough.rs +++ b/crates/aether-ai/formats/src/formats/shared/passthrough.rs @@ -1,11 +1,12 @@ use crate::contracts::{ - CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND, - CLAUDE_CLI_SYNC_PLAN_KIND, GEMINI_CHAT_STREAM_PLAN_KIND, GEMINI_CHAT_SYNC_PLAN_KIND, - GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_SYNC_PLAN_KIND, GEMINI_EMBEDDING_SYNC_PLAN_KIND, - GEMINI_EMBEDDING_SYNC_SUCCESS_REPORT_KIND, GEMINI_INTERACTIONS_STREAM_PLAN_KIND, - GEMINI_INTERACTIONS_STREAM_SUCCESS_REPORT_KIND, GEMINI_INTERACTIONS_SYNC_PLAN_KIND, - GEMINI_INTERACTIONS_SYNC_SUCCESS_REPORT_KIND, OPENAI_EMBEDDING_SYNC_PLAN_KIND, - OPENAI_RERANK_SYNC_PLAN_KIND, OPENAI_SEARCH_SYNC_PLAN_KIND, + ApiOperation, CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND, + CLAUDE_CLI_STREAM_PLAN_KIND, CLAUDE_CLI_SYNC_PLAN_KIND, CLAUDE_COUNT_TOKENS_SYNC_PLAN_KIND, + CLAUDE_COUNT_TOKENS_SYNC_SUCCESS_REPORT_KIND, GEMINI_CHAT_STREAM_PLAN_KIND, + GEMINI_CHAT_SYNC_PLAN_KIND, GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_SYNC_PLAN_KIND, + GEMINI_EMBEDDING_SYNC_PLAN_KIND, GEMINI_EMBEDDING_SYNC_SUCCESS_REPORT_KIND, + GEMINI_INTERACTIONS_STREAM_PLAN_KIND, GEMINI_INTERACTIONS_STREAM_SUCCESS_REPORT_KIND, + GEMINI_INTERACTIONS_SYNC_PLAN_KIND, GEMINI_INTERACTIONS_SYNC_SUCCESS_REPORT_KIND, + OPENAI_EMBEDDING_SYNC_PLAN_KIND, OPENAI_RERANK_SYNC_PLAN_KIND, OPENAI_SEARCH_SYNC_PLAN_KIND, OPENAI_SEARCH_SYNC_SUCCESS_REPORT_KIND, }; @@ -22,6 +23,7 @@ pub struct LocalSameFormatProviderSpec { pub report_kind: &'static str, pub family: LocalSameFormatProviderFamily, pub require_streaming: bool, + pub operation: Option, } pub fn resolve_sync_spec(plan_kind: &str) -> Option { @@ -32,6 +34,7 @@ pub fn resolve_sync_spec(plan_kind: &str) -> Option report_kind: "claude_chat_sync_success", family: LocalSameFormatProviderFamily::Standard, require_streaming: false, + operation: Some(ApiOperation::ClaudeMessagesCreate), }), CLAUDE_CLI_SYNC_PLAN_KIND => Some(LocalSameFormatProviderSpec { api_format: "claude:messages", @@ -39,6 +42,15 @@ pub fn resolve_sync_spec(plan_kind: &str) -> Option report_kind: "claude_cli_sync_success", family: LocalSameFormatProviderFamily::Standard, require_streaming: false, + operation: Some(ApiOperation::ClaudeMessagesCreate), + }), + CLAUDE_COUNT_TOKENS_SYNC_PLAN_KIND => Some(LocalSameFormatProviderSpec { + api_format: "claude:messages", + decision_kind: CLAUDE_COUNT_TOKENS_SYNC_PLAN_KIND, + report_kind: CLAUDE_COUNT_TOKENS_SYNC_SUCCESS_REPORT_KIND, + family: LocalSameFormatProviderFamily::Standard, + require_streaming: false, + operation: Some(ApiOperation::ClaudeCountTokens), }), GEMINI_CHAT_SYNC_PLAN_KIND => Some(LocalSameFormatProviderSpec { api_format: "gemini:generate_content", @@ -46,6 +58,7 @@ pub fn resolve_sync_spec(plan_kind: &str) -> Option report_kind: "gemini_chat_sync_success", family: LocalSameFormatProviderFamily::Gemini, require_streaming: false, + operation: None, }), GEMINI_CLI_SYNC_PLAN_KIND => Some(LocalSameFormatProviderSpec { api_format: "gemini:generate_content", @@ -53,6 +66,7 @@ pub fn resolve_sync_spec(plan_kind: &str) -> Option report_kind: "gemini_cli_sync_success", family: LocalSameFormatProviderFamily::Gemini, require_streaming: false, + operation: None, }), GEMINI_EMBEDDING_SYNC_PLAN_KIND => Some(LocalSameFormatProviderSpec { api_format: "gemini:embedding", @@ -60,6 +74,7 @@ pub fn resolve_sync_spec(plan_kind: &str) -> Option report_kind: GEMINI_EMBEDDING_SYNC_SUCCESS_REPORT_KIND, family: LocalSameFormatProviderFamily::Gemini, require_streaming: false, + operation: None, }), GEMINI_INTERACTIONS_SYNC_PLAN_KIND => Some(LocalSameFormatProviderSpec { api_format: "gemini:interactions", @@ -67,6 +82,7 @@ pub fn resolve_sync_spec(plan_kind: &str) -> Option report_kind: GEMINI_INTERACTIONS_SYNC_SUCCESS_REPORT_KIND, family: LocalSameFormatProviderFamily::Gemini, require_streaming: false, + operation: None, }), OPENAI_EMBEDDING_SYNC_PLAN_KIND => Some(LocalSameFormatProviderSpec { api_format: "openai:embedding", @@ -74,6 +90,7 @@ pub fn resolve_sync_spec(plan_kind: &str) -> Option report_kind: "openai_embedding_sync_success", family: LocalSameFormatProviderFamily::Standard, require_streaming: false, + operation: None, }), OPENAI_RERANK_SYNC_PLAN_KIND => Some(LocalSameFormatProviderSpec { api_format: "openai:rerank", @@ -81,6 +98,7 @@ pub fn resolve_sync_spec(plan_kind: &str) -> Option report_kind: "openai_rerank_sync_success", family: LocalSameFormatProviderFamily::Standard, require_streaming: false, + operation: None, }), OPENAI_SEARCH_SYNC_PLAN_KIND => Some(LocalSameFormatProviderSpec { api_format: "openai:search", @@ -88,6 +106,7 @@ pub fn resolve_sync_spec(plan_kind: &str) -> Option report_kind: OPENAI_SEARCH_SYNC_SUCCESS_REPORT_KIND, family: LocalSameFormatProviderFamily::Standard, require_streaming: false, + operation: None, }), _ => None, } @@ -101,6 +120,7 @@ pub fn resolve_stream_spec(plan_kind: &str) -> Option Some(LocalSameFormatProviderSpec { api_format: "claude:messages", @@ -108,6 +128,7 @@ pub fn resolve_stream_spec(plan_kind: &str) -> Option Some(LocalSameFormatProviderSpec { api_format: "gemini:generate_content", @@ -115,6 +136,7 @@ pub fn resolve_stream_spec(plan_kind: &str) -> Option Some(LocalSameFormatProviderSpec { api_format: "gemini:generate_content", @@ -122,6 +144,7 @@ pub fn resolve_stream_spec(plan_kind: &str) -> Option Some(LocalSameFormatProviderSpec { api_format: "gemini:interactions", @@ -129,6 +152,7 @@ pub fn resolve_stream_spec(plan_kind: &str) -> Option None, } diff --git a/crates/aether-ai/formats/src/formats/shared/routing.rs b/crates/aether-ai/formats/src/formats/shared/routing.rs index fab37faa0..a7f03a4df 100644 --- a/crates/aether-ai/formats/src/formats/shared/routing.rs +++ b/crates/aether-ai/formats/src/formats/shared/routing.rs @@ -2,20 +2,21 @@ use http::Method; use url::form_urlencoded; use crate::contracts::{ - CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND, - CLAUDE_CLI_SYNC_PLAN_KIND, GEMINI_CHAT_STREAM_PLAN_KIND, GEMINI_CHAT_SYNC_PLAN_KIND, - GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_SYNC_PLAN_KIND, GEMINI_EMBEDDING_SYNC_PLAN_KIND, - GEMINI_FILES_DELETE_PLAN_KIND, GEMINI_FILES_DOWNLOAD_PLAN_KIND, GEMINI_FILES_GET_PLAN_KIND, - GEMINI_FILES_LIST_PLAN_KIND, GEMINI_FILES_UPLOAD_PLAN_KIND, - GEMINI_INTERACTIONS_STREAM_PLAN_KIND, GEMINI_INTERACTIONS_SYNC_PLAN_KIND, - GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND, GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND, - OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND, OPENAI_EMBEDDING_SYNC_PLAN_KIND, - OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND, OPENAI_RERANK_SYNC_PLAN_KIND, - OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND, - OPENAI_RESPONSES_STREAM_PLAN_KIND, OPENAI_RESPONSES_SYNC_PLAN_KIND, - OPENAI_SEARCH_SYNC_PLAN_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, - OPENAI_VIDEO_CONTENT_PLAN_KIND, OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND, - OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND, OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND, + ClientSurface, CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND, + CLAUDE_CLI_STREAM_PLAN_KIND, CLAUDE_CLI_SYNC_PLAN_KIND, CLAUDE_COUNT_TOKENS_SYNC_PLAN_KIND, + GEMINI_CHAT_STREAM_PLAN_KIND, GEMINI_CHAT_SYNC_PLAN_KIND, GEMINI_CLI_STREAM_PLAN_KIND, + GEMINI_CLI_SYNC_PLAN_KIND, GEMINI_EMBEDDING_SYNC_PLAN_KIND, GEMINI_FILES_DELETE_PLAN_KIND, + GEMINI_FILES_DOWNLOAD_PLAN_KIND, GEMINI_FILES_GET_PLAN_KIND, GEMINI_FILES_LIST_PLAN_KIND, + GEMINI_FILES_UPLOAD_PLAN_KIND, GEMINI_INTERACTIONS_STREAM_PLAN_KIND, + GEMINI_INTERACTIONS_SYNC_PLAN_KIND, GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND, + GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND, + OPENAI_EMBEDDING_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND, + OPENAI_RERANK_SYNC_PLAN_KIND, OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, + OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND, OPENAI_RESPONSES_STREAM_PLAN_KIND, + OPENAI_RESPONSES_SYNC_PLAN_KIND, OPENAI_SEARCH_SYNC_PLAN_KIND, + OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_VIDEO_CONTENT_PLAN_KIND, + OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND, + OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND, }; use crate::formats::openai::image::request::is_openai_image_stream_request; @@ -26,6 +27,26 @@ pub fn resolve_execution_runtime_stream_plan_kind( request_auth_channel: Option<&str>, method: &Method, path: &str, +) -> Option<&'static str> { + resolve_execution_runtime_stream_plan_kind_with_client_surface( + route_class, + route_family, + route_kind, + legacy_client_surface_from_auth_channel(request_auth_channel), + request_auth_channel, + method, + path, + ) +} + +pub fn resolve_execution_runtime_stream_plan_kind_with_client_surface( + route_class: Option<&str>, + route_family: Option<&str>, + route_kind: Option<&str>, + client_surface: Option, + request_auth_channel: Option<&str>, + method: &Method, + path: &str, ) -> Option<&'static str> { if route_class != Some("ai_public") { return None; @@ -53,7 +74,7 @@ pub fn resolve_execution_runtime_stream_plan_kind( && path == "/v1/messages" { return Some(resolve_claude_messages_plan_kind( - request_auth_channel, + client_surface, CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND, )); @@ -121,6 +142,26 @@ pub fn resolve_execution_runtime_sync_plan_kind( request_auth_channel: Option<&str>, method: &Method, path: &str, +) -> Option<&'static str> { + resolve_execution_runtime_sync_plan_kind_with_client_surface( + route_class, + route_family, + route_kind, + legacy_client_surface_from_auth_channel(request_auth_channel), + request_auth_channel, + method, + path, + ) +} + +pub fn resolve_execution_runtime_sync_plan_kind_with_client_surface( + route_class: Option<&str>, + route_family: Option<&str>, + route_kind: Option<&str>, + client_surface: Option, + request_auth_channel: Option<&str>, + method: &Method, + path: &str, ) -> Option<&'static str> { if route_class != Some("ai_public") { return None; @@ -248,13 +289,21 @@ pub fn resolve_execution_runtime_sync_plan_kind( return Some(OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND); } + if route_family == Some("claude") + && route_kind == Some("count_tokens") + && *method == Method::POST + && path == "/v1/messages/count_tokens" + { + return Some(CLAUDE_COUNT_TOKENS_SYNC_PLAN_KIND); + } + if route_family == Some("claude") && is_claude_messages_route_kind(route_kind) && *method == Method::POST && path == "/v1/messages" { return Some(resolve_claude_messages_plan_kind( - request_auth_channel, + client_surface, CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CLI_SYNC_PLAN_KIND, )); @@ -313,17 +362,27 @@ fn is_gemini_generate_content_route_kind(route_kind: Option<&str>) -> bool { } fn resolve_claude_messages_plan_kind( - request_auth_channel: Option<&str>, + client_surface: Option, chat_plan_kind: &'static str, cli_plan_kind: &'static str, ) -> &'static str { - if request_auth_channel == Some("bearer_like") { + if client_surface == Some(ClientSurface::ClaudeCode) { cli_plan_kind } else { chat_plan_kind } } +fn legacy_client_surface_from_auth_channel( + request_auth_channel: Option<&str>, +) -> Option { + Some(if request_auth_channel == Some("bearer_like") { + ClientSurface::ClaudeCode + } else { + ClientSurface::GenericCompatible + }) +} + fn resolve_gemini_generate_content_plan_kind( request_auth_channel: Option<&str>, chat_plan_kind: &'static str, @@ -447,6 +506,7 @@ pub fn supports_sync_execution_decision_kind(plan_kind: &str) -> bool { | OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND | CLAUDE_CHAT_SYNC_PLAN_KIND | CLAUDE_CLI_SYNC_PLAN_KIND + | CLAUDE_COUNT_TOKENS_SYNC_PLAN_KIND | GEMINI_CHAT_SYNC_PLAN_KIND | GEMINI_CLI_SYNC_PLAN_KIND | GEMINI_EMBEDDING_SYNC_PLAN_KIND diff --git a/crates/aether-ai/formats/src/formats/shared/stream_rewrite.rs b/crates/aether-ai/formats/src/formats/shared/stream_rewrite.rs index 351f57a49..537fe9639 100644 --- a/crates/aether-ai/formats/src/formats/shared/stream_rewrite.rs +++ b/crates/aether-ai/formats/src/formats/shared/stream_rewrite.rs @@ -10,6 +10,7 @@ use crate::formats::shared::response::{ }; use crate::formats::shared::sse::encode_json_sse; use crate::formats::shared::stream_core::StreamingStandardFormatMatrix; +use crate::formats::shared::sync_products::anthropic_legacy_compatibility_enabled; use crate::formats::shared::AiSurfaceFinalizeError; use crate::provider_compat::kiro_stream::KiroToClaudeCliStreamState; use crate::provider_compat::private_envelope::transform_provider_private_stream_line; @@ -117,7 +118,12 @@ pub fn resolve_finalize_stream_rewrite_mode( if provider_stream_event_api_format == "claude:messages" && client_api_format == "claude:messages" { - return Some(FinalizeStreamRewriteMode::ClaudeReadToolSanitize); + return if anthropic_legacy_compatibility_enabled(report_context) { + Some(FinalizeStreamRewriteMode::ClaudeReadToolSanitize) + } else { + model_directive_display_model_from_report_context(report_context) + .map(|_| FinalizeStreamRewriteMode::ModelDirectiveDisplay) + }; } return model_directive_display_model_from_report_context(report_context) .map(|_| FinalizeStreamRewriteMode::ModelDirectiveDisplay); @@ -144,7 +150,11 @@ pub fn resolve_finalize_stream_rewrite_mode( ) { if provider_stream_event_api_format == "claude:messages" { - return Some(FinalizeStreamRewriteMode::ClaudeReadToolSanitize); + return Some(if anthropic_legacy_compatibility_enabled(report_context) { + FinalizeStreamRewriteMode::ClaudeReadToolSanitize + } else { + FinalizeStreamRewriteMode::ModelDirectiveDisplay + }); } return Some(FinalizeStreamRewriteMode::ModelDirectiveDisplay); } @@ -152,7 +162,8 @@ pub fn resolve_finalize_stream_rewrite_mode( if provider_stream_event_api_format == "claude:messages" && client_api_format == "claude:messages" { - return Some(FinalizeStreamRewriteMode::ClaudeReadToolSanitize); + return anthropic_legacy_compatibility_enabled(report_context) + .then_some(FinalizeStreamRewriteMode::ClaudeReadToolSanitize); } if is_same_format_family( @@ -1160,13 +1171,19 @@ data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_123\",\"objec "needs_conversion": true, "model": "claude-sonnet-4.5-high", "mapped_model": "claude-sonnet-4.5", + "provider_type": "claude_code", + "anthropic_compatibility_profile": "native_transparent", }); + assert_eq!( + resolve_finalize_stream_rewrite_mode(&report_context), + Some(FinalizeStreamRewriteMode::ModelDirectiveDisplay) + ); let mut rewriter = maybe_build_ai_surface_stream_rewriter(Some(&report_context)) .expect("rewriter should exist"); let output = rewriter .push_chunk( b"event: content_block_delta\n\ -data: {\"type\":\"content_block_delta\",\"index\":1,\"delta\":{\"type\":\"thinking_delta\",\"thinking\":\"Let me reason...\"}}\n\n", +data: {\"type\":\"content_block_delta\",\"index\":1,\"future_event_field\":{\"keep\":true},\"delta\":{\"type\":\"thinking_delta\",\"thinking\":\"Let me reason...\",\"future_delta_field\":42}}\n\n", ) .expect("rewrite should succeed"); let output = String::from_utf8(output).expect("output should be utf8"); @@ -1175,20 +1192,51 @@ data: {\"type\":\"content_block_delta\",\"index\":1,\"delta\":{\"type\":\"thinki assert!(output.contains("event: content_block_delta")); assert!(output.contains("\"thinking\":\"Let me reason...\"")); assert!(output.contains("\"type\":\"thinking_delta\"")); + assert!(output.contains("\"future_event_field\":{\"keep\":true}")); + assert!(output.contains("\"future_delta_field\":42")); } #[test] - fn same_format_claude_uses_read_tool_sanitizer_without_display_model() { - // Claude→Claude needs a narrow sanitizer for Claude Code Read input. + fn native_same_format_claude_without_display_model_is_verbatim() { let report_context = json!({ "provider_api_format": "claude:messages", "client_api_format": "claude:messages", "needs_conversion": true, }); - assert_eq!( - resolve_finalize_stream_rewrite_mode(&report_context), - Some(FinalizeStreamRewriteMode::ClaudeReadToolSanitize) - ); + assert_eq!(resolve_finalize_stream_rewrite_mode(&report_context), None); + assert!(maybe_build_ai_surface_stream_rewriter(Some(&report_context)).is_none()); + } + + #[test] + fn claude_code_legacy_profile_and_provider_fallback_enable_read_tool_sanitizer() { + for report_context in [ + json!({ + "provider_api_format": "claude:messages", + "client_api_format": "claude:messages", + "needs_conversion": false, + "anthropic_compatibility_profile": "claude_code_legacy", + }), + json!({ + "provider_api_format": "claude:messages", + "client_api_format": "claude:messages", + "needs_conversion": false, + "provider_type": "claude_code", + }), + ] { + assert_eq!( + resolve_finalize_stream_rewrite_mode(&report_context), + Some(FinalizeStreamRewriteMode::ClaudeReadToolSanitize) + ); + } + + let explicit_native = json!({ + "provider_api_format": "claude:messages", + "client_api_format": "claude:messages", + "needs_conversion": false, + "provider_type": "claude_code", + "anthropic_compatibility_profile": "native_transparent", + }); + assert_eq!(resolve_finalize_stream_rewrite_mode(&explicit_native), None); } #[test] @@ -1197,6 +1245,7 @@ data: {\"type\":\"content_block_delta\",\"index\":1,\"delta\":{\"type\":\"thinki "provider_api_format": "claude:messages", "client_api_format": "claude:messages", "needs_conversion": false, + "anthropic_compatibility_profile": "claude_code_legacy", }); let mut rewriter = maybe_build_ai_surface_stream_rewriter(Some(&report_context)) .expect("same-format claude sanitizer should exist"); @@ -1219,6 +1268,7 @@ data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\": "provider_api_format": "claude:messages", "client_api_format": "claude:messages", "needs_conversion": false, + "anthropic_compatibility_profile": "claude_code_legacy", }); let mut rewriter = maybe_build_ai_surface_stream_rewriter(Some(&report_context)) .expect("same-format claude sanitizer should exist"); @@ -1263,6 +1313,7 @@ data: {\"type\":\"content_block_stop\",\"index\":0}\n\n", "provider_api_format": "claude:messages", "client_api_format": "claude:messages", "needs_conversion": false, + "anthropic_compatibility_profile": "claude_code_legacy", }); let mut rewriter = maybe_build_ai_surface_stream_rewriter(Some(&report_context)) .expect("same-format claude sanitizer should exist"); 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 373a00717..e2e451c23 100644 --- a/crates/aether-ai/formats/src/formats/shared/sync_products.rs +++ b/crates/aether-ai/formats/src/formats/shared/sync_products.rs @@ -25,10 +25,7 @@ use crate::formats::claude::messages::stream::ClaudeProviderState; use crate::formats::gemini::generate_content::stream::GeminiProviderState; use crate::formats::openai::chat::stream::{OpenAIChatProviderState, OpenAIResponsesProviderState}; use crate::formats::shared::model_directives::model_directive_display_model_from_report_context; -use crate::formats::shared::response::{ - remove_empty_pages_from_tool_arguments, remove_empty_pages_from_tool_input_value, - sanitize_claude_read_tool_inputs, -}; +use crate::formats::shared::response::sanitize_claude_read_tool_inputs; use crate::formats::shared::stream_core::common::{ content_part_from_openai_image_generation_item, gemini_usage_metadata_from_usage, map_openai_finish_reason_to_gemini, parse_json_arguments_value, @@ -685,7 +682,9 @@ fn maybe_build_standard_same_format_sync_body( } let mut body_json = body_json.clone(); - if expected_api_format == "claude:messages" { + if expected_api_format == "claude:messages" + && anthropic_legacy_compatibility_enabled(report_context) + { sanitize_claude_read_tool_inputs(&mut body_json); } @@ -744,13 +743,18 @@ fn maybe_build_standard_same_format_stream_sync_body( let body_bytes = base64::engine::general_purpose::STANDARD.decode(body_base64)?; let provider_stream_event_api_format = provider_stream_event_api_format_for_report_context(report_context, &provider_api_format); - let Some(body) = try_aggregate_standard_chat_stream_sync_response( + let Some(mut body) = try_aggregate_standard_chat_stream_sync_response( &body_bytes, &provider_stream_event_api_format, )? else { return Ok(None); }; + if provider_stream_event_api_format == "claude:messages" + && anthropic_legacy_compatibility_enabled(report_context) + { + sanitize_claude_read_tool_inputs(&mut body); + } if api_format_is_gemini_generate_content(&provider_stream_event_api_format) && !gemini_generate_content_body_has_visible_output(&body) { @@ -772,6 +776,21 @@ fn maybe_build_standard_same_format_stream_sync_body( ))) } +pub(super) fn anthropic_legacy_compatibility_enabled(report_context: &Value) -> bool { + if let Some(profile) = report_context.get("anthropic_compatibility_profile") { + return profile + .as_str() + .map(str::trim) + .is_some_and(|profile| profile.eq_ignore_ascii_case("claude_code_legacy")); + } + + report_context + .get("provider_type") + .and_then(Value::as_str) + .map(str::trim) + .is_some_and(|provider_type| provider_type.eq_ignore_ascii_case("claude_code")) +} + fn maybe_build_openai_responses_same_family_sync_body( report_kind: &str, status_code: u16, @@ -3400,22 +3419,9 @@ pub fn aggregate_claude_stream_sync_response(body: &[u8]) -> Option { } } "tool_use" => { - let tool_name = block - .get("name") - .and_then(Value::as_str) - .unwrap_or_default() - .to_string(); - if let Some(input) = block.get("input") { - let sanitized = remove_empty_pages_from_tool_input_value(&tool_name, input); - if sanitized != *input { - block.insert("input".to_string(), sanitized); - } - } if !state.partial_json.is_empty() { - let arguments = - remove_empty_pages_from_tool_arguments(&tool_name, &state.partial_json); - let input = serde_json::from_str::(&arguments) - .unwrap_or(Value::String(arguments)); + let input = serde_json::from_str::(&state.partial_json) + .unwrap_or(Value::String(state.partial_json)); block.insert("input".to_string(), input); } } @@ -4094,7 +4100,7 @@ mod tests { } #[test] - fn aggregates_claude_stream_removes_empty_pages_from_tool_input() { + fn aggregates_native_claude_stream_preserves_empty_pages_from_tool_input() { let body = concat!( "event: message_start\n", "data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_123\",\"type\":\"message\",\"role\":\"assistant\",\"model\":\"claude-sonnet-4-5\",\"content\":[],\"stop_reason\":null,\"stop_sequence\":null}}\n\n", @@ -4118,12 +4124,13 @@ mod tests { "file_path": "/tmp/a.txt", "offset": 1, "limit": 20, + "pages": "", }) ); } #[test] - fn aggregates_claude_stream_removes_empty_pages_from_start_tool_input() { + fn aggregates_native_claude_stream_preserves_empty_pages_from_start_tool_input() { let body = concat!( "event: message_start\n", "data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_123\",\"type\":\"message\",\"role\":\"assistant\",\"model\":\"claude-sonnet-4-5\",\"content\":[],\"stop_reason\":null,\"stop_sequence\":null}}\n\n", @@ -4143,6 +4150,7 @@ mod tests { json!({ "file_path": "/tmp/a.txt", "limit": 20, + "pages": "", }) ); } @@ -4496,8 +4504,8 @@ mod tests { } #[test] - fn same_format_claude_sync_body_sanitizes_read_tool_input() { - let report_context = json!({ + fn same_format_claude_sync_body_sanitizes_only_for_legacy_compatibility() { + let native_report_context = json!({ "provider_api_format": "claude:messages", "client_api_format": "claude:messages", "needs_conversion": false, @@ -4507,6 +4515,7 @@ mod tests { "type": "message", "role": "assistant", "model": "claude-sonnet-4-6", + "future_message_field": {"preserve": true}, "content": [ { "type": "tool_use", @@ -4515,8 +4524,10 @@ mod tests { "input": { "file_path": "/tmp/a.txt", "limit": 20, - "pages": "" - } + "pages": "", + "future_input_field": 42 + }, + "future_block_field": "keep" }, { "type": "tool_use", @@ -4530,10 +4541,28 @@ mod tests { ] }); + let native_body = maybe_build_standard_same_format_sync_body_from_normalized_payload( + "claude_chat_sync_finalize", + 200, + Some(&native_report_context), + Some(&provider_body_json), + None, + ) + .expect("native same-format sync body should succeed") + .expect("native body should exist"); + assert_eq!(native_body, provider_body_json); + + let legacy_report_context = json!({ + "provider_api_format": "claude:messages", + "client_api_format": "claude:messages", + "needs_conversion": false, + "anthropic_compatibility_profile": "claude_code_legacy", + }); + let body_json = maybe_build_standard_same_format_sync_body_from_normalized_payload( "claude_chat_sync_finalize", 200, - Some(&report_context), + Some(&legacy_report_context), Some(&provider_body_json), None, ) @@ -4545,8 +4574,11 @@ mod tests { json!({ "file_path": "/tmp/a.txt", "limit": 20, + "future_input_field": 42, }) ); + assert_eq!(body_json["content"][0]["future_block_field"], "keep"); + assert_eq!(body_json["future_message_field"]["preserve"], true); assert_eq!( body_json["content"][1]["input"], json!({ @@ -4556,6 +4588,60 @@ mod tests { ); } + #[test] + fn same_format_claude_stream_sync_body_sanitizes_only_for_legacy_compatibility() { + let stream = concat!( + "event: message_start\n", + "data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_read\",\"type\":\"message\",\"role\":\"assistant\",\"model\":\"claude-sonnet-4-6\",\"content\":[],\"future_message_field\":{\"keep\":true}}}\n\n", + "event: content_block_start\n", + "data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"tool_use\",\"id\":\"toolu_read\",\"name\":\"Read\",\"input\":{},\"future_block_field\":\"keep\"}}\n\n", + "event: content_block_delta\n", + "data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"{\\\"file_path\\\":\\\"/tmp/a.txt\\\",\\\"pages\\\":\\\"\\\"}\"}}\n\n", + "event: content_block_stop\n", + "data: {\"type\":\"content_block_stop\",\"index\":0}\n\n", + "event: message_stop\n", + "data: {\"type\":\"message_stop\"}\n\n", + ); + let body_base64 = base64::engine::general_purpose::STANDARD.encode(stream); + let native_report_context = json!({ + "provider_api_format": "claude:messages", + "client_api_format": "claude:messages", + "needs_conversion": false, + }); + + let native_body = maybe_build_standard_same_format_sync_body_from_normalized_payload( + "claude_chat_sync_finalize", + 200, + Some(&native_report_context), + None, + Some(&body_base64), + ) + .expect("native stream sync body should succeed") + .expect("native stream sync body should exist"); + assert_eq!(native_body["content"][0]["input"]["pages"], ""); + assert_eq!(native_body["content"][0]["future_block_field"], "keep"); + assert_eq!(native_body["future_message_field"]["keep"], true); + + let legacy_report_context = json!({ + "provider_api_format": "claude:messages", + "client_api_format": "claude:messages", + "needs_conversion": false, + "provider_type": "claude_code", + }); + let legacy_body = maybe_build_standard_same_format_sync_body_from_normalized_payload( + "claude_chat_sync_finalize", + 200, + Some(&legacy_report_context), + None, + Some(&body_base64), + ) + .expect("legacy stream sync body should succeed") + .expect("legacy stream sync body should exist"); + assert!(legacy_body["content"][0]["input"].get("pages").is_none()); + assert_eq!(legacy_body["content"][0]["future_block_field"], "keep"); + assert_eq!(legacy_body["future_message_field"]["keep"], true); + } + #[test] fn same_format_sync_response_restores_model_directive_display_model() { let report_context = json!({ diff --git a/crates/aether-ai/formats/src/formats/shared/sync_to_stream.rs b/crates/aether-ai/formats/src/formats/shared/sync_to_stream.rs index 1533a4a60..bd5291d71 100644 --- a/crates/aether-ai/formats/src/formats/shared/sync_to_stream.rs +++ b/crates/aether-ai/formats/src/formats/shared/sync_to_stream.rs @@ -1814,6 +1814,12 @@ mod tests { "event: message_stop\n", "data: {\"type\":\"message_stop\"}\n\n", ); + let report_context = json!({ + "provider_api_format": "claude:messages", + "client_api_format": "claude:messages", + "needs_conversion": false, + "anthropic_compatibility_profile": "claude_code_legacy", + }); let outcome = maybe_bridge_standard_sync_json_to_stream( &json!({ "status_code": 200, @@ -1825,7 +1831,7 @@ mod tests { }), "openai:responses", "claude:messages", - None, + Some(&report_context), ) .expect("bridge should succeed") .expect("capture should bridge"); diff --git a/crates/aether-ai/formats/src/lib.rs b/crates/aether-ai/formats/src/lib.rs index 892943480..17f1654ff 100644 --- a/crates/aether-ai/formats/src/lib.rs +++ b/crates/aether-ai/formats/src/lib.rs @@ -6,6 +6,8 @@ pub mod formats; pub mod protocol; pub mod provider_compat; +pub use contracts::{ApiOperation, ClientSurface}; + pub use formats::context::{ ConversionFieldRecord, ConversionFieldStatus, ConversionReport, Converted, FormatContext, FormatError, diff --git a/crates/aether-ai/serving/Cargo.toml b/crates/aether-ai/serving/Cargo.toml index 3b779d98c..1860c104b 100644 --- a/crates/aether-ai/serving/Cargo.toml +++ b/crates/aether-ai/serving/Cargo.toml @@ -13,6 +13,7 @@ aether-data-contracts.workspace = true aether-pool-core.workspace = true aether-scheduler-core.workspace = true async-trait.workspace = true +base64.workspace = true http.workspace = true serde.workspace = true serde_json.workspace = true diff --git a/crates/aether-ai/serving/src/attempt_loop.rs b/crates/aether-ai/serving/src/attempt_loop.rs index 60684ae75..c2df49cfc 100644 --- a/crates/aether-ai/serving/src/attempt_loop.rs +++ b/crates/aether-ai/serving/src/attempt_loop.rs @@ -18,10 +18,45 @@ pub trait AiExecutionAttempt { #[derive(Debug)] pub enum AiAttemptLoopOutcome { Responded(Response), + Deferred(Response), Exhausted(Exhaustion), NoPath, } +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub enum AiAttemptRetryScope { + #[default] + Candidate, + Credential, + Endpoint, + Provider, +} + +#[derive(Debug)] +pub enum AiAttemptExecutionOutcome { + Responded(Response), + Retry { + scope: AiAttemptRetryScope, + fallback_response: Option, + }, +} + +impl AiAttemptExecutionOutcome { + pub fn retry(scope: AiAttemptRetryScope) -> Self { + Self::Retry { + scope, + fallback_response: None, + } + } + + pub fn from_optional_response(response: Option) -> Self { + match response { + Some(response) => Self::Responded(response), + None => Self::retry(AiAttemptRetryScope::Candidate), + } + } +} + #[async_trait] pub trait AiAttemptLoopPort: Send + Sync where @@ -34,7 +69,7 @@ where async fn execute_attempt( &self, attempt: &Attempt, - ) -> Result, Self::Error>; + ) -> Result, Self::Error>; async fn should_skip_attempt(&self, _attempt: &Attempt) -> Result { Ok(false) @@ -67,32 +102,52 @@ where { let mut remaining = attempts.into_iter(); let mut last_attempted = None; + let mut retry_filters: Vec = Vec::new(); + let mut fallback_response = None; while let Some(attempt) = remaining.next() { - if port.should_skip_attempt(&attempt).await? { + if retry_filters.iter().any(|filter| filter.matches(&attempt)) + || port.should_skip_attempt(&attempt).await? + { port.mark_unused_attempts(vec![attempt]).await?; continue; } port.record_attempt_started(&attempt).await?; - let response = match port.execute_attempt(&attempt).await { - Ok(response) => response, + let execution = match port.execute_attempt(&attempt).await { + Ok(execution) => execution, Err(err) => { port.mark_unused_attempts(remaining.collect()).await?; return Err(err); } }; - if let Some(response) = response { - port.mark_unused_attempts(remaining.collect()).await?; - return Ok(AiAttemptLoopOutcome::Responded(response)); + match execution { + AiAttemptExecutionOutcome::Responded(response) => { + port.mark_unused_attempts(remaining.collect()).await?; + return Ok(AiAttemptLoopOutcome::Responded(response)); + } + AiAttemptExecutionOutcome::Retry { + scope, + fallback_response: attempt_fallback_response, + } => { + port.record_attempt_failed(&attempt).await?; + if attempt_fallback_response.is_some() { + fallback_response = attempt_fallback_response; + } + if scope != AiAttemptRetryScope::Candidate { + retry_filters.push(AiAttemptRetryFilter::new(&attempt, scope)); + } + } } - port.record_attempt_failed(&attempt).await?; - // Exhaustion diagnostics are only needed after an attempt fails. Keep // the common successful path free of a deep plan/report-context clone. last_attempted = Some((attempt.execution_plan().clone(), attempt.report_context())); } + if let Some(response) = fallback_response { + return Ok(AiAttemptLoopOutcome::Deferred(response)); + } + let Some((last_plan, last_report_context)) = last_attempted else { return Ok(AiAttemptLoopOutcome::NoPath); }; @@ -103,6 +158,36 @@ where )) } +#[derive(Debug)] +struct AiAttemptRetryFilter { + scope: AiAttemptRetryScope, + provider_id: String, + endpoint_id: String, + key_id: String, +} + +impl AiAttemptRetryFilter { + fn new(attempt: &Attempt, scope: AiAttemptRetryScope) -> Self { + let plan = attempt.execution_plan(); + Self { + scope, + provider_id: plan.provider_id.clone(), + endpoint_id: plan.endpoint_id.clone(), + key_id: plan.key_id.clone(), + } + } + + fn matches(&self, attempt: &Attempt) -> bool { + let plan = attempt.execution_plan(); + match self.scope { + AiAttemptRetryScope::Candidate => false, + AiAttemptRetryScope::Credential => plan.key_id == self.key_id, + AiAttemptRetryScope::Endpoint => plan.endpoint_id == self.endpoint_id, + AiAttemptRetryScope::Provider => plan.provider_id == self.provider_id, + } + } +} + impl AiExecutionAttempt for crate::dto::AiSyncAttempt { fn execution_plan(&self) -> &aether_contracts::ExecutionPlan { &self.plan @@ -146,7 +231,10 @@ mod tests { use async_trait::async_trait; - use super::{run_ai_attempt_loop, AiAttemptLoopPort, AiExecutionAttempt}; + use super::{ + run_ai_attempt_loop, AiAttemptExecutionOutcome, AiAttemptLoopPort, AiAttemptRetryScope, + AiExecutionAttempt, + }; #[derive(Clone)] struct TestAttempt { @@ -173,6 +261,60 @@ mod tests { unused: Mutex>, } + struct ScopedRetryPort { + executed: Mutex>, + unused: Mutex>, + } + + #[async_trait] + impl AiAttemptLoopPort for ScopedRetryPort { + type Response = &'static str; + type Exhaustion = (); + type Error = &'static str; + + async fn execute_attempt( + &self, + attempt: &TestAttempt, + ) -> Result, Self::Error> { + self.executed + .lock() + .expect("executed attempts should lock") + .push(attempt.id); + Ok(match attempt.id { + "endpoint-failure" => { + AiAttemptExecutionOutcome::retry(AiAttemptRetryScope::Endpoint) + } + "credential-failure" => { + AiAttemptExecutionOutcome::retry(AiAttemptRetryScope::Credential) + } + "provider-failure" => AiAttemptExecutionOutcome::Retry { + scope: AiAttemptRetryScope::Provider, + fallback_response: Some("provider-error"), + }, + _ => AiAttemptExecutionOutcome::Responded(attempt.id), + }) + } + + async fn mark_unused_attempts( + &self, + attempts: Vec, + ) -> Result<(), Self::Error> { + self.unused + .lock() + .expect("unused attempts should lock") + .extend(attempts.into_iter().map(|attempt| attempt.id)); + Ok(()) + } + + async fn build_exhaustion( + &self, + _last_plan: aether_contracts::ExecutionPlan, + _last_report_context: Option, + ) -> Result { + Ok(()) + } + } + #[async_trait] impl AiAttemptLoopPort for FailingPort { type Response = (); @@ -182,11 +324,13 @@ mod tests { async fn execute_attempt( &self, attempt: &TestAttempt, - ) -> Result, Self::Error> { + ) -> Result, Self::Error> { if attempt.id == self.fail_on { Err("attempt failed") } else { - Ok(None) + Ok(AiAttemptExecutionOutcome::retry( + AiAttemptRetryScope::Candidate, + )) } } @@ -237,6 +381,19 @@ mod tests { } } + fn routed_attempt( + id: &'static str, + provider_id: &str, + endpoint_id: &str, + key_id: &str, + ) -> TestAttempt { + let mut attempt = attempt(id); + attempt.plan.provider_id = provider_id.to_string(); + attempt.plan.endpoint_id = endpoint_id.to_string(); + attempt.plan.key_id = key_id.to_string(); + attempt + } + #[tokio::test] async fn marks_unattempted_candidates_unused_when_execution_returns_error() { let port = FailingPort { @@ -261,4 +418,69 @@ mod tests { vec!["candidate-3"] ); } + + #[tokio::test] + async fn retry_scopes_skip_matching_static_candidates() { + let port = ScopedRetryPort { + executed: Mutex::new(Vec::new()), + unused: Mutex::new(Vec::new()), + }; + let attempts = vec![ + routed_attempt("endpoint-failure", "provider-a", "endpoint-a", "key-a"), + routed_attempt("same-endpoint", "provider-a", "endpoint-a", "key-b"), + routed_attempt("credential-failure", "provider-a", "endpoint-b", "key-c"), + routed_attempt("same-credential", "provider-a", "endpoint-c", "key-c"), + routed_attempt("provider-failure", "provider-b", "endpoint-d", "key-d"), + routed_attempt("same-provider", "provider-b", "endpoint-e", "key-e"), + routed_attempt("success", "provider-c", "endpoint-f", "key-f"), + ]; + + let outcome = run_ai_attempt_loop(&port, attempts) + .await + .expect("scoped retry loop should succeed"); + + assert!(matches!( + outcome, + super::AiAttemptLoopOutcome::Responded("success") + )); + assert_eq!( + *port.executed.lock().expect("executed attempts should lock"), + vec![ + "endpoint-failure", + "credential-failure", + "provider-failure", + "success" + ] + ); + assert_eq!( + *port.unused.lock().expect("unused attempts should lock"), + vec!["same-endpoint", "same-credential", "same-provider"] + ); + } + + #[tokio::test] + async fn returns_preserved_upstream_response_after_candidates_exhaust() { + let port = ScopedRetryPort { + executed: Mutex::new(Vec::new()), + unused: Mutex::new(Vec::new()), + }; + let outcome = run_ai_attempt_loop( + &port, + vec![ + routed_attempt("provider-failure", "provider-a", "endpoint-a", "key-a"), + routed_attempt("same-provider", "provider-a", "endpoint-b", "key-b"), + ], + ) + .await + .expect("fallback response loop should succeed"); + + assert!(matches!( + outcome, + super::AiAttemptLoopOutcome::Deferred("provider-error") + )); + assert_eq!( + *port.unused.lock().expect("unused attempts should lock"), + vec!["same-provider"] + ); + } } diff --git a/crates/aether-ai/serving/src/dto.rs b/crates/aether-ai/serving/src/dto.rs index 686132cef..8fc0274ee 100644 --- a/crates/aether-ai/serving/src/dto.rs +++ b/crates/aether-ai/serving/src/dto.rs @@ -44,6 +44,28 @@ impl ConversionMode { } } +/// Request/response adaptation applied after candidate selection. +/// +/// This is independent from format conversion: a same-format request may be +/// byte-transparent or may intentionally apply provider compatibility edits. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum AdaptationMode { + NativeTransparent, + SameFormatCompat, + CrossFormat, +} + +impl AdaptationMode { + pub const fn as_str(self) -> &'static str { + match self { + Self::NativeTransparent => "native_transparent", + Self::SameFormatCompat => "same_format_compat", + Self::CrossFormat => "cross_format", + } + } +} + #[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)] pub struct AiRequestGzipPolicy { #[serde(default, skip_serializing_if = "Option::is_none")] diff --git a/crates/aether-ai/serving/src/execution_path.rs b/crates/aether-ai/serving/src/execution_path.rs index a4cf2fcc8..ba73415d0 100644 --- a/crates/aether-ai/serving/src/execution_path.rs +++ b/crates/aether-ai/serving/src/execution_path.rs @@ -3,6 +3,7 @@ use async_trait::async_trait; #[derive(Debug)] pub enum AiServingExecutionOutcome { Responded(Response), + Deferred(Response), Exhausted(Exhaustion), NoPath, } @@ -98,9 +99,15 @@ where Port: AiSyncExecutionPathPort, { let mut exhausted = None; + let mut deferred = None; - if let Some(response) = - absorb_sync_step(port, AiSyncExecutionStep::VideoTaskFollowUp, &mut exhausted).await? + if let Some(response) = absorb_sync_step( + port, + AiSyncExecutionStep::VideoTaskFollowUp, + &mut deferred, + &mut exhausted, + ) + .await? { return Ok(response); } @@ -116,12 +123,18 @@ where AiSyncExecutionStep::LocalGeminiFiles, AiSyncExecutionStep::RemoteDecision, ] { - if let Some(response) = absorb_sync_step(port, step, &mut exhausted).await? { + if let Some(response) = + absorb_sync_step(port, step, &mut deferred, &mut exhausted).await? + { return Ok(response); } } } + if let Some(response) = deferred { + return Ok(AiServingExecutionOutcome::Deferred(response)); + } + if let Some(outcome) = exhausted { return Ok(AiServingExecutionOutcome::Exhausted(outcome)); } @@ -135,6 +148,9 @@ where AiServingExecutionOutcome::Responded(response) => { Ok(AiServingExecutionOutcome::Responded(response)) } + AiServingExecutionOutcome::Deferred(response) => { + Ok(AiServingExecutionOutcome::Deferred(response)) + } AiServingExecutionOutcome::Exhausted(outcome) => { Ok(AiServingExecutionOutcome::Exhausted(outcome)) } @@ -149,17 +165,24 @@ where Port: AiStreamExecutionPathPort, { let mut exhausted = None; + let mut deferred = None; for step in port.stream_execution_steps() { if *step != AiStreamExecutionStep::LocalVideoContent && !port.scheduler_decision_supported() { continue; } - if let Some(response) = absorb_stream_step(port, *step, &mut exhausted).await? { + if let Some(response) = + absorb_stream_step(port, *step, &mut deferred, &mut exhausted).await? + { return Ok(response); } } + if let Some(response) = deferred { + return Ok(AiServingExecutionOutcome::Deferred(response)); + } + if let Some(outcome) = exhausted { return Ok(AiServingExecutionOutcome::Exhausted(outcome)); } @@ -173,6 +196,9 @@ where AiServingExecutionOutcome::Responded(response) => { Ok(AiServingExecutionOutcome::Responded(response)) } + AiServingExecutionOutcome::Deferred(response) => { + Ok(AiServingExecutionOutcome::Deferred(response)) + } AiServingExecutionOutcome::Exhausted(outcome) => { Ok(AiServingExecutionOutcome::Exhausted(outcome)) } @@ -183,6 +209,7 @@ where async fn absorb_sync_step( port: &Port, step: AiSyncExecutionStep, + deferred: &mut Option, exhausted: &mut Option, ) -> Result>, Port::Error> where @@ -192,6 +219,10 @@ where AiServingExecutionOutcome::Responded(response) => { Ok(Some(AiServingExecutionOutcome::Responded(response))) } + AiServingExecutionOutcome::Deferred(response) => { + *deferred = Some(response); + Ok(None) + } AiServingExecutionOutcome::Exhausted(outcome) => { *exhausted = Some(outcome); Ok(None) @@ -203,6 +234,7 @@ where async fn absorb_stream_step( port: &Port, step: AiStreamExecutionStep, + deferred: &mut Option, exhausted: &mut Option, ) -> Result>, Port::Error> where @@ -212,6 +244,10 @@ where AiServingExecutionOutcome::Responded(response) => { Ok(Some(AiServingExecutionOutcome::Responded(response))) } + AiServingExecutionOutcome::Deferred(response) => { + *deferred = Some(response); + Ok(None) + } AiServingExecutionOutcome::Exhausted(outcome) => { *exhausted = Some(outcome); Ok(None) @@ -390,6 +426,30 @@ mod tests { ); } + #[tokio::test] + async fn sync_path_keeps_deferred_error_until_a_later_step_succeeds() { + let port = TestSyncPort { + scheduler_supported: true, + outcomes: Mutex::new(VecDeque::from([ + AiServingExecutionOutcome::NoPath, + AiServingExecutionOutcome::Deferred("preserved_upstream_error"), + AiServingExecutionOutcome::Responded("later_success"), + ])), + calls: Mutex::default(), + }; + + let outcome = run_ai_sync_execution_path(&port).await.unwrap(); + + assert!(matches!( + outcome, + AiServingExecutionOutcome::Responded("later_success") + )); + assert_eq!( + port.calls.lock().unwrap().as_slice(), + ["VideoTaskFollowUp", "LocalVideo", "LocalImage"] + ); + } + #[tokio::test] async fn stream_path_skips_scheduler_steps_when_unsupported() { let port = TestStreamPort { @@ -451,6 +511,34 @@ mod tests { assert_eq!(port.calls.lock().unwrap().as_slice(), ["LocalOpenAiChat"]); } + #[tokio::test] + async fn stream_path_returns_deferred_error_only_after_steps_exhaust() { + const TWO_STEPS: &[AiStreamExecutionStep] = &[ + AiStreamExecutionStep::LocalOpenAiChat, + AiStreamExecutionStep::LocalSameFormatProvider, + ]; + let port = TestStreamPort { + scheduler_supported: true, + stream_steps: Some(TWO_STEPS), + outcomes: Mutex::new(VecDeque::from([ + AiServingExecutionOutcome::Deferred("preserved_upstream_error"), + AiServingExecutionOutcome::NoPath, + ])), + calls: Mutex::default(), + }; + + let outcome = run_ai_stream_execution_path(&port).await.unwrap(); + + assert!(matches!( + outcome, + AiServingExecutionOutcome::Deferred("preserved_upstream_error") + )); + assert_eq!( + port.calls.lock().unwrap().as_slice(), + ["LocalOpenAiChat", "LocalSameFormatProvider"] + ); + } + #[tokio::test] async fn stream_path_returns_last_exhaustion_without_plan_fallback() { let port = TestStreamPort { diff --git a/crates/aether-ai/serving/src/lib.rs b/crates/aether-ai/serving/src/lib.rs index 7e31e87c5..5c44bcaf3 100644 --- a/crates/aether-ai/serving/src/lib.rs +++ b/crates/aether-ai/serving/src/lib.rs @@ -14,6 +14,7 @@ pub mod decision_payload; pub mod dto; pub mod execution_path; pub mod failure_diagnostic; +pub mod payload_fidelity; pub mod plan_payload; pub mod ports; pub mod ranking_metadata; @@ -52,7 +53,8 @@ pub use aether_pool_core::{ REQUEST_FAILURE_PENALTY, UNSCHEDULABLE_SCORE_CAP, }; pub use attempt_loop::{ - run_ai_attempt_loop, AiAttemptLoopOutcome, AiAttemptLoopPort, AiExecutionAttempt, + run_ai_attempt_loop, AiAttemptExecutionOutcome, AiAttemptLoopOutcome, AiAttemptLoopPort, + AiAttemptRetryScope, AiExecutionAttempt, }; pub use attempt_plan::{ build_ai_execution_decision_from_plan, build_ai_execution_plan_from_decision, @@ -106,7 +108,7 @@ pub use decision_payload::{ AiExecutionDecisionResponseParts, }; pub use dto::{ - augment_sync_report_context, generic_decision_missing_exact_provider_request, + augment_sync_report_context, generic_decision_missing_exact_provider_request, AdaptationMode, AiExecutionDecision, AiExecutionPlanPayload, AiRequestGzipPolicy, AiStreamAttempt, AiSyncAttempt, ConversionMode, ExecutionStrategy, }; @@ -116,6 +118,7 @@ pub use execution_path::{ AiSyncExecutionPathPort, AiSyncExecutionStep, DEFAULT_STREAM_EXECUTION_STEPS, }; pub use failure_diagnostic::{CandidateFailureDiagnostic, CandidateFailureDiagnosticKind}; +pub use payload_fidelity::OriginalRequestPayload; pub use plan_payload::{ build_ai_stream_execution_plan_payload, build_ai_sync_execution_plan_payload, }; diff --git a/crates/aether-ai/serving/src/payload_fidelity.rs b/crates/aether-ai/serving/src/payload_fidelity.rs new file mode 100644 index 000000000..71c12ef26 --- /dev/null +++ b/crates/aether-ai/serving/src/payload_fidelity.rs @@ -0,0 +1,77 @@ +use std::sync::Arc; + +use base64::Engine as _; + +/// The frontdoor-normalized JSON request as both a parsed value and its exact +/// decoded bytes. This is carried only inside the local process through HTTP +/// request extensions; serialized execution contracts continue to use +/// `RequestBody::body_bytes_b64`. +#[derive(Debug, Clone)] +pub struct OriginalRequestPayload { + body_json: Arc, + body_bytes: Arc<[u8]>, +} + +impl OriginalRequestPayload { + pub fn from_parsed_json(body_json: serde_json::Value, body_bytes: &[u8]) -> Self { + Self { + body_json: Arc::new(body_json), + body_bytes: Arc::from(body_bytes), + } + } + + /// Returns the original body only when the terminal provider JSON is + /// semantically unchanged. Object key order and whitespace are preserved by + /// returning the captured bytes rather than serializing `provider_body`. + pub fn body_bytes_base64_if_unchanged( + &self, + provider_body: &serde_json::Value, + ) -> Option { + if self.body_bytes.is_empty() || provider_body != self.body_json.as_ref() { + return None; + } + + Some(base64::engine::general_purpose::STANDARD.encode(self.body_bytes.as_ref())) + } +} + +#[cfg(test)] +mod tests { + use base64::Engine as _; + use serde_json::json; + + use super::OriginalRequestPayload; + + #[test] + fn preserves_exact_json_bytes_when_terminal_value_is_unchanged() { + let raw = br#"{ "unknown": true, "model": "claude-sonnet-4" }"#; + let parsed: serde_json::Value = serde_json::from_slice(raw).expect("request should parse"); + let payload = OriginalRequestPayload::from_parsed_json(parsed.clone(), raw); + + let encoded = payload + .body_bytes_base64_if_unchanged(&parsed) + .expect("unchanged body should preserve bytes"); + + assert_eq!( + base64::engine::general_purpose::STANDARD + .decode(encoded) + .expect("body should decode"), + raw + ); + } + + #[test] + fn rejects_original_bytes_when_terminal_value_changed() { + let raw = br#"{"model":"claude-sonnet-4","messages":[]}"#; + let parsed: serde_json::Value = serde_json::from_slice(raw).expect("request should parse"); + let payload = OriginalRequestPayload::from_parsed_json(parsed, raw); + + assert_eq!( + payload.body_bytes_base64_if_unchanged(&json!({ + "model": "claude-sonnet-4-5", + "messages": [] + })), + None + ); + } +} diff --git a/crates/aether-contracts/src/lib.rs b/crates/aether-contracts/src/lib.rs index c450eac37..8dd71ebbe 100644 --- a/crates/aether-contracts/src/lib.rs +++ b/crates/aether-contracts/src/lib.rs @@ -9,9 +9,10 @@ mod usage; pub use error::{ExecutionError, ExecutionErrorKind, ExecutionPhase}; pub use frame::{StreamFrame, StreamFramePayload, StreamFrameType}; pub use plan::{ - ExecutionPlan, ExecutionTimeouts, ProxySnapshot, RequestBody, ResolvedTransportProfile, - EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER, EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER, - EXECUTION_REQUEST_HTTP1_ONLY_HEADER, MAX_EXECUTION_REQUEST_TIMEOUT_MS, + ExecutionPlan, ExecutionResponseBodyMode, ExecutionTimeouts, ProxySnapshot, RequestBody, + ResolvedTransportProfile, EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER, + EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER, EXECUTION_REQUEST_HTTP1_ONLY_HEADER, + EXECUTION_RESPONSE_BODY_MODE_HEADER, MAX_EXECUTION_REQUEST_TIMEOUT_MS, MAX_EXECUTION_REQUEST_TIMEOUT_SECS, MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_MS, MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_SECS, TRANSPORT_BACKEND_BROWSER_WREQ, TRANSPORT_BACKEND_HYPER_RUSTLS, TRANSPORT_BACKEND_REQWEST_RUSTLS, TRANSPORT_HTTP_MODE_AUTO, diff --git a/crates/aether-contracts/src/plan.rs b/crates/aether-contracts/src/plan.rs index 3609fb561..443256521 100644 --- a/crates/aether-contracts/src/plan.rs +++ b/crates/aether-contracts/src/plan.rs @@ -7,12 +7,39 @@ pub const EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER: &str = "x-aether-execution- pub const EXECUTION_REQUEST_HTTP1_ONLY_HEADER: &str = "x-aether-execution-http1-only"; pub const EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER: &str = "x-aether-execution-accept-invalid-certs"; +pub const EXECUTION_RESPONSE_BODY_MODE_HEADER: &str = "x-aether-execution-response-body-mode"; pub const MAX_EXECUTION_REQUEST_TIMEOUT_SECS: u64 = 1_200; pub const MAX_EXECUTION_REQUEST_TIMEOUT_MS: u64 = MAX_EXECUTION_REQUEST_TIMEOUT_SECS * 1_000; pub const MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_SECS: u64 = 300; pub const MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_MS: u64 = MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_SECS * 1_000; +#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum ExecutionResponseBodyMode { + #[default] + StructuredJson, + PreserveBytes, +} + +impl ExecutionResponseBodyMode { + pub const fn as_str(self) -> &'static str { + match self { + Self::StructuredJson => "structured_json", + Self::PreserveBytes => "preserve_bytes", + } + } + + pub fn from_header_value(value: Option<&str>) -> Self { + match value.map(str::trim) { + Some(value) if value.eq_ignore_ascii_case(Self::PreserveBytes.as_str()) => { + Self::PreserveBytes + } + _ => Self::StructuredJson, + } + } +} + #[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] #[serde(default)] pub struct ExecutionTimeouts { @@ -139,6 +166,22 @@ pub struct ExecutionPlan { mod tests { use super::*; + #[test] + fn response_body_mode_is_fail_closed_to_structured_json() { + assert_eq!( + ExecutionResponseBodyMode::from_header_value(Some(" preserve_bytes ")), + ExecutionResponseBodyMode::PreserveBytes + ); + assert_eq!( + ExecutionResponseBodyMode::from_header_value(Some("unexpected")), + ExecutionResponseBodyMode::StructuredJson + ); + assert_eq!( + ExecutionResponseBodyMode::from_header_value(None), + ExecutionResponseBodyMode::StructuredJson + ); + } + #[test] fn serializes_plan_with_json_body() { let plan = ExecutionPlan { diff --git a/crates/aether-data/adapters/mysql/src/candidate_selection.rs b/crates/aether-data/adapters/mysql/src/candidate_selection.rs index 97248ac5c..d3d6879f4 100644 --- a/crates/aether-data/adapters/mysql/src/candidate_selection.rs +++ b/crates/aether-data/adapters/mysql/src/candidate_selection.rs @@ -460,22 +460,16 @@ fn key_auth_channel_matches(row: &CandidateSelectionRow, api_format: &str) -> bo matches!(auth_type.as_str(), "oauth" | "api_key" | "bearer") && api_format == "openai:chat" } - "vertex_ai" => { - (auth_type == "api_key" - && matches!( - api_format.as_str(), - "gemini:generate_content" | "gemini:embedding" - )) - || (matches!(auth_type.as_str(), "service_account" | "vertex_ai") - && matches!( - api_format.as_str(), - "claude:messages" | "gemini:generate_content" | "gemini:embedding" - )) - } + "vertex_ai" => vertex_key_auth_channel_matches(&auth_type, &api_format), _ => auth_type != "oauth", } } +fn vertex_key_auth_channel_matches(auth_type: &str, api_format: &str) -> bool { + matches!(auth_type, "api_key" | "service_account" | "vertex_ai") + && matches!(api_format, "gemini:generate_content" | "gemini:embedding") +} + fn dedupe_candidate_selection_rows( rows: Vec, ) -> Vec { @@ -797,7 +791,7 @@ fn sql_match_aliases(api_formats: &[String]) -> Vec { #[cfg(test)] mod tests { - use super::MysqlMinimalCandidateSelectionReadRepository; + use super::{vertex_key_auth_channel_matches, MysqlMinimalCandidateSelectionReadRepository}; #[tokio::test] async fn repository_builds_from_lazy_pool() { @@ -809,4 +803,22 @@ mod tests { let _repository = MysqlMinimalCandidateSelectionReadRepository::new(pool); } + + #[test] + fn vertex_auth_matrix_rejects_retired_claude_format() { + for auth_type in ["api_key", "service_account", "vertex_ai"] { + assert!(!vertex_key_auth_channel_matches( + auth_type, + "claude:messages" + )); + assert!(vertex_key_auth_channel_matches( + auth_type, + "gemini:generate_content" + )); + assert!(vertex_key_auth_channel_matches( + auth_type, + "gemini:embedding" + )); + } + } } diff --git a/crates/aether-data/adapters/mysql/src/provider_catalog.rs b/crates/aether-data/adapters/mysql/src/provider_catalog.rs index 1efa21335..38089af08 100644 --- a/crates/aether-data/adapters/mysql/src/provider_catalog.rs +++ b/crates/aether-data/adapters/mysql/src/provider_catalog.rs @@ -1281,6 +1281,19 @@ WHERE id = ? "provider catalog OAuth api_key update must not be empty".to_string(), )); } + if update.expected_credential.as_ref().is_some_and(|expected| { + expected + .encrypted_api_key + .as_deref() + .is_some_and(|value| value.trim().is_empty()) + || expected.auth_type.trim().is_empty() + || expected.provider_id.trim().is_empty() + || expected.provider_type.trim().is_empty() + }) { + return Err(DataLayerError::InvalidInput( + "provider catalog OAuth credential fence must not contain empty fields".to_string(), + )); + } if !update.status_snapshot_patch.is_object() { return Err(DataLayerError::InvalidInput( "provider catalog status snapshot patch must be an object".to_string(), @@ -1335,8 +1348,24 @@ WHERE id = ? ) .push(" WHERE id = ") .push_bind(&update.key_id) - .push(" AND auth_config <=> ") + .push(" AND BINARY auth_config <=> BINARY ") .push_bind(update.expected_encrypted_auth_config.as_deref()); + if let Some(expected) = update.expected_credential.as_ref() { + builder + .push(" AND BINARY api_key <=> BINARY ") + .push_bind(expected.encrypted_api_key.as_deref()) + .push(" AND BINARY auth_type = BINARY ") + .push_bind(&expected.auth_type) + .push(" AND BINARY provider_id = BINARY ") + .push_bind(&expected.provider_id) + .push( + " AND EXISTS (SELECT 1 FROM providers WHERE \ + BINARY providers.id = BINARY provider_api_keys.provider_id \ + AND BINARY providers.provider_type = BINARY ", + ) + .push_bind(&expected.provider_type) + .push(")"); + } let rows_affected = builder .build() .execute(&self.pool) diff --git a/crates/aether-data/adapters/postgres/src/candidate_selection.rs b/crates/aether-data/adapters/postgres/src/candidate_selection.rs index e27284e72..b34179137 100644 --- a/crates/aether-data/adapters/postgres/src/candidate_selection.rs +++ b/crates/aether-data/adapters/postgres/src/candidate_selection.rs @@ -114,16 +114,8 @@ INNER JOIN LATERAL ( ) OR ( LOWER(BTRIM(p.provider_type)) = 'vertex_ai' - AND ( - ( - LOWER(BTRIM(pak.auth_type)) = 'api_key' - AND LOWER($3) IN ('gemini:generate_content', 'gemini:embedding') - ) - OR ( - LOWER(BTRIM(pak.auth_type)) IN ('service_account', 'vertex_ai') - AND LOWER($3) IN ('claude:messages', 'gemini:generate_content', 'gemini:embedding') - ) - ) + AND LOWER(BTRIM(pak.auth_type)) IN ('api_key', 'service_account', 'vertex_ai') + AND LOWER($3) IN ('gemini:generate_content', 'gemini:embedding') ) OR ( LOWER(BTRIM(p.provider_type)) NOT IN ( @@ -207,16 +199,8 @@ WHERE p.is_active = TRUE ) OR ( LOWER(BTRIM(p.provider_type)) = 'vertex_ai' - AND ( - ( - LOWER(BTRIM(pak.auth_type)) = 'api_key' - AND LOWER($3) IN ('gemini:generate_content', 'gemini:embedding') - ) - OR ( - LOWER(BTRIM(pak.auth_type)) IN ('service_account', 'vertex_ai') - AND LOWER($3) IN ('claude:messages', 'gemini:generate_content', 'gemini:embedding') - ) - ) + AND LOWER(BTRIM(pak.auth_type)) IN ('api_key', 'service_account', 'vertex_ai') + AND LOWER($3) IN ('gemini:generate_content', 'gemini:embedding') ) OR ( LOWER(BTRIM(p.provider_type)) NOT IN ( @@ -393,16 +377,8 @@ INNER JOIN LATERAL ( ) OR ( LOWER(BTRIM(p.provider_type)) = 'vertex_ai' - AND ( - ( - LOWER(BTRIM(pak.auth_type)) = 'api_key' - AND LOWER($4) IN ('gemini:generate_content', 'gemini:embedding') - ) - OR ( - LOWER(BTRIM(pak.auth_type)) IN ('service_account', 'vertex_ai') - AND LOWER($4) IN ('claude:messages', 'gemini:generate_content', 'gemini:embedding') - ) - ) + AND LOWER(BTRIM(pak.auth_type)) IN ('api_key', 'service_account', 'vertex_ai') + AND LOWER($4) IN ('gemini:generate_content', 'gemini:embedding') ) OR ( LOWER(BTRIM(p.provider_type)) NOT IN ( @@ -487,16 +463,8 @@ WHERE p.is_active = TRUE ) OR ( LOWER(BTRIM(p.provider_type)) = 'vertex_ai' - AND ( - ( - LOWER(BTRIM(pak.auth_type)) = 'api_key' - AND LOWER($4) IN ('gemini:generate_content', 'gemini:embedding') - ) - OR ( - LOWER(BTRIM(pak.auth_type)) IN ('service_account', 'vertex_ai') - AND LOWER($4) IN ('claude:messages', 'gemini:generate_content', 'gemini:embedding') - ) - ) + AND LOWER(BTRIM(pak.auth_type)) IN ('api_key', 'service_account', 'vertex_ai') + AND LOWER($4) IN ('gemini:generate_content', 'gemini:embedding') ) OR ( LOWER(BTRIM(p.provider_type)) NOT IN ( @@ -681,16 +649,8 @@ WHERE p.is_active = TRUE ) OR ( LOWER(BTRIM(p.provider_type)) = 'vertex_ai' - AND ( - ( - LOWER(BTRIM(pak.auth_type)) = 'api_key' - AND LOWER($6) IN ('gemini:generate_content', 'gemini:embedding') - ) - OR ( - LOWER(BTRIM(pak.auth_type)) IN ('service_account', 'vertex_ai') - AND LOWER($6) IN ('claude:messages', 'gemini:generate_content', 'gemini:embedding') - ) - ) + AND LOWER(BTRIM(pak.auth_type)) IN ('api_key', 'service_account', 'vertex_ai') + AND LOWER($6) IN ('gemini:generate_content', 'gemini:embedding') ) OR ( LOWER(BTRIM(p.provider_type)) NOT IN ( @@ -1579,18 +1539,39 @@ mod tests { } #[test] - fn candidate_selection_sql_allows_vertex_embedding_auth() { + fn candidate_selection_sql_rejects_retired_vertex_claude_auth() { let requested_model_sql = requested_model_selection_sql(); - for sql in [ - LIST_FOR_EXACT_API_FORMAT_SQL, - LIST_FOR_EXACT_API_FORMAT_AND_GLOBAL_MODEL_SQL, - LIST_POOL_KEYS_FOR_GROUP_SQL, - requested_model_sql.as_str(), + for (name, sql, expected_occurrences) in [ + ("exact", LIST_FOR_EXACT_API_FORMAT_SQL, 2), + ( + "global_model", + LIST_FOR_EXACT_API_FORMAT_AND_GLOBAL_MODEL_SQL, + 2, + ), + ("pool_keys", LIST_POOL_KEYS_FOR_GROUP_SQL, 1), + ("requested_model", requested_model_sql.as_str(), 2), ] { - assert!(sql.contains("LOWER(BTRIM(p.provider_type)) = 'vertex_ai'")); - assert!(sql.contains("gemini:embedding")); - assert!(sql.contains("gemini:generate_content")); - assert!(sql.contains("claude:messages")); + let mut remaining = sql; + let mut occurrences = 0; + while let Some((_, suffix)) = + remaining.split_once("LOWER(BTRIM(p.provider_type)) = 'vertex_ai'") + { + let (vertex_clause, rest) = suffix + .split_once("LOWER(BTRIM(p.provider_type)) NOT IN") + .expect("each Vertex auth clause should have a following fallback clause"); + assert!(vertex_clause.contains( + "LOWER(BTRIM(pak.auth_type)) IN ('api_key', 'service_account', 'vertex_ai')" + )); + assert!(vertex_clause.contains("gemini:embedding")); + assert!(vertex_clause.contains("gemini:generate_content")); + assert!(!vertex_clause.contains("claude:messages")); + occurrences += 1; + remaining = rest; + } + assert_eq!( + occurrences, expected_occurrences, + "unexpected Vertex auth clause count in {name} SQL" + ); } } diff --git a/crates/aether-data/adapters/postgres/src/provider_catalog.rs b/crates/aether-data/adapters/postgres/src/provider_catalog.rs index 52172160e..e52284d5d 100644 --- a/crates/aether-data/adapters/postgres/src/provider_catalog.rs +++ b/crates/aether-data/adapters/postgres/src/provider_catalog.rs @@ -892,6 +892,15 @@ WHERE id = $1 .encrypted_api_key_update .as_deref() .is_some_and(|value| value.trim().is_empty()) + || update.expected_credential.as_ref().is_some_and(|expected| { + expected + .encrypted_api_key + .as_deref() + .is_some_and(|value| value.trim().is_empty()) + || expected.auth_type.trim().is_empty() + || expected.provider_id.trim().is_empty() + || expected.provider_type.trim().is_empty() + }) || !update.status_snapshot_patch.is_object() || update .upstream_metadata_patch @@ -934,6 +943,18 @@ SET END WHERE id = $1 AND auth_config IS NOT DISTINCT FROM $12 + AND ($13::boolean IS FALSE OR api_key IS NOT DISTINCT FROM $14) + AND ($15::text IS NULL OR auth_type = $15) + AND ($16::text IS NULL OR provider_id = $16) + AND ( + $17::text IS NULL + OR EXISTS ( + SELECT 1 + FROM providers + WHERE providers.id = provider_api_keys.provider_id + AND providers.provider_type = $17 + ) + ) "#, ) .bind(&update.key_id) @@ -953,6 +974,31 @@ WHERE id = $1 .bind(update.reset_error_count) .bind(update.updated_at_unix_secs.map(|value| value as f64)) .bind(update.expected_encrypted_auth_config.as_deref()) + .bind(update.expected_credential.is_some()) + .bind( + update + .expected_credential + .as_ref() + .and_then(|expected| expected.encrypted_api_key.as_deref()), + ) + .bind( + update + .expected_credential + .as_ref() + .map(|expected| expected.auth_type.as_str()), + ) + .bind( + update + .expected_credential + .as_ref() + .map(|expected| expected.provider_id.as_str()), + ) + .bind( + update + .expected_credential + .as_ref() + .map(|expected| expected.provider_type.as_str()), + ) .execute(&self.pool) .await .map_postgres_err()? diff --git a/crates/aether-data/adapters/sqlite/src/candidate_selection.rs b/crates/aether-data/adapters/sqlite/src/candidate_selection.rs index ad5e749ac..d09d78197 100644 --- a/crates/aether-data/adapters/sqlite/src/candidate_selection.rs +++ b/crates/aether-data/adapters/sqlite/src/candidate_selection.rs @@ -562,24 +562,12 @@ fn push_key_auth_channel_sql_filter( ) OR ( LOWER(TRIM(p.provider_type)) = 'vertex_ai' - AND ( - ( - LOWER(TRIM(pak.auth_type)) = 'api_key' - AND "#, + AND LOWER(TRIM(pak.auth_type)) IN ('api_key', 'service_account', 'vertex_ai') + AND "#, ); builder.push_bind(api_format.clone()); builder.push( - r#" = 'gemini:generate_content' - ) - OR ( - LOWER(TRIM(pak.auth_type)) IN ('service_account', 'vertex_ai') - AND "#, - ); - builder.push_bind(api_format.clone()); - builder.push( - r#" IN ('claude:messages', 'gemini:generate_content') - ) - ) + r#" IN ('gemini:generate_content', 'gemini:embedding') ) OR ( LOWER(TRIM(p.provider_type)) NOT IN ( @@ -862,22 +850,16 @@ fn key_auth_channel_matches(row: &CandidateSelectionRow, api_format: &str) -> bo matches!(auth_type.as_str(), "oauth" | "api_key" | "bearer") && api_format == "openai:chat" } - "vertex_ai" => { - (auth_type == "api_key" - && matches!( - api_format.as_str(), - "gemini:generate_content" | "gemini:embedding" - )) - || (matches!(auth_type.as_str(), "service_account" | "vertex_ai") - && matches!( - api_format.as_str(), - "claude:messages" | "gemini:generate_content" | "gemini:embedding" - )) - } + "vertex_ai" => vertex_key_auth_channel_matches(&auth_type, &api_format), _ => auth_type != "oauth", } } +fn vertex_key_auth_channel_matches(auth_type: &str, api_format: &str) -> bool { + matches!(auth_type, "api_key" | "service_account" | "vertex_ai") + && matches!(api_format, "gemini:generate_content" | "gemini:embedding") +} + fn dedupe_candidate_selection_rows( rows: Vec, ) -> Vec { @@ -1194,13 +1176,46 @@ fn sql_match_aliases(api_formats: &[String]) -> Vec { #[cfg(test)] mod tests { - use super::SqliteMinimalCandidateSelectionReadRepository; + use super::{ + push_key_auth_channel_sql_filter, vertex_key_auth_channel_matches, + SqliteMinimalCandidateSelectionReadRepository, + }; use crate::run_migrations; use aether_data_contracts::repository::candidate_selection::{ MinimalCandidateSelectionReadRepository, StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsQuery, StoredRequestedModelCandidateRowsQuery, }; + #[test] + fn vertex_auth_matrix_rejects_retired_claude_format() { + for auth_type in ["api_key", "service_account", "vertex_ai"] { + assert!(!vertex_key_auth_channel_matches( + auth_type, + "claude:messages" + )); + assert!(vertex_key_auth_channel_matches( + auth_type, + "gemini:generate_content" + )); + assert!(vertex_key_auth_channel_matches( + auth_type, + "gemini:embedding" + )); + } + + let mut builder = sqlx::QueryBuilder::::new("SELECT 1 WHERE 1 = 1"); + push_key_auth_channel_sql_filter(&mut builder, "claude:messages"); + let sql = builder.sql(); + let vertex_clause = sql + .split_once("LOWER(TRIM(p.provider_type)) = 'vertex_ai'") + .and_then(|(_, suffix)| suffix.split_once("LOWER(TRIM(p.provider_type)) NOT IN")) + .map(|(clause, _)| clause) + .expect("Vertex auth clause should exist"); + assert!(!vertex_clause.contains("claude:messages")); + assert!(vertex_clause.contains("gemini:generate_content")); + assert!(vertex_clause.contains("gemini:embedding")); + } + #[tokio::test] async fn sqlite_repository_reads_candidate_selection_rows() { let pool = sqlx::sqlite::SqlitePoolOptions::new() diff --git a/crates/aether-data/adapters/sqlite/src/provider_catalog.rs b/crates/aether-data/adapters/sqlite/src/provider_catalog.rs index ee3e5c5a9..19680448a 100644 --- a/crates/aether-data/adapters/sqlite/src/provider_catalog.rs +++ b/crates/aether-data/adapters/sqlite/src/provider_catalog.rs @@ -1462,6 +1462,19 @@ WHERE id = ? "provider catalog OAuth api_key update must not be empty".to_string(), )); } + if update.expected_credential.as_ref().is_some_and(|expected| { + expected + .encrypted_api_key + .as_deref() + .is_some_and(|value| value.trim().is_empty()) + || expected.auth_type.trim().is_empty() + || expected.provider_id.trim().is_empty() + || expected.provider_type.trim().is_empty() + }) { + return Err(DataLayerError::InvalidInput( + "provider catalog OAuth credential fence must not contain empty fields".to_string(), + )); + } if !update.status_snapshot_patch.is_object() { return Err(DataLayerError::InvalidInput( "provider catalog status snapshot patch must be an object".to_string(), @@ -1518,6 +1531,22 @@ WHERE id = ? .push_bind(&update.key_id) .push(" AND auth_config IS ") .push_bind(update.expected_encrypted_auth_config.as_deref()); + if let Some(expected) = update.expected_credential.as_ref() { + builder + .push(" AND api_key IS ") + .push_bind(expected.encrypted_api_key.as_deref()) + .push(" AND auth_type = ") + .push_bind(&expected.auth_type) + .push(" AND provider_id = ") + .push_bind(&expected.provider_id) + .push( + " AND EXISTS (SELECT 1 FROM providers WHERE \ + providers.id = provider_api_keys.provider_id \ + AND providers.provider_type = ", + ) + .push_bind(&expected.provider_type) + .push(")"); + } let rows_affected = builder .build() .execute(&self.pool) @@ -2938,9 +2967,10 @@ mod tests { use aether_data_contracts::repository::provider_catalog::{ ProviderCatalogKeyAdaptiveState, ProviderCatalogKeyAdaptiveStateUpdate, ProviderCatalogKeyHealthStateUpdate, ProviderCatalogKeyListOrder, - ProviderCatalogKeyListQuery, ProviderCatalogKeyOAuthRuntimeStateCasUpdate, - ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogUpstreamMetadataNamespaceUpdate, - StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider, + ProviderCatalogKeyListQuery, ProviderCatalogKeyOAuthCredentialFence, + ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate, + ProviderCatalogUpstreamMetadataNamespaceUpdate, StoredProviderCatalogEndpoint, + StoredProviderCatalogKey, StoredProviderCatalogProvider, }; use serde_json::json; @@ -3246,6 +3276,12 @@ mod tests { let update = ProviderCatalogKeyOAuthRuntimeStateCasUpdate { key_id: key.id.clone(), expected_encrypted_auth_config: Some("encrypted-auth-v1".to_string()), + expected_credential: Some(ProviderCatalogKeyOAuthCredentialFence { + encrypted_api_key: Some("encrypted-api-key".to_string()), + auth_type: "oauth".to_string(), + provider_id: "oauth-cas-provider".to_string(), + provider_type: "codex".to_string(), + }), encrypted_auth_config: "encrypted-auth-v2".to_string(), encrypted_api_key_update: Some("encrypted-api-v2".to_string()), expires_at_unix_secs_update: Some(Some(4_102_555_900)), @@ -3304,6 +3340,24 @@ mod tests { assert_eq!(status["admin"], json!({"label": "keep"})); assert_eq!(status["runtime"], json!({"generation": 2})); + let stale_api_key_update = ProviderCatalogKeyOAuthRuntimeStateCasUpdate { + expected_encrypted_auth_config: Some("encrypted-auth-v2".to_string()), + expected_credential: Some(ProviderCatalogKeyOAuthCredentialFence { + encrypted_api_key: Some("encrypted-api-key".to_string()), + auth_type: "oauth".to_string(), + provider_id: "oauth-cas-provider".to_string(), + provider_type: "codex".to_string(), + }), + encrypted_auth_config: "encrypted-auth-v3".to_string(), + status_snapshot_patch: json!({}), + updated_at_unix_secs: Some(201), + ..update.clone() + }; + assert!(!repository + .compare_and_update_key_oauth_runtime_state(&stale_api_key_update) + .await + .expect("stale API key generation should conflict")); + let stale_update = ProviderCatalogKeyOAuthRuntimeStateCasUpdate { expected_encrypted_auth_config: Some("encrypted-auth-v1".to_string()), encrypted_auth_config: "encrypted-auth-v3".to_string(), diff --git a/crates/aether-data/contracts/src/repository/provider_catalog/mod.rs b/crates/aether-data/contracts/src/repository/provider_catalog/mod.rs index 03bbcee4b..7c09c24f9 100644 --- a/crates/aether-data/contracts/src/repository/provider_catalog/mod.rs +++ b/crates/aether-data/contracts/src/repository/provider_catalog/mod.rs @@ -5,10 +5,10 @@ pub use snapshot::ProviderCatalogSnapshot; pub use types::{ ProviderCatalogKeyAdaptiveState, ProviderCatalogKeyAdaptiveStateUpdate, ProviderCatalogKeyHealthStateUpdate, ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, - ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate, - ProviderCatalogKeyStatusSnapshotUpdate, ProviderCatalogReadRepository, - ProviderCatalogUpstreamMetadataNamespaceUpdate, ProviderCatalogWriteRepository, - StoredProviderCatalogEndpoint, StoredProviderCatalogKey, + ProviderCatalogKeyOAuthCredentialFence, ProviderCatalogKeyOAuthRuntimeStateCasUpdate, + ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate, + ProviderCatalogReadRepository, ProviderCatalogUpstreamMetadataNamespaceUpdate, + ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider, }; diff --git a/crates/aether-data/contracts/src/repository/provider_catalog/types.rs b/crates/aether-data/contracts/src/repository/provider_catalog/types.rs index d520b4d27..cca730e61 100644 --- a/crates/aether-data/contracts/src/repository/provider_catalog/types.rs +++ b/crates/aether-data/contracts/src/repository/provider_catalog/types.rs @@ -67,13 +67,28 @@ pub struct ProviderCatalogKeyStatusSnapshotUpdate { pub updated_at_unix_secs: Option, } +/// Credential context observed before an OAuth refresh started. Repositories +/// compare every field atomically with the runtime-state update so an +/// administrator replacement cannot be overwritten by an older refresh. +#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] +pub struct ProviderCatalogKeyOAuthCredentialFence { + /// Exact nullable ciphertext stored in `provider_api_keys.api_key`. + pub encrypted_api_key: Option, + pub auth_type: String, + pub provider_id: String, + pub provider_type: String, +} + /// Agent/runtime-owned OAuth state update fenced by the exact encrypted -/// auth_config observed before the refresh started. Repositories must update -/// only these fields and return `false` when the expected config changed. +/// auth_config and, when supplied, credential context observed before the +/// refresh started. Repositories must update only these fields and return +/// `false` when an expected value changed. #[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] pub struct ProviderCatalogKeyOAuthRuntimeStateCasUpdate { pub key_id: String, pub expected_encrypted_auth_config: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub expected_credential: Option, pub encrypted_auth_config: String, /// Optional access-token ciphertext replacement owned by refresh success. #[serde(default, skip_serializing_if = "Option::is_none")] diff --git a/crates/aether-data/runtime/src/repository/candidate_selection/memory.rs b/crates/aether-data/runtime/src/repository/candidate_selection/memory.rs index 25b811606..2906bf693 100644 --- a/crates/aether-data/runtime/src/repository/candidate_selection/memory.rs +++ b/crates/aether-data/runtime/src/repository/candidate_selection/memory.rs @@ -335,16 +335,13 @@ fn key_auth_channel_matches(row: &StoredMinimalCandidateSelectionRow, api_format && api_format == "openai:chat" } "vertex_ai" => { - (auth_type == "api_key" - && matches!( - api_format.as_str(), - "gemini:generate_content" | "gemini:embedding" - )) - || (matches!(auth_type.as_str(), "service_account" | "vertex_ai") - && matches!( - api_format.as_str(), - "claude:messages" | "gemini:generate_content" | "gemini:embedding" - )) + matches!( + auth_type.as_str(), + "api_key" | "service_account" | "vertex_ai" + ) && matches!( + api_format.as_str(), + "gemini:generate_content" | "gemini:embedding" + ) } _ => auth_type != "oauth", } @@ -640,6 +637,48 @@ mod tests { ); } + #[tokio::test] + async fn vertex_auth_matrix_rejects_claude_and_keeps_gemini_candidates() { + let mut candidates = Vec::new(); + for auth_type in ["api_key", "service_account", "vertex_ai"] { + for api_format in [ + "claude:messages", + "gemini:generate_content", + "gemini:embedding", + ] { + let provider_id = format!( + "vertex-{}-{}", + auth_type, + api_format.replace(':', "-").replace('_', "-") + ); + let mut row = sample_row(&provider_id, api_format, "vertex-model", 10); + row.provider_type = "vertex_ai".to_string(); + row.key_auth_type = auth_type.to_string(); + candidates.push(row); + } + } + let repository = InMemoryMinimalCandidateSelectionReadRepository::seed(candidates); + + let claude_rows = repository + .list_for_exact_api_format("claude:messages") + .await + .expect("list should succeed"); + assert!(claude_rows.is_empty()); + + for api_format in ["gemini:generate_content", "gemini:embedding"] { + let rows = repository + .list_for_exact_api_format(api_format) + .await + .expect("list should succeed"); + let mut auth_types = rows + .into_iter() + .map(|row| row.key_auth_type) + .collect::>(); + auth_types.sort(); + assert_eq!(auth_types, ["api_key", "service_account", "vertex_ai"]); + } + } + #[tokio::test] async fn filters_by_exact_api_format_only() { let repository = InMemoryMinimalCandidateSelectionReadRepository::seed(vec![ diff --git a/crates/aether-data/runtime/src/repository/provider_catalog/memory.rs b/crates/aether-data/runtime/src/repository/provider_catalog/memory.rs index ac06f6d73..b5a9c6af6 100644 --- a/crates/aether-data/runtime/src/repository/provider_catalog/memory.rs +++ b/crates/aether-data/runtime/src/repository/provider_catalog/memory.rs @@ -779,6 +779,15 @@ impl ProviderCatalogWriteRepository for InMemoryProviderCatalogReadRepository { .encrypted_api_key_update .as_deref() .is_some_and(|value| value.trim().is_empty()) + || update.expected_credential.as_ref().is_some_and(|expected| { + expected + .encrypted_api_key + .as_deref() + .is_some_and(|value| value.trim().is_empty()) + || expected.auth_type.trim().is_empty() + || expected.provider_id.trim().is_empty() + || expected.provider_type.trim().is_empty() + }) || !update.status_snapshot_patch.is_object() || update .upstream_metadata_patch @@ -799,13 +808,30 @@ impl ProviderCatalogWriteRepository for InMemoryProviderCatalogReadRepository { .index .write() .expect("provider catalog repository lock"); - let Some(key) = index.keys.get_mut(&update.key_id) else { + let Some(key) = index.keys.get(&update.key_id) else { return Ok(false); }; if key.encrypted_auth_config.as_deref() != update.expected_encrypted_auth_config.as_deref() { return Ok(false); } + if let Some(expected) = update.expected_credential.as_ref() { + let provider_type_matches = index + .providers + .get(&key.provider_id) + .is_some_and(|provider| provider.provider_type == expected.provider_type); + if key.encrypted_api_key != expected.encrypted_api_key + || key.auth_type != expected.auth_type + || key.provider_id != expected.provider_id + || !provider_type_matches + { + return Ok(false); + } + } + let key = index + .keys + .get_mut(&update.key_id) + .expect("provider catalog key was checked under the same write lock"); if let Some(encrypted_api_key) = update.encrypted_api_key_update.as_ref() { key.encrypted_api_key = Some(encrypted_api_key.clone()); } @@ -1170,10 +1196,10 @@ mod tests { use crate::repository::provider_catalog::{ ProviderCatalogKeyAdaptiveState, ProviderCatalogKeyAdaptiveStateUpdate, ProviderCatalogKeyHealthStateUpdate, ProviderCatalogKeyListOrder, - ProviderCatalogKeyListQuery, ProviderCatalogKeyOAuthRuntimeStateCasUpdate, - ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogReadRepository, - ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey, - StoredProviderCatalogProvider, + ProviderCatalogKeyListQuery, ProviderCatalogKeyOAuthCredentialFence, + ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate, + ProviderCatalogReadRepository, ProviderCatalogWriteRepository, + StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider, }; use crate::repository::usage::ProviderApiKeyUsageDelta; use serde_json::{json, Value}; @@ -1347,6 +1373,12 @@ mod tests { let update = ProviderCatalogKeyOAuthRuntimeStateCasUpdate { key_id: "key-1".to_string(), expected_encrypted_auth_config: Some("ciphertext-auth-1".to_string()), + expected_credential: Some(ProviderCatalogKeyOAuthCredentialFence { + encrypted_api_key: Some("ciphertext-placeholder".to_string()), + auth_type: "api_key".to_string(), + provider_id: "provider-1".to_string(), + provider_type: "custom".to_string(), + }), encrypted_auth_config: "ciphertext-auth-2".to_string(), encrypted_api_key_update: Some("ciphertext-api-2".to_string()), expires_at_unix_secs_update: Some(Some(456)), @@ -1391,6 +1423,99 @@ mod tests { ); } + #[tokio::test] + async fn oauth_runtime_cas_rejects_changed_credential_context() { + let repository = || { + InMemoryProviderCatalogReadRepository::seed( + vec![sample_provider("provider-1")], + vec![], + vec![sample_key("key-1", "provider-1") + .with_transport_fields( + None, + "ciphertext-api-1".to_string(), + Some("ciphertext-auth-1".to_string()), + None, + None, + None, + None, + None, + None, + ) + .expect("key transport should build")], + ) + }; + let update = || ProviderCatalogKeyOAuthRuntimeStateCasUpdate { + key_id: "key-1".to_string(), + expected_encrypted_auth_config: Some("ciphertext-auth-1".to_string()), + expected_credential: Some(ProviderCatalogKeyOAuthCredentialFence { + encrypted_api_key: Some("ciphertext-api-1".to_string()), + auth_type: "api_key".to_string(), + provider_id: "provider-1".to_string(), + provider_type: "custom".to_string(), + }), + encrypted_auth_config: "ciphertext-auth-2".to_string(), + encrypted_api_key_update: Some("ciphertext-api-2".to_string()), + expires_at_unix_secs_update: None, + oauth_invalid_at_unix_secs: None, + oauth_invalid_reason: None, + upstream_metadata_patch: None, + status_snapshot_patch: json!({}), + reset_error_count: false, + updated_at_unix_secs: Some(123), + }; + + let api_key_repository = repository(); + let mut key = api_key_repository + .list_keys_by_ids(&["key-1".to_string()]) + .await + .expect("key should load") + .pop() + .expect("key should exist"); + key.encrypted_api_key = Some("ciphertext-admin".to_string()); + api_key_repository + .update_key(&key) + .await + .expect("api key replacement should persist"); + assert!(!api_key_repository + .compare_and_update_key_oauth_runtime_state(&update()) + .await + .expect("API key mismatch should be a CAS miss")); + + let auth_type_repository = repository(); + let mut key = auth_type_repository + .list_keys_by_ids(&["key-1".to_string()]) + .await + .expect("key should load") + .pop() + .expect("key should exist"); + key.auth_type = "oauth".to_string(); + auth_type_repository + .update_key(&key) + .await + .expect("auth type replacement should persist"); + assert!(!auth_type_repository + .compare_and_update_key_oauth_runtime_state(&update()) + .await + .expect("auth type mismatch should be a CAS miss")); + + let provider_repository = repository(); + let mut provider = provider_repository + .list_providers_by_ids(&["provider-1".to_string()]) + .await + .expect("provider should load") + .pop() + .expect("provider should exist"); + provider.provider_type = "codex".to_string(); + provider_repository + .update_provider(&provider) + .await + .expect("provider type replacement should persist"); + assert!(!provider_repository + .compare_and_update_key_oauth_runtime_state(&update()) + .await + .expect("provider type mismatch should be a CAS miss")); + } + #[tokio::test] async fn materializes_codex_window_usage_stats_delta_in_memory() { let mut key = sample_key("key-1", "provider-1"); diff --git a/crates/aether-data/runtime/src/repository/provider_catalog/mod.rs b/crates/aether-data/runtime/src/repository/provider_catalog/mod.rs index 8a293b93e..e9b4e7d29 100644 --- a/crates/aether-data/runtime/src/repository/provider_catalog/mod.rs +++ b/crates/aether-data/runtime/src/repository/provider_catalog/mod.rs @@ -4,8 +4,9 @@ mod memory; pub(crate) use aether_data_contracts::repository::provider_catalog::{ ProviderCatalogKeyAdaptiveState, ProviderCatalogKeyAdaptiveStateUpdate, ProviderCatalogKeyHealthStateUpdate, ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, - ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate, - ProviderCatalogKeyStatusSnapshotUpdate, ProviderCatalogReadRepository, ProviderCatalogSnapshot, + ProviderCatalogKeyOAuthCredentialFence, ProviderCatalogKeyOAuthRuntimeStateCasUpdate, + ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate, + ProviderCatalogReadRepository, ProviderCatalogSnapshot, ProviderCatalogUpstreamMetadataNamespaceUpdate, ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage, diff --git a/crates/aether-provider/transport/src/anthropic_compat.rs b/crates/aether-provider/transport/src/anthropic_compat.rs new file mode 100644 index 000000000..d137889d2 --- /dev/null +++ b/crates/aether-provider/transport/src/anthropic_compat.rs @@ -0,0 +1,376 @@ +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +use crate::snapshot::GatewayProviderTransportSnapshot; + +/// Provider-side compatibility applied to otherwise same-format Anthropic requests. +/// +/// Native Anthropic endpoints should remain transparent. The legacy Claude Code +/// profile is opt-in, except for the existing `claude_code` provider type where it +/// remains the backwards-compatible default. Endpoint config takes precedence over +/// provider config. The canonical field is `anthropic.compatibility_profile`; +/// explicitly Anthropic-namespaced legacy spellings remain accepted. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Deserialize, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum AnthropicCompatibilityProfile { + #[default] + #[serde( + alias = "native", + alias = "anthropic", + alias = "none", + alias = "transparent" + )] + NativeTransparent, + #[serde( + alias = "claude_code", + alias = "claude-code", + alias = "legacy", + alias = "same_format_compat" + )] + ClaudeCodeLegacy, +} + +impl AnthropicCompatibilityProfile { + pub const fn as_str(self) -> &'static str { + match self { + Self::NativeTransparent => "native_transparent", + Self::ClaudeCodeLegacy => "claude_code_legacy", + } + } + + pub const fn uses_claude_code_compatibility(self) -> bool { + matches!(self, Self::ClaudeCodeLegacy) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] +#[error("unknown Anthropic compatibility profile")] +pub struct AnthropicCompatibilityProfileConfigError; + +/// Validate the optional Anthropic compatibility fields in a provider or +/// endpoint config without resolving provider defaults or mutating config. +pub fn validate_anthropic_compatibility_profile_config( + config: Option<&Value>, +) -> Result<(), AnthropicCompatibilityProfileConfigError> { + match profile_from_config(config) { + ConfiguredProfile::Absent | ConfiguredProfile::Valid(_) => Ok(()), + ConfiguredProfile::Invalid => Err(AnthropicCompatibilityProfileConfigError), + } +} + +pub fn resolve_anthropic_compatibility_profile( + transport: &GatewayProviderTransportSnapshot, + provider_api_format: &str, +) -> AnthropicCompatibilityProfile { + if !aether_ai_formats::api_format_alias_matches(provider_api_format, "claude:messages") { + return AnthropicCompatibilityProfile::NativeTransparent; + } + + for (scope, resolution) in [ + ( + "endpoint", + profile_from_config(transport.endpoint.config.as_ref()), + ), + ( + "provider", + profile_from_config(transport.provider.config.as_ref()), + ), + ] { + match resolution { + ConfiguredProfile::Valid(profile) => return profile, + ConfiguredProfile::Invalid => { + tracing::warn!( + event_name = "anthropic_compatibility_profile_invalid", + log_type = "ops", + provider_id = %transport.provider.id, + endpoint_id = %transport.endpoint.id, + config_scope = scope, + "invalid Anthropic compatibility profile; using native transparent behavior" + ); + return AnthropicCompatibilityProfile::NativeTransparent; + } + ConfiguredProfile::Absent => {} + } + } + + if transport + .provider + .provider_type + .trim() + .eq_ignore_ascii_case("claude_code") + { + AnthropicCompatibilityProfile::ClaudeCodeLegacy + } else { + AnthropicCompatibilityProfile::NativeTransparent + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ConfiguredProfile { + Absent, + Valid(AnthropicCompatibilityProfile), + Invalid, +} + +fn profile_from_config(config: Option<&Value>) -> ConfiguredProfile { + let Some(config) = config.and_then(Value::as_object) else { + return ConfiguredProfile::Absent; + }; + + for (container, fields) in [ + ( + config.get("anthropic").and_then(Value::as_object), + &["compatibility_profile", "compatibilityProfile", "profile"][..], + ), + ( + config + .get("anthropic_compatibility") + .and_then(Value::as_object), + &["profile", "compatibility_profile", "compatibilityProfile"][..], + ), + ( + config + .get("anthropicCompatibility") + .and_then(Value::as_object), + &["profile", "compatibilityProfile", "compatibility_profile"][..], + ), + ] { + let Some(container) = container else { + continue; + }; + for &field in fields { + if let Some(profile) = container.get(field).and_then(parse_profile_value) { + return ConfiguredProfile::Valid(profile); + } + if container.contains_key(field) { + return ConfiguredProfile::Invalid; + } + } + } + + for field in [ + "anthropic_compatibility_profile", + "anthropicCompatibilityProfile", + ] { + if let Some(value) = config.get(field) { + return parse_profile_value(value) + .map(ConfiguredProfile::Valid) + .unwrap_or(ConfiguredProfile::Invalid); + } + } + + let Some(claude_code_advanced) = config + .get("claude_code_advanced") + .and_then(Value::as_object) + else { + return ConfiguredProfile::Absent; + }; + for field in ["compatibility_profile", "compatibilityProfile"] { + if let Some(value) = claude_code_advanced.get(field) { + return parse_profile_value(value) + .map(ConfiguredProfile::Valid) + .unwrap_or(ConfiguredProfile::Invalid); + } + } + ConfiguredProfile::Absent +} + +fn parse_profile_value(value: &Value) -> Option { + if let Some(object) = value.as_object() { + return ["kind", "name", "profile"] + .into_iter() + .find_map(|field| object.get(field).and_then(parse_profile_value)); + } + serde_json::from_value(value.clone()).ok() +} + +#[cfg(test)] +mod tests { + use super::{ + resolve_anthropic_compatibility_profile, validate_anthropic_compatibility_profile_config, + AnthropicCompatibilityProfile, + }; + use crate::snapshot::{ + GatewayProviderTransportEndpoint, GatewayProviderTransportKey, + GatewayProviderTransportProvider, GatewayProviderTransportSnapshot, + }; + use serde_json::json; + + fn sample_transport(provider_type: &str) -> GatewayProviderTransportSnapshot { + GatewayProviderTransportSnapshot { + provider: GatewayProviderTransportProvider { + id: "provider-1".to_string(), + name: "provider".to_string(), + provider_type: provider_type.to_string(), + website: None, + is_active: true, + keep_priority_on_conversion: false, + enable_format_conversion: false, + concurrent_limit: None, + max_retries: None, + proxy: None, + request_timeout_secs: None, + stream_first_byte_timeout_secs: None, + config: None, + }, + endpoint: GatewayProviderTransportEndpoint { + id: "endpoint-1".to_string(), + provider_id: "provider-1".to_string(), + api_format: "claude:messages".to_string(), + api_family: Some("claude".to_string()), + endpoint_kind: Some("chat".to_string()), + is_active: true, + base_url: "https://api.anthropic.com".to_string(), + header_rules: None, + body_rules: None, + max_retries: None, + custom_path: None, + config: None, + format_acceptance_config: None, + proxy: None, + }, + key: GatewayProviderTransportKey { + id: "key-1".to_string(), + provider_id: "provider-1".to_string(), + name: "key".to_string(), + auth_type: "api_key".to_string(), + is_active: true, + api_formats: None, + auth_type_by_format: None, + allow_auth_channel_mismatch_formats: None, + allowed_models: None, + capabilities: None, + rate_multipliers: None, + global_priority_by_format: None, + expires_at_unix_secs: None, + proxy: None, + fingerprint: None, + upstream_metadata: None, + decrypted_api_key: "secret".to_string(), + decrypted_auth_config: None, + }, + } + } + + #[test] + fn validates_configured_profile_without_resolving_provider_defaults() { + assert!(validate_anthropic_compatibility_profile_config(None).is_ok()); + assert!( + validate_anthropic_compatibility_profile_config(Some(&json!({ + "anthropic_compatibility": {"profile": "native_transparent"} + }))) + .is_ok() + ); + assert!( + validate_anthropic_compatibility_profile_config(Some(&json!({ + "anthropic": {"compatibility_profile": "claude_code_legacy"} + }))) + .is_ok() + ); + + let error = validate_anthropic_compatibility_profile_config(Some(&json!({ + "anthropic_compatibility": {"profile": "claude_cod_typo"} + }))) + .expect_err("unknown profile should be rejected"); + assert_eq!(error.to_string(), "unknown Anthropic compatibility profile"); + } + + #[test] + fn ignores_generic_compatibility_fields_owned_by_other_transports() { + let config = json!({ + "compatibility_profile": "strict", + "compatibility": {"profile": "v2"}, + "adaptation": {"profile": "legacy"} + }); + + assert!(validate_anthropic_compatibility_profile_config(Some(&config)).is_ok()); + let mut transport = sample_transport("custom"); + transport.provider.config = Some(config); + assert_eq!( + resolve_anthropic_compatibility_profile(&transport, "claude:messages"), + AnthropicCompatibilityProfile::NativeTransparent + ); + } + + #[test] + fn native_anthropic_is_transparent_by_default() { + let transport = sample_transport("custom"); + + assert_eq!( + resolve_anthropic_compatibility_profile(&transport, "claude:messages"), + AnthropicCompatibilityProfile::NativeTransparent + ); + } + + #[test] + fn claude_code_keeps_legacy_compatibility_by_default() { + let transport = sample_transport("claude_code"); + + assert_eq!( + resolve_anthropic_compatibility_profile(&transport, "claude:messages"), + AnthropicCompatibilityProfile::ClaudeCodeLegacy + ); + } + + #[test] + fn endpoint_profile_overrides_provider_and_legacy_defaults() { + let mut transport = sample_transport("claude_code"); + transport.provider.config = Some(json!({ + "anthropic": {"compatibility_profile": "claude_code_legacy"} + })); + transport.endpoint.config = Some(json!({ + "anthropic_compatibility": { + "profile": "native_transparent" + } + })); + + assert_eq!( + resolve_anthropic_compatibility_profile(&transport, "claude:messages"), + AnthropicCompatibilityProfile::NativeTransparent + ); + } + + #[test] + fn provider_profile_can_explicitly_enable_legacy_compatibility() { + let mut transport = sample_transport("custom"); + transport.provider.config = Some(json!({ + "anthropic": { + "compatibility_profile": "claude_code" + } + })); + + assert_eq!( + resolve_anthropic_compatibility_profile(&transport, "claude:messages"), + AnthropicCompatibilityProfile::ClaudeCodeLegacy + ); + } + + #[test] + fn anthropic_profile_is_ignored_for_non_anthropic_formats() { + let mut transport = sample_transport("claude_code"); + transport.endpoint.config = Some(json!({ + "anthropic": {"compatibility_profile": "claude_code_legacy"} + })); + + assert_eq!( + resolve_anthropic_compatibility_profile(&transport, "openai:chat"), + AnthropicCompatibilityProfile::NativeTransparent + ); + } + + #[test] + fn invalid_explicit_profile_fails_closed_instead_of_using_legacy_default() { + let mut transport = sample_transport("claude_code"); + transport.endpoint.config = Some(json!({ + "anthropic_compatibility": {"profile": "native_transparnt"} + })); + transport.provider.config = Some(json!({ + "anthropic_compatibility": {"profile": "claude_code_legacy"} + })); + + assert_eq!( + resolve_anthropic_compatibility_profile(&transport, "claude:messages"), + AnthropicCompatibilityProfile::NativeTransparent + ); + } +} diff --git a/crates/aether-provider/transport/src/auth.rs b/crates/aether-provider/transport/src/auth.rs index 707686ea8..d7a8b5dc2 100644 --- a/crates/aether-provider/transport/src/auth.rs +++ b/crates/aether-provider/transport/src/auth.rs @@ -1,8 +1,8 @@ use std::collections::BTreeMap; use super::headers::{ - normalize_upstream_accept_encoding, should_skip_upstream_complete_passthrough_header, - should_skip_upstream_passthrough_header, + is_aether_internal_header, is_upstream_credential_header, normalize_upstream_accept_encoding, + should_skip_upstream_complete_passthrough_header, should_skip_upstream_passthrough_header, }; use super::snapshot::GatewayProviderTransportSnapshot; @@ -30,6 +30,9 @@ fn collect_passthrough_headers( for (key, value) in extra_headers { let normalized_key = key.to_ascii_lowercase(); + if should_skip_upstream_passthrough_header(&normalized_key) { + continue; + } let Some(value) = normalize_passthrough_header_value(&normalized_key, value) else { continue; }; @@ -60,6 +63,9 @@ fn collect_complete_passthrough_headers( for (key, value) in extra_headers { let normalized_key = key.to_ascii_lowercase(); + if should_skip_upstream_complete_passthrough_header(&normalized_key) { + continue; + } let Some(value) = normalize_passthrough_header_value(&normalized_key, value) else { continue; }; @@ -136,7 +142,7 @@ pub fn build_complete_passthrough_headers_with_auth( content_type: Option<&str>, ) -> BTreeMap { let mut out = build_complete_passthrough_headers(headers, extra_headers, content_type); - ensure_upstream_auth_header(&mut out, auth_header, auth_value); + replace_upstream_auth_headers(&mut out, auth_header, auth_value); out } @@ -155,6 +161,24 @@ pub fn build_claude_passthrough_headers( content_type, ); + for (name, value) in extra_headers { + let key = name.to_ascii_lowercase(); + let value = value.trim(); + if value.is_empty() || !should_restore_claude_passthrough_header(&key) { + continue; + } + + if key == "anthropic-beta" { + let merged = merge_comma_header_values(out.get(&key).map(String::as_str), Some(value)); + if let Some(merged) = merged { + out.insert(key, merged); + } + continue; + } + + out.insert(key, value.to_string()); + } + for (name, value) in headers.iter() { let Ok(value) = value.to_str() else { continue; @@ -188,7 +212,7 @@ pub fn build_passthrough_headers_with_auth( extra_headers: &BTreeMap, ) -> BTreeMap { let mut out = collect_passthrough_headers(headers, extra_headers); - ensure_upstream_auth_header(&mut out, auth_header, auth_value); + replace_upstream_auth_headers(&mut out, auth_header, auth_value); out.remove("content-length"); out } @@ -200,7 +224,9 @@ pub fn ensure_upstream_auth_header( ) { let header_name = auth_header.trim().to_ascii_lowercase(); let header_value = auth_value.trim(); - if header_name.is_empty() || header_value.is_empty() { + headers.retain(|name, _| !is_aether_internal_header(name)); + if header_name.is_empty() || header_value.is_empty() || is_aether_internal_header(&header_name) + { return; } @@ -213,6 +239,16 @@ pub fn ensure_upstream_auth_header( } } +pub(crate) fn replace_upstream_auth_headers( + headers: &mut BTreeMap, + auth_header: &str, + auth_value: &str, +) { + headers + .retain(|name, _| !is_upstream_credential_header(name) && !is_aether_internal_header(name)); + ensure_upstream_auth_header(headers, auth_header, auth_value); +} + fn should_restore_claude_passthrough_header(name: &str) -> bool { name.starts_with("anthropic-") || name.starts_with("x-stainless-") || name == "x-app" } @@ -405,7 +441,14 @@ mod tests { &headers, "x-api-key", "sk-upstream-claude", - &BTreeMap::from([("anthropic-beta".to_string(), "custom-beta".to_string())]), + &BTreeMap::from([ + ("anthropic-beta".to_string(), "custom-beta".to_string()), + ("authorization".to_string(), "Bearer extra".to_string()), + ( + "x-aether-auth-user-id".to_string(), + "user-private".to_string(), + ), + ]), Some("application/json"), ); @@ -422,6 +465,8 @@ mod tests { Some("v22.14.0") ); assert_eq!(built.get("x-app").map(String::as_str), Some("cli")); + assert_eq!(built.get("authorization"), None); + assert!(built.keys().all(|name| !name.starts_with("x-aether-"))); assert_eq!( built.get("x-api-key").map(String::as_str), Some("sk-upstream-claude") @@ -488,12 +533,34 @@ mod tests { "authorization", http::HeaderValue::from_static("Bearer client-token"), ); + headers.insert("api-key", http::HeaderValue::from_static("client-api-key")); + headers.insert( + "x-api-key", + http::HeaderValue::from_static("client-x-api-key"), + ); + headers.insert("cookie", http::HeaderValue::from_static("session=client")); + headers.insert( + "proxy-authorization", + http::HeaderValue::from_static("Basic client-proxy"), + ); + headers.insert( + "x-aether-auth-user-id", + http::HeaderValue::from_static("user-private"), + ); let built = build_complete_passthrough_headers_with_auth( &headers, "x-api-key", "sk-upstream", - &BTreeMap::new(), + &BTreeMap::from([ + ("authorization".to_string(), "Bearer extra".to_string()), + ("cookie".to_string(), "session=extra".to_string()), + ("x-api-key".to_string(), "extra-x-api-key".to_string()), + ( + "x-aether-auth-api-key-id".to_string(), + "key-private".to_string(), + ), + ]), Some("application/json"), ); @@ -507,6 +574,10 @@ mod tests { ); assert_eq!(built.get("x-app").map(String::as_str), Some("cli")); assert_eq!(built.get("authorization"), None); + assert_eq!(built.get("api-key"), None); + assert_eq!(built.get("cookie"), None); + assert_eq!(built.get("proxy-authorization"), None); + assert!(built.keys().all(|name| !name.starts_with("x-aether-"))); assert_eq!( built.get("x-api-key").map(String::as_str), Some("sk-upstream") diff --git a/crates/aether-provider/transport/src/claude_code/fingerprint.rs b/crates/aether-provider/transport/src/claude_code/fingerprint.rs index cbe87daa0..c72e58e93 100644 --- a/crates/aether-provider/transport/src/claude_code/fingerprint.rs +++ b/crates/aether-provider/transport/src/claude_code/fingerprint.rs @@ -2,149 +2,19 @@ use aether_contracts::{ TRANSPORT_BACKEND_REQWEST_RUSTLS, TRANSPORT_HTTP_MODE_AUTO, TRANSPORT_POOL_SCOPE_KEY, }; use serde_json::{Map, Value}; -use sha2::{Digest, Sha256}; use uuid::Uuid; -// Chrome impersonate profiles -const CHROME_IMPERSONATE_PROFILES: &[&str] = &[ - "chrome110", - "chrome116", - "chrome119", - "chrome120", - "chrome123", - "chrome124", - "chrome131", - "chrome133", -]; +use super::profile::current_claude_code_transport_identity_profile; -const CHROME_VERSIONS: &[(&str, &str)] = &[ - ("chrome110", "110.0.5481.177"), - ("chrome116", "116.0.5845.188"), - ("chrome119", "119.0.6045.214"), - ("chrome120", "120.0.6099.216"), - ("chrome123", "123.0.6312.122"), - ("chrome124", "124.0.6367.243"), - ("chrome131", "131.0.6778.265"), - ("chrome133", "133.0.6943.142"), -]; - -// (os, arch, platform_token, platform_info) -const PLATFORM_VARIANTS: &[(&str, &str, &str, &str)] = &[ - ("Linux", "x64", "X11; Linux x86_64", "Linux x86_64"), - ("Linux", "arm64", "X11; Linux arm64", "Linux arm64"), - ( - "Windows", - "x64", - "Windows NT 10.0; Win64; x64", - "Windows x64", - ), - ( - "MacOS", - "x64", - "Macintosh; Intel Mac OS X 10_15_7", - "Darwin x64", - ), - ( - "MacOS", - "arm64", - "Macintosh; ARM Mac OS X 14_0_0", - "Darwin arm64", - ), -]; - -const STAINLESS_PACKAGE_VERSIONS: &[&str] = &["0.68.0", "0.69.0", "0.70.0", "0.71.0"]; -const NODE_VERSIONS: &[&str] = &["v20.18.1", "v22.12.0", "v22.14.0", "v24.13.0"]; -const ELECTRON_VERSIONS: &[&str] = &["35.5.1", "36.7.1", "37.3.0", "38.7.0", "39.2.3"]; -const STAINLESS_TIMEOUTS: &[&str] = &["600", "900"]; -const CLAUDE_CODE_TRANSPORT_PROFILE_ID: &str = "claude_code_nodejs"; - -/// Deterministic hash-based index picker, compatible with Python implementation. -/// Each `slot` produces a different selection from the same seed. -struct SeededPicker { - seed_bytes: [u8; 32], -} - -impl SeededPicker { - fn new(seed: &str) -> Self { - let mut hasher = Sha256::new(); - hasher.update(seed.as_bytes()); - Self { - seed_bytes: hasher.finalize().into(), - } - } - - /// Pick an index from `[0, len)` using a specific slot. - /// Different slots produce independent-looking selections from the same seed. - fn pick(&self, slot: u8, len: usize) -> usize { - if len == 0 { - return 0; - } - // Hash seed_bytes + slot to get a new digest, take first 8 bytes as u64 - let mut hasher = Sha256::new(); - hasher.update(self.seed_bytes); - hasher.update([slot]); - let hash = hasher.finalize(); - let value = u64::from_be_bytes(hash[..8].try_into().unwrap()); - (value % len as u64) as usize - } -} - -fn chrome_version_for_profile(profile: &str) -> &'static str { - for (p, v) in CHROME_VERSIONS { - if p.eq_ignore_ascii_case(profile) { - return v; - } - } - "120.0.6099.216" -} - -fn build_user_agent(platform_token: &str, chrome_version: &str, electron_version: &str) -> String { - format!( - "Mozilla/5.0 ({platform_token}) AppleWebKit/537.36 (KHTML, like Gecko) \ - Chrome/{chrome_version} Electron/{electron_version} Safari/537.36" - ) -} - -fn resolve_platform_token(os: &str, arch: &str) -> &'static str { - let os_lower = os.to_ascii_lowercase(); - let arch_lower = arch.to_ascii_lowercase(); - - if os_lower.starts_with("win") { - return "Windows NT 10.0; Win64; x64"; - } - if matches!(os_lower.as_str(), "darwin" | "mac" | "macos") { - return if matches!(arch_lower.as_str(), "arm64" | "aarch64") { - "Macintosh; ARM Mac OS X 14_0_0" - } else { - "Macintosh; Intel Mac OS X 10_15_7" - }; - } - if matches!(arch_lower.as_str(), "arm64" | "aarch64") { - "X11; Linux arm64" - } else { - "X11; Linux x86_64" - } -} - -/// Generate a complete Claude Code transport fingerprint from a seed. +/// Generate the transport-profile metadata and per-key pool partition from a +/// stable key seed. HTTP identity headers are owned by the versioned profile +/// and are not overridden from this stored value. pub fn generate_fingerprint(seed: &str) -> Value { wrap_header_fingerprint(generate_header_fingerprint(seed)) } fn generate_header_fingerprint(seed: &str) -> Value { - let picker = SeededPicker::new(seed); - - let impersonate = - CHROME_IMPERSONATE_PROFILES[picker.pick(0, CHROME_IMPERSONATE_PROFILES.len())]; - let chrome_version = chrome_version_for_profile(impersonate); - let node_version = NODE_VERSIONS[picker.pick(1, NODE_VERSIONS.len())]; - let electron_version = ELECTRON_VERSIONS[picker.pick(2, ELECTRON_VERSIONS.len())]; - let platform = PLATFORM_VARIANTS[picker.pick(3, PLATFORM_VARIANTS.len())]; - let (stainless_os, stainless_arch, platform_token, platform_info) = platform; - let stainless_package_version = - STAINLESS_PACKAGE_VERSIONS[picker.pick(4, STAINLESS_PACKAGE_VERSIONS.len())]; - let stainless_timeout = STAINLESS_TIMEOUTS[picker.pick(5, STAINLESS_TIMEOUTS.len())]; - + let profile = *current_claude_code_transport_identity_profile(); let vscode_session_id = Uuid::new_v5( &Uuid::NAMESPACE_URL, format!("aether:fingerprint:{seed}").as_bytes(), @@ -152,32 +22,36 @@ fn generate_header_fingerprint(seed: &str) -> Value { .simple() .to_string(); - let user_agent = build_user_agent(platform_token, chrome_version, electron_version); - serde_json::json!({ - "impersonate": impersonate, - "stainless_package_version": stainless_package_version, - "stainless_os": stainless_os, - "stainless_arch": stainless_arch, - "stainless_runtime_version": node_version, - "stainless_timeout": stainless_timeout, - "node_version": node_version, - "chrome_version": chrome_version, - "electron_version": electron_version, + "identity_profile_version": profile.version().as_str(), + "cli_version": profile.cli_version(), + "billing_cli_version": profile.billing_cli_version(), + "stainless_lang": profile.stainless_lang(), + "stainless_package_version": profile.stainless_package_version(), + "stainless_os": profile.stainless_os(), + "stainless_arch": profile.stainless_arch(), + "stainless_runtime": profile.stainless_runtime(), + "stainless_runtime_version": profile.stainless_runtime_version(), + "stainless_retry_count": profile.stainless_retry_count(), + "stainless_timeout": profile.stainless_timeout(), "vscode_session_id": vscode_session_id, - "platform_info": platform_info, - "user_agent": user_agent, + "user_agent": profile.user_agent(), }) } fn wrap_header_fingerprint(header_fingerprint: Value) -> Value { + let profile = *current_claude_code_transport_identity_profile(); serde_json::json!({ "transport_profile": { - "profile_id": CLAUDE_CODE_TRANSPORT_PROFILE_ID, + "profile_id": profile.transport_profile_id(), "backend": TRANSPORT_BACKEND_REQWEST_RUSTLS, "http_mode": TRANSPORT_HTTP_MODE_AUTO, "pool_scope": TRANSPORT_POOL_SCOPE_KEY, "header_fingerprint": header_fingerprint, + "extra": { + "claude_code_identity_profile_version": profile.version().as_str(), + "claude_code_cli_version": profile.cli_version(), + } } }) } @@ -192,82 +66,31 @@ pub fn header_fingerprint_from_fingerprint(fingerprint: &Value) -> Option<&Map Value { - let random_seed = Uuid::new_v4().to_string(); - generate_fingerprint(&random_seed) + generate_fingerprint(&Uuid::new_v4().to_string()) } -/// Sanitize an existing fingerprint JSON, filling missing fields with -/// deterministic fallbacks derived from `key_id`. +/// Upgrade stored transport metadata to the current typed identity profile. +/// Only the per-key pool partition (kept under its legacy session-id field) is +/// retained; fixed CLI and Stainless values remain one coherent version set. pub fn sanitize_fingerprint(raw: &Value, key_id: &str) -> Value { - let generated = generate_header_fingerprint(key_id); - let gen_map = generated.as_object().unwrap(); - let raw_map = header_fingerprint_from_fingerprint(raw); + let mut generated = generate_header_fingerprint(key_id); + let Some(generated) = generated.as_object_mut() else { + return generate_fingerprint(key_id); + }; - let mut out = Map::new(); - - // Start with generated values, then overlay non-empty raw values - for (key, gen_value) in gen_map { - let value = raw_map - .and_then(|raw_map| raw_map.get(key)) - .and_then(Value::as_str) - .map(str::trim) - .filter(|v| !v.is_empty()) - .map(|v| Value::String(v.to_string())) - .unwrap_or_else(|| gen_value.clone()); - out.insert(key.clone(), value); - } - - // Normalize impersonate to known profile - let impersonate = out - .get("impersonate") - .and_then(Value::as_str) - .unwrap_or_default() - .to_ascii_lowercase(); - let is_known = CHROME_IMPERSONATE_PROFILES - .iter() - .any(|p| p.eq_ignore_ascii_case(&impersonate)); - if !is_known { - out.insert("impersonate".to_string(), gen_map["impersonate"].clone()); - } - - // Ensure chrome_version matches impersonate profile - let profile = out - .get("impersonate") - .and_then(Value::as_str) - .unwrap_or_default(); - let chrome_version = chrome_version_for_profile(profile); - out.insert( - "chrome_version".to_string(), - Value::String(chrome_version.to_string()), - ); - - // Rebuild user_agent if missing - let has_ua = out - .get("user_agent") + if let Some(session_id) = header_fingerprint_from_fingerprint(raw) + .and_then(|raw| raw.get("vscode_session_id")) .and_then(Value::as_str) .map(str::trim) - .is_some_and(|v| !v.is_empty()); - if !has_ua { - let os = out - .get("stainless_os") - .and_then(Value::as_str) - .unwrap_or("Linux"); - let arch = out - .get("stainless_arch") - .and_then(Value::as_str) - .unwrap_or("x64"); - let electron = out - .get("electron_version") - .and_then(Value::as_str) - .unwrap_or("38.7.0"); - let platform_token = resolve_platform_token(os, arch); - out.insert( - "user_agent".to_string(), - Value::String(build_user_agent(platform_token, chrome_version, electron)), + .filter(|value| !value.is_empty()) + { + generated.insert( + "vscode_session_id".to_string(), + Value::String(session_id.to_string()), ); } - wrap_header_fingerprint(Value::Object(out)) + wrap_header_fingerprint(Value::Object(generated.clone())) } #[cfg(test)] @@ -282,101 +105,51 @@ mod tests { } #[test] - fn different_seeds_produce_different_fingerprints() { + fn different_seeds_produce_different_session_fingerprints() { let fp1 = generate_fingerprint("key-1"); let fp2 = generate_fingerprint("key-2"); - // At least one field should differ (statistically near-certain) assert_ne!(fp1, fp2); } #[test] - fn generated_fingerprint_has_all_fields() { + fn generated_fingerprint_matches_current_identity_profile() { let fp = generate_fingerprint("test-key"); - let expected_keys = [ - "impersonate", - "stainless_package_version", - "stainless_os", - "stainless_arch", - "stainless_runtime_version", - "stainless_timeout", - "node_version", - "chrome_version", - "electron_version", - "vscode_session_id", - "platform_info", - "user_agent", - ]; - let map = header_fingerprint_from_fingerprint(&fp).unwrap(); - for key in expected_keys { - assert!(map.contains_key(key), "missing field: {key}"); - let value = map[key].as_str().unwrap(); - assert!(!value.is_empty(), "empty field: {key}"); - } - } + let map = header_fingerprint_from_fingerprint(&fp).expect("header fingerprint"); - #[test] - fn sanitize_preserves_user_overrides() { - let raw = serde_json::json!({ - "transport_profile": { - "profile_id": "claude_code_nodejs", - "header_fingerprint": { - "stainless_os": "MacOS", - "stainless_arch": "arm64", - "stainless_timeout": "900", - "user_agent": "Custom-Agent/1.0" - } - } - }); - let sanitized = sanitize_fingerprint(&raw, "test-key"); - let map = header_fingerprint_from_fingerprint(&sanitized).unwrap(); - assert_eq!(map["stainless_os"].as_str(), Some("MacOS")); - assert_eq!(map["stainless_arch"].as_str(), Some("arm64")); - assert_eq!(map["stainless_timeout"].as_str(), Some("900")); - assert_eq!(map["user_agent"].as_str(), Some("Custom-Agent/1.0")); - // Other fields should be filled from generation - assert!(map.contains_key("impersonate")); - assert!(map.contains_key("stainless_package_version")); - } - - #[test] - fn sanitize_fills_missing_fields_from_seed() { - let raw = serde_json::json!({}); - let sanitized = sanitize_fingerprint(&raw, "test-key"); - let generated = generate_header_fingerprint("test-key"); - // All fields should match generated since raw is empty - let s = header_fingerprint_from_fingerprint(&sanitized).unwrap(); - let g = generated.as_object().unwrap(); - for key in g.keys() { - assert!(s.contains_key(key), "sanitized missing key: {key}"); - assert!( - !s[key].as_str().unwrap().is_empty(), - "sanitized empty key: {key}" - ); - } - } - - #[test] - fn sanitize_normalizes_unknown_impersonate_profile() { - let raw = serde_json::json!({ - "transport_profile": { - "profile_id": "claude_code_nodejs", - "header_fingerprint": { - "impersonate": "firefox99" - } - } - }); - let sanitized = sanitize_fingerprint(&raw, "test-key"); - let profile = header_fingerprint_from_fingerprint(&sanitized).unwrap()["impersonate"] - .as_str() - .unwrap(); - assert!( - CHROME_IMPERSONATE_PROFILES - .iter() - .any(|p| p.eq_ignore_ascii_case(profile)), - "should normalize to known profile, got: {profile}" + assert_eq!(map["identity_profile_version"], "2026-04"); + assert_eq!(map["cli_version"], "2.1.161"); + assert_eq!(map["billing_cli_version"], "2.1.161"); + assert_eq!(map["stainless_package_version"], "0.94.0"); + assert_eq!(map["stainless_runtime_version"], "v24.3.0"); + assert_eq!(map["user_agent"], "claude-cli/2.1.161 (external, cli)"); + assert_eq!( + fp["transport_profile"]["extra"]["claude_code_identity_profile_version"], + "2026-04" ); } + #[test] + fn sanitize_upgrades_stale_identity_as_one_version_set() { + let raw = serde_json::json!({ + "transport_profile": { + "profile_id": "claude_code_nodejs", + "header_fingerprint": { + "stainless_package_version": "0.68.0", + "stainless_runtime_version": "v20.18.1", + "user_agent": "Mozilla/5.0 stale", + "vscode_session_id": "existing-session" + } + } + }); + + let sanitized = sanitize_fingerprint(&raw, "test-key"); + let map = header_fingerprint_from_fingerprint(&sanitized).expect("header fingerprint"); + assert_eq!(map["stainless_package_version"], "0.94.0"); + assert_eq!(map["stainless_runtime_version"], "v24.3.0"); + assert_eq!(map["user_agent"], "claude-cli/2.1.161 (external, cli)"); + assert_eq!(map["vscode_session_id"], "existing-session"); + } + #[test] fn random_fingerprint_differs_each_call() { let fp1 = generate_random_fingerprint(); diff --git a/crates/aether-provider/transport/src/claude_code/mod.rs b/crates/aether-provider/transport/src/claude_code/mod.rs index 8595121b9..7317f54a0 100644 --- a/crates/aether-provider/transport/src/claude_code/mod.rs +++ b/crates/aether-provider/transport/src/claude_code/mod.rs @@ -1,6 +1,7 @@ mod auth; mod fingerprint; mod policy; +mod profile; mod request; mod url; @@ -13,5 +14,13 @@ pub use policy::{ local_claude_code_transport_unsupported_reason_with_network, supports_local_claude_code_transport_with_network, }; -pub use request::{build_claude_code_passthrough_headers, sanitize_claude_code_request_body}; +pub use profile::{ + current_claude_code_transport_identity_profile, ClaudeCodeBodyCapabilityGate, + ClaudeCodeTransportIdentityProfile, ClaudeCodeTransportIdentityProfileVersion, + CLAUDE_CODE_CONTEXT_MANAGEMENT_BETA, CLAUDE_CODE_TRANSPORT_IDENTITY_2026_04, +}; +pub use request::{ + build_claude_code_passthrough_headers, sanitize_claude_code_request_body, + sanitize_claude_code_request_body_for_beta_header, +}; pub use url::build_claude_code_messages_url; diff --git a/crates/aether-provider/transport/src/claude_code/profile.rs b/crates/aether-provider/transport/src/claude_code/profile.rs new file mode 100644 index 000000000..45d15d340 --- /dev/null +++ b/crates/aether-provider/transport/src/claude_code/profile.rs @@ -0,0 +1,322 @@ +use std::collections::{BTreeMap, BTreeSet}; + +use aether_ai_formats::ApiOperation; + +pub const CLAUDE_CODE_CONTEXT_MANAGEMENT_BETA: &str = "context-management-2025-06-27"; + +const MESSAGE_BETAS_2026_04: &[&str] = &[ + "claude-code-20250219", + "oauth-2025-04-20", + "interleaved-thinking-2025-05-14", + "prompt-caching-scope-2026-01-05", + "effort-2025-11-24", + CLAUDE_CODE_CONTEXT_MANAGEMENT_BETA, + "extended-cache-ttl-2025-04-11", +]; +const COUNT_TOKENS_BETAS_2026_04: &[&str] = &[ + "claude-code-20250219", + "oauth-2025-04-20", + "interleaved-thinking-2025-05-14", + "prompt-caching-scope-2026-01-05", + "effort-2025-11-24", + CLAUDE_CODE_CONTEXT_MANAGEMENT_BETA, + "extended-cache-ttl-2025-04-11", + "token-counting-2024-11-01", +]; +const DROPPED_BETAS_2026_04: &[&str] = &[]; +const BODY_CAPABILITY_GATES_2026_04: &[ClaudeCodeBodyCapabilityGate] = + &[ClaudeCodeBodyCapabilityGate { + body_field: "context_management", + beta_token: CLAUDE_CODE_CONTEXT_MANAGEMENT_BETA, + inject_when_thinking_enabled: true, + default_edit_type: Some("clear_thinking_20251015"), + }]; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ClaudeCodeTransportIdentityProfileVersion { + V2026_04, +} + +impl ClaudeCodeTransportIdentityProfileVersion { + pub const fn as_str(self) -> &'static str { + match self { + Self::V2026_04 => "2026-04", + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ClaudeCodeBodyCapabilityGate { + pub body_field: &'static str, + pub beta_token: &'static str, + pub inject_when_thinking_enabled: bool, + pub default_edit_type: Option<&'static str>, +} + +/// Versioned upstream identity used when Aether is intentionally acting as a +/// Claude Code transport. Native Anthropic transports never resolve this +/// profile and therefore keep their original headers and body untouched. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ClaudeCodeTransportIdentityProfile { + version: ClaudeCodeTransportIdentityProfileVersion, + transport_profile_id: &'static str, + anthropic_version: &'static str, + cli_version: &'static str, + stainless_lang: &'static str, + stainless_package_version: &'static str, + stainless_os: &'static str, + stainless_arch: &'static str, + stainless_runtime: &'static str, + stainless_runtime_version: &'static str, + stainless_retry_count: &'static str, + stainless_timeout: &'static str, + message_required_betas: &'static [&'static str], + count_tokens_required_betas: &'static [&'static str], + preserve_incoming_betas: bool, + dropped_betas: &'static [&'static str], + body_capability_gates: &'static [ClaudeCodeBodyCapabilityGate], +} + +pub const CLAUDE_CODE_TRANSPORT_IDENTITY_2026_04: ClaudeCodeTransportIdentityProfile = + ClaudeCodeTransportIdentityProfile { + version: ClaudeCodeTransportIdentityProfileVersion::V2026_04, + transport_profile_id: "claude_code_nodejs", + anthropic_version: "2023-06-01", + cli_version: "2.1.161", + stainless_lang: "js", + stainless_package_version: "0.94.0", + stainless_os: "Linux", + stainless_arch: "arm64", + stainless_runtime: "node", + stainless_runtime_version: "v24.3.0", + stainless_retry_count: "0", + stainless_timeout: "600", + message_required_betas: MESSAGE_BETAS_2026_04, + count_tokens_required_betas: COUNT_TOKENS_BETAS_2026_04, + preserve_incoming_betas: true, + dropped_betas: DROPPED_BETAS_2026_04, + body_capability_gates: BODY_CAPABILITY_GATES_2026_04, + }; + +pub const fn current_claude_code_transport_identity_profile( +) -> &'static ClaudeCodeTransportIdentityProfile { + &CLAUDE_CODE_TRANSPORT_IDENTITY_2026_04 +} + +impl ClaudeCodeTransportIdentityProfile { + pub const fn version(self) -> ClaudeCodeTransportIdentityProfileVersion { + self.version + } + + pub const fn transport_profile_id(self) -> &'static str { + self.transport_profile_id + } + + pub const fn cli_version(self) -> &'static str { + self.cli_version + } + + pub const fn billing_cli_version(self) -> &'static str { + self.cli_version + } + + pub fn user_agent(self) -> String { + format!("claude-cli/{} (external, cli)", self.cli_version) + } + + pub const fn stainless_package_version(self) -> &'static str { + self.stainless_package_version + } + + pub const fn stainless_lang(self) -> &'static str { + self.stainless_lang + } + + pub const fn stainless_os(self) -> &'static str { + self.stainless_os + } + + pub const fn stainless_arch(self) -> &'static str { + self.stainless_arch + } + + pub const fn stainless_runtime(self) -> &'static str { + self.stainless_runtime + } + + pub const fn stainless_runtime_version(self) -> &'static str { + self.stainless_runtime_version + } + + pub const fn stainless_retry_count(self) -> &'static str { + self.stainless_retry_count + } + + pub const fn stainless_timeout(self) -> &'static str { + self.stainless_timeout + } + + pub fn required_beta_tokens(self, operation: Option) -> &'static [&'static str] { + if operation == Some(ApiOperation::ClaudeCountTokens) { + self.count_tokens_required_betas + } else { + self.message_required_betas + } + } + + pub const fn preserves_incoming_betas(self) -> bool { + self.preserve_incoming_betas + } + + pub const fn dropped_beta_tokens(self) -> &'static [&'static str] { + self.dropped_betas + } + + pub const fn body_capability_gates(self) -> &'static [ClaudeCodeBodyCapabilityGate] { + self.body_capability_gates + } + + pub fn body_capability_gate(self, field: &str) -> Option { + self.body_capability_gates + .iter() + .copied() + .find(|gate| gate.body_field == field) + } + + pub fn apply_fixed_headers(self, headers: &mut BTreeMap, stream: bool) { + for (name, value) in [ + ("accept", "application/json"), + ("anthropic-version", self.anthropic_version), + ("anthropic-dangerous-direct-browser-access", "true"), + ("x-app", "cli"), + ("x-stainless-lang", self.stainless_lang), + ( + "x-stainless-package-version", + self.stainless_package_version, + ), + ("x-stainless-os", self.stainless_os), + ("x-stainless-arch", self.stainless_arch), + ("x-stainless-runtime", self.stainless_runtime), + ( + "x-stainless-runtime-version", + self.stainless_runtime_version, + ), + ("x-stainless-retry-count", self.stainless_retry_count), + ("x-stainless-timeout", self.stainless_timeout), + ] { + headers.insert(name.to_string(), value.to_string()); + } + headers.insert("user-agent".to_string(), self.user_agent()); + if stream { + headers.insert( + "x-stainless-helper-method".to_string(), + "stream".to_string(), + ); + } else { + headers.remove("x-stainless-helper-method"); + } + } + + pub fn apply_beta_policy( + self, + headers: &mut BTreeMap, + operation: Option, + ) { + let incoming = headers.get("anthropic-beta").map(String::as_str); + let merged = self.merge_beta_tokens(incoming, operation); + if merged.is_empty() { + headers.remove("anthropic-beta"); + } else { + headers.insert("anthropic-beta".to_string(), merged); + } + } + + pub fn merge_beta_tokens( + self, + incoming: Option<&str>, + operation: Option, + ) -> String { + let mut seen = BTreeSet::new(); + let mut merged = Vec::new(); + + for token in self.required_beta_tokens(operation) { + self.append_beta_token(&mut seen, &mut merged, token); + } + if self.preserve_incoming_betas { + for token in incoming.unwrap_or_default().split(',') { + self.append_beta_token(&mut seen, &mut merged, token); + } + } + merged.join(",") + } + + pub fn beta_header_enables_body_field(self, beta_header: &str, field: &str) -> bool { + let Some(gate) = self.body_capability_gate(field) else { + return true; + }; + beta_header + .split(',') + .map(str::trim) + .any(|token| token.eq_ignore_ascii_case(gate.beta_token)) + } + + fn append_beta_token(self, seen: &mut BTreeSet, merged: &mut Vec, token: &str) { + let token = token.trim(); + if token.is_empty() + || self + .dropped_betas + .iter() + .any(|dropped| token.eq_ignore_ascii_case(dropped)) + { + return; + } + if seen.insert(token.to_ascii_lowercase()) { + merged.push(token.to_string()); + } + } +} + +#[cfg(test)] +mod tests { + use super::current_claude_code_transport_identity_profile; + use aether_ai_formats::ApiOperation; + + #[test] + fn profile_versions_cli_user_agent_stainless_and_billing_together() { + let profile = *current_claude_code_transport_identity_profile(); + + assert_eq!(profile.version().as_str(), "2026-04"); + assert_eq!(profile.cli_version(), "2.1.161"); + assert_eq!(profile.billing_cli_version(), profile.cli_version()); + assert_eq!( + profile.user_agent(), + format!("claude-cli/{} (external, cli)", profile.cli_version()) + ); + assert_eq!(profile.stainless_package_version(), "0.94.0"); + assert_eq!(profile.stainless_runtime_version(), "v24.3.0"); + } + + #[test] + fn profile_preserves_context_1m_and_adds_operation_specific_betas() { + let profile = *current_claude_code_transport_identity_profile(); + let messages = profile.merge_beta_tokens(Some("context-1m-2025-08-07,custom"), None); + + assert!(messages + .split(',') + .any(|token| token == "context-1m-2025-08-07")); + assert!(messages.split(',').any(|token| token == "custom")); + assert!(!messages + .split(',') + .any(|token| token == "token-counting-2024-11-01")); + + let count_tokens = profile.merge_beta_tokens( + Some("context-1m-2025-08-07"), + Some(ApiOperation::ClaudeCountTokens), + ); + assert!(count_tokens + .split(',') + .any(|token| token == "token-counting-2024-11-01")); + assert!(profile.dropped_beta_tokens().is_empty()); + assert!(profile.preserves_incoming_betas()); + } +} diff --git a/crates/aether-provider/transport/src/claude_code/request.rs b/crates/aether-provider/transport/src/claude_code/request.rs index 8320f6b92..8e07352c7 100644 --- a/crates/aether-provider/transport/src/claude_code/request.rs +++ b/crates/aether-provider/transport/src/claude_code/request.rs @@ -1,31 +1,15 @@ -use std::collections::{BTreeMap, BTreeSet}; +use std::collections::BTreeMap; +use std::sync::OnceLock; +use regex::Regex; use serde_json::{Map, Value}; use super::super::auth::build_openai_passthrough_headers; -use super::fingerprint::header_fingerprint_from_fingerprint; +use super::profile::{ + current_claude_code_transport_identity_profile, ClaudeCodeTransportIdentityProfile, +}; -const DEFAULT_ANTHROPIC_VERSION: &str = "2023-06-01"; -const DEFAULT_ACCEPT: &str = "application/json"; -const STREAM_HELPER_METHOD: &str = "stream"; const DUMMY_THINKING_SIGNATURE: &str = "skip_thought_signature_validator"; -const REQUIRED_BETA_TOKENS: &[&str] = &[ - "claude-code-20250219", - "oauth-2025-04-20", - "interleaved-thinking-2025-05-14", -]; -const EXCLUDED_BETA_TOKENS: &[&str] = &["context-1m-2025-08-07"]; - -/// Fingerprint field -> HTTP header mapping. -/// Every stainless / identity dimension that can vary per-key is listed here. -const FINGERPRINT_HEADER_MAP: &[(&str, &str)] = &[ - ("stainless_package_version", "x-stainless-package-version"), - ("stainless_os", "x-stainless-os"), - ("stainless_arch", "x-stainless-arch"), - ("stainless_runtime_version", "x-stainless-runtime-version"), - ("stainless_timeout", "x-stainless-timeout"), - ("user_agent", "user-agent"), -]; pub fn build_claude_code_passthrough_headers( headers: &http::HeaderMap, @@ -33,7 +17,6 @@ pub fn build_claude_code_passthrough_headers( auth_value: &str, extra_headers: &BTreeMap, stream: bool, - fingerprint: Option<&Value>, ) -> BTreeMap { let mut out = build_openai_passthrough_headers( headers, @@ -43,77 +26,58 @@ pub fn build_claude_code_passthrough_headers( Some("application/json"), ); - // -- Anthropic protocol headers -- - out.insert("accept".to_string(), DEFAULT_ACCEPT.to_string()); - out.insert( - "anthropic-version".to_string(), - DEFAULT_ANTHROPIC_VERSION.to_string(), - ); - // Read incoming anthropic-beta directly from the original HeaderMap because the - // upstream passthrough filter now strips `anthropic-*` headers to avoid leaking - // them to non-Anthropic upstreams. - let incoming_anthropic_beta = headers + // The common passthrough filter intentionally strips Anthropic identity + // headers. Restore only the client beta input; the versioned profile owns + // all fixed identity values and the final beta policy. + let mut incoming_beta_values = headers .get("anthropic-beta") - .and_then(|value| value.to_str().ok()); - out.insert( - "anthropic-beta".to_string(), - merge_anthropic_beta_tokens(incoming_anthropic_beta), + .and_then(|value| value.to_str().ok()) + .map(str::trim) + .filter(|value| !value.is_empty()) + .into_iter() + .map(ToOwned::to_owned) + .collect::>(); + incoming_beta_values.extend( + extra_headers + .iter() + .filter(|(name, _)| name.eq_ignore_ascii_case("anthropic-beta")) + .map(|(_, value)| value.trim()) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned), ); - out.insert( - "anthropic-dangerous-direct-browser-access".to_string(), - "true".to_string(), - ); - out.insert("x-app".to_string(), "cli".to_string()); - - // -- Stainless SDK identity headers -- - // Fixed values: these don't vary per fingerprint. - out.insert("x-stainless-lang".to_string(), "js".to_string()); - out.insert("x-stainless-runtime".to_string(), "node".to_string()); - out.insert("x-stainless-retry-count".to_string(), "0".to_string()); - - // Defaults for fingerprint-overridable fields (used when no fingerprint is present). - out.insert( - "x-stainless-package-version".to_string(), - "0.70.0".to_string(), - ); - out.insert("x-stainless-os".to_string(), "Linux".to_string()); - out.insert("x-stainless-arch".to_string(), "arm64".to_string()); - out.insert( - "x-stainless-runtime-version".to_string(), - "v24.13.0".to_string(), - ); - out.insert("x-stainless-timeout".to_string(), "600".to_string()); - - if stream { - out.insert( - "x-stainless-helper-method".to_string(), - STREAM_HELPER_METHOD.to_string(), - ); - } else { - out.remove("x-stainless-helper-method"); + if !incoming_beta_values.is_empty() { + out.insert("anthropic-beta".to_string(), incoming_beta_values.join(",")); } - // Override from the formal transport profile header fingerprint. - if let Some(fp) = fingerprint.and_then(header_fingerprint_from_fingerprint) { - for &(fp_key, header_key) in FINGERPRINT_HEADER_MAP { - if let Some(value) = fp - .get(fp_key) - .and_then(Value::as_str) - .map(str::trim) - .filter(|v| !v.is_empty()) - { - out.insert(header_key.to_string(), value.to_string()); - } - } - } + let profile = *current_claude_code_transport_identity_profile(); + profile.apply_fixed_headers(&mut out, stream); + profile.apply_beta_policy(&mut out, None); out } pub fn sanitize_claude_code_request_body(body: &mut Value) { + let profile = *current_claude_code_transport_identity_profile(); + let beta_header = profile.merge_beta_tokens(None, None); + sanitize_claude_code_request_body_for_beta_header(body, &beta_header, profile); +} + +pub fn sanitize_claude_code_request_body_for_beta_header( + body: &mut Value, + beta_header: &str, + profile: ClaudeCodeTransportIdentityProfile, +) { let Some(body_object) = body.as_object_mut() else { return; }; + + synchronize_billing_header_version(body_object, profile.billing_cli_version()); + for gate in profile.body_capability_gates() { + if !profile.beta_header_enables_body_field(beta_header, gate.body_field) { + body_object.remove(gate.body_field); + } + } + let thinking_enabled = body_object .get("thinking") .and_then(Value::as_object) @@ -122,6 +86,26 @@ pub fn sanitize_claude_code_request_body(body: &mut Value) { .map(str::trim) .is_some_and(|value| matches!(value.to_ascii_lowercase().as_str(), "enabled" | "adaptive")); + if let Some(gate) = profile.body_capability_gate("context_management") { + if thinking_enabled + && gate.inject_when_thinking_enabled + && profile.beta_header_enables_body_field(beta_header, gate.body_field) + && !body_object.contains_key(gate.body_field) + { + if let Some(default_edit_type) = gate.default_edit_type { + body_object.insert( + gate.body_field.to_string(), + serde_json::json!({ + "edits": [{ + "type": default_edit_type, + "keep": "all" + }] + }), + ); + } + } + } + let Some(messages) = body_object .get_mut("messages") .and_then(Value::as_array_mut) @@ -160,6 +144,37 @@ pub fn sanitize_claude_code_request_body(body: &mut Value) { } } +fn synchronize_billing_header_version(body: &mut Map, cli_version: &str) { + static CC_VERSION: OnceLock = OnceLock::new(); + let Some(system) = body.get_mut("system").and_then(Value::as_array_mut) else { + return; + }; + let regex = CC_VERSION.get_or_init(|| { + Regex::new(r"cc_version=\d+\.\d+\.\d+").expect("billing version regex must compile") + }); + let replacement = format!("cc_version={cli_version}"); + + for block in system { + let Some(block) = block.as_object_mut() else { + continue; + }; + let Some(text) = block + .get("text") + .and_then(Value::as_str) + .map(ToOwned::to_owned) + else { + continue; + }; + if !text.starts_with("x-anthropic-billing-header") { + continue; + } + let updated = regex.replace_all(&text, replacement.as_str()); + if updated != text { + block.insert("text".to_string(), Value::String(updated.into_owned())); + } + } +} + fn keep_claude_code_block( block_object: &Map, role: &str, @@ -187,46 +202,18 @@ fn keep_claude_code_block( true } -fn merge_anthropic_beta_tokens(incoming: Option<&str>) -> String { - let mut seen = BTreeSet::new(); - let mut merged = Vec::new(); - - for token in REQUIRED_BETA_TOKENS { - append_beta_token(&mut seen, &mut merged, token); - } - for token in incoming.unwrap_or_default().split(',') { - let token = token.trim(); - if EXCLUDED_BETA_TOKENS - .iter() - .any(|excluded| token.eq_ignore_ascii_case(excluded)) - { - continue; - } - append_beta_token(&mut seen, &mut merged, token); - } - - merged.join(",") -} - -fn append_beta_token(seen: &mut BTreeSet, merged: &mut Vec, token: &str) { - let normalized = token.trim(); - if normalized.is_empty() { - return; - } - let key = normalized.to_ascii_lowercase(); - if seen.insert(key) { - merged.push(normalized.to_string()); - } -} - #[cfg(test)] mod tests { - use super::{build_claude_code_passthrough_headers, sanitize_claude_code_request_body}; + use super::{ + build_claude_code_passthrough_headers, sanitize_claude_code_request_body, + sanitize_claude_code_request_body_for_beta_header, + }; + use crate::claude_code::current_claude_code_transport_identity_profile; use serde_json::json; use std::collections::BTreeMap; #[test] - fn claude_code_headers_use_transport_profile_header_fingerprint_and_merge_required_betas() { + fn claude_code_headers_use_versioned_identity_and_merge_preserved_betas() { let mut headers = http::HeaderMap::new(); headers.insert( "anthropic-beta", @@ -242,23 +229,12 @@ mod tests { "Bearer upstream-token", &BTreeMap::new(), true, - Some(&json!({ - "transport_profile": { - "profile_id": "claude_code_nodejs", - "header_fingerprint": { - "user_agent":"Claude-Code/9.9", - "stainless_package_version":"1.0.5", - "stainless_runtime_version":"v22.12.0", - "stainless_timeout":"900" - } - } - })), ); assert_eq!( built.get("anthropic-beta").map(String::as_str), Some( - "claude-code-20250219,oauth-2025-04-20,interleaved-thinking-2025-05-14,custom-beta" + "claude-code-20250219,oauth-2025-04-20,interleaved-thinking-2025-05-14,prompt-caching-scope-2026-01-05,effort-2025-11-24,context-management-2025-06-27,extended-cache-ttl-2025-04-11,context-1m-2025-08-07,custom-beta" ) ); assert_eq!( @@ -276,19 +252,19 @@ mod tests { assert_eq!(built.get("x-app").map(String::as_str), Some("cli")); assert_eq!( built.get("x-stainless-package-version").map(String::as_str), - Some("1.0.5") + Some("0.94.0") ); assert_eq!( built.get("x-stainless-runtime-version").map(String::as_str), - Some("v22.12.0") + Some("v24.3.0") ); assert_eq!( built.get("x-stainless-timeout").map(String::as_str), - Some("900") + Some("600") ); assert_eq!( built.get("user-agent").map(String::as_str), - Some("Claude-Code/9.9") + Some("claude-cli/2.1.161 (external, cli)") ); assert_eq!( built.get("authorization").map(String::as_str), @@ -324,4 +300,74 @@ mod tests { ]) ); } + + #[test] + fn context_management_body_is_gated_by_the_matching_beta_token() { + let profile = *current_claude_code_transport_identity_profile(); + let original = json!({ + "context_management": { + "edits": [{"type":"clear_thinking_20251015", "keep":"all"}] + }, + "messages": [] + }); + + let mut without_beta = original.clone(); + sanitize_claude_code_request_body_for_beta_header( + &mut without_beta, + "oauth-2025-04-20", + profile, + ); + assert!(without_beta.get("context_management").is_none()); + + let mut with_beta = original.clone(); + sanitize_claude_code_request_body_for_beta_header( + &mut with_beta, + "oauth-2025-04-20, context-management-2025-06-27", + profile, + ); + assert_eq!(with_beta, original); + } + + #[test] + fn default_profile_keeps_header_and_injected_context_management_in_sync() { + let headers = build_claude_code_passthrough_headers( + &http::HeaderMap::new(), + "authorization", + "Bearer upstream-token", + &BTreeMap::new(), + false, + ); + let mut body = json!({ + "thinking": {"type":"adaptive"}, + "messages": [] + }); + + sanitize_claude_code_request_body(&mut body); + + assert!(headers["anthropic-beta"] + .split(',') + .any(|token| token == "context-management-2025-06-27")); + assert_eq!( + body["context_management"], + json!({"edits":[{"type":"clear_thinking_20251015", "keep":"all"}]}) + ); + } + + #[test] + fn billing_attribution_uses_the_profile_cli_version() { + let mut body = json!({ + "system": [{ + "type":"text", + "text":"x-anthropic-billing-header: cc_version=2.0.0.abc; cc_entrypoint=cli;" + }], + "messages": [] + }); + + sanitize_claude_code_request_body(&mut body); + + assert_eq!( + body["system"][0]["text"], + "x-anthropic-billing-header: cc_version=2.1.161.abc; cc_entrypoint=cli;" + ); + } } diff --git a/crates/aether-provider/transport/src/claude_code/url.rs b/crates/aether-provider/transport/src/claude_code/url.rs index 00e0f8797..ea22367da 100644 --- a/crates/aether-provider/transport/src/claude_code/url.rs +++ b/crates/aether-provider/transport/src/claude_code/url.rs @@ -5,14 +5,19 @@ use url::form_urlencoded; pub fn build_claude_code_messages_url(upstream_base_url: &str, query: Option<&str>) -> String { let (trimmed_base_url, base_query) = split_query(upstream_base_url.trim()); let trimmed_base_url = trimmed_base_url.trim_end_matches('/'); - let mut url = - if trimmed_base_url.ends_with("/v1/messages") || trimmed_base_url.ends_with("/messages") { - trimmed_base_url.to_string() - } else if trimmed_base_url.ends_with("/v1") { - format!("{trimmed_base_url}/messages") - } else { - format!("{trimmed_base_url}/v1/messages") - }; + let mut url = if trimmed_base_url.ends_with("/messages/count_tokens") { + trimmed_base_url + .strip_suffix("/count_tokens") + .unwrap_or(trimmed_base_url) + .to_string() + } else if trimmed_base_url.ends_with("/v1/messages") || trimmed_base_url.ends_with("/messages") + { + trimmed_base_url.to_string() + } else if trimmed_base_url.ends_with("/v1") { + format!("{trimmed_base_url}/messages") + } else { + format!("{trimmed_base_url}/v1/messages") + }; append_merged_query(&mut url, base_query, query); url } @@ -66,6 +71,13 @@ mod tests { build_claude_code_messages_url("https://api.anthropic.com/v1/messages", None), "https://api.anthropic.com/v1/messages" ); + assert_eq!( + build_claude_code_messages_url( + "https://api.anthropic.com/v1/messages/count_tokens?tenant=base", + Some("trace=1"), + ), + "https://api.anthropic.com/v1/messages?tenant=base&trace=1" + ); } #[test] diff --git a/crates/aether-provider/transport/src/generic_oauth/mod.rs b/crates/aether-provider/transport/src/generic_oauth/mod.rs index c7061fcb5..0a24f7ff5 100644 --- a/crates/aether-provider/transport/src/generic_oauth/mod.rs +++ b/crates/aether-provider/transport/src/generic_oauth/mod.rs @@ -26,6 +26,42 @@ pub fn supports_local_generic_oauth_request_auth_resolution( && generic_provider_type(transport.provider.provider_type.as_str()).is_some() } +pub fn resolve_local_generic_oauth_transport_authorization( + transport: &GatewayProviderTransportSnapshot, +) -> Option { + if !supports_local_generic_oauth_request_auth_resolution(transport) { + return None; + } + if let Some(value) = + auth_config_authorization_header(transport.key.decrypted_auth_config.as_deref()) + { + return if bearer_access_token(&value).is_some() { + Some(value) + } else { + None + }; + } + + let auth_config = GenericOAuthRefreshAdapter::auth_config_from_transport(transport); + let refreshable = auth_config + .as_ref() + .and_then(refresh_token_from_auth_config) + .is_some(); + if refreshable && auth_config_expires_soon(auth_config.as_ref()) { + return None; + } + + let secret = transport.key.decrypted_api_key.trim(); + if !secret.is_empty() && secret != PLACEHOLDER_API_KEY { + return Some(format!("Bearer {secret}")); + } + + auth_config + .as_ref() + .and_then(access_token_from_auth_config) + .map(|token| format!("Bearer {token}")) +} + #[derive(Debug, Clone, Default)] pub struct GenericOAuthRefreshAdapter { token_url_overrides: BTreeMap, @@ -114,48 +150,23 @@ impl GenericOAuthRefreshAdapter { &self, transport: &GatewayProviderTransportSnapshot, ) -> Option { - if !supports_local_generic_oauth_request_auth_resolution(transport) { - return None; - } - - if let Some(value) = - auth_config_authorization_header(transport.key.decrypted_auth_config.as_deref()) - { - return Some(LocalResolvedOAuthRequestAuth::Header { - name: AUTH_HEADER_NAME.to_string(), - value, - }); - } - - let secret = transport.key.decrypted_api_key.trim(); - if secret.is_empty() || secret == PLACEHOLDER_API_KEY { - return None; - } - - let auth_config = Self::auth_config_from_transport(transport); - let refreshable = auth_config - .as_ref() - .and_then(refresh_token_from_auth_config) - .is_some(); - if refreshable && auth_config_expires_soon(auth_config.as_ref()) { - return None; - } - Some(LocalResolvedOAuthRequestAuth::Header { name: AUTH_HEADER_NAME.to_string(), - value: format!("Bearer {secret}"), + value: resolve_local_generic_oauth_transport_authorization(transport)?, }) } fn build_cached_entry( provider_type: &'static str, transport: &GatewayProviderTransportSnapshot, - refreshed: ProviderOAuthTokenSet, + mut refreshed: ProviderOAuthTokenSet, ) -> CachedOAuthEntry { + let auth_header_value = refreshed.token_set.bearer_header_value(); + synchronize_authorization_overrides(&mut refreshed.auth_config, &auth_header_value); CachedOAuthEntry { provider_type: provider_type.to_string(), auth_header_name: AUTH_HEADER_NAME.to_string(), - auth_header_value: refreshed.token_set.bearer_header_value(), + auth_header_value, expires_at_unix_secs: refreshed.token_set.expires_at_unix_secs, metadata: Some(refreshed.auth_config), source_fingerprint: Some(generic_oauth_transport_source_fingerprint(transport)), @@ -184,12 +195,14 @@ impl LocalOAuthRefreshAdapter for GenericOAuthRefreshAdapter { { return None; } - if let Some(value) = - auth_config_authorization_header(transport.key.decrypted_auth_config.as_deref()) + if auth_config_authorization_header(transport.key.decrypted_auth_config.as_deref()) + .is_some() { - return Some(LocalResolvedOAuthRequestAuth::Header { - name: AUTH_HEADER_NAME.to_string(), - value, + return resolve_local_generic_oauth_transport_authorization(transport).map(|value| { + LocalResolvedOAuthRequestAuth::Header { + name: AUTH_HEADER_NAME.to_string(), + value, + } }); } if !generic_oauth_cached_entry_matches_transport(transport, entry) { @@ -211,6 +224,26 @@ impl LocalOAuthRefreshAdapter for GenericOAuthRefreshAdapter { }) } + fn resolve_fenced_cached( + &self, + transport: &GatewayProviderTransportSnapshot, + entry: &CachedOAuthEntry, + ) -> Option { + generic_oauth_successor_entry_matches_transport(transport, entry) + .then(|| resolved_entry_header(entry)) + .flatten() + } + + fn resolve_refreshed( + &self, + transport: &GatewayProviderTransportSnapshot, + entry: &CachedOAuthEntry, + ) -> Option { + generic_oauth_entry_belongs_to_transport(transport, entry) + .then(|| resolved_entry_header(entry)) + .flatten() + } + fn resolve_without_refresh( &self, transport: &GatewayProviderTransportSnapshot, @@ -240,6 +273,39 @@ impl LocalOAuthRefreshAdapter for GenericOAuthRefreshAdapter { .is_some() } + fn refresh_fingerprint( + &self, + transport: &GatewayProviderTransportSnapshot, + entry: Option<&CachedOAuthEntry>, + ) -> Option { + generic_oauth_refresh_fingerprint(transport, entry) + } + + fn cached_entry_from_transport( + &self, + transport: &GatewayProviderTransportSnapshot, + ) -> Option { + let provider_type = generic_provider_type(transport.provider.provider_type.as_str())?; + let LocalResolvedOAuthRequestAuth::Header { name, value } = + self.resolve_direct_header(transport)? + else { + return None; + }; + let metadata = Self::auth_config_from_transport(transport); + let expires_at_unix_secs = transport + .key + .expires_at_unix_secs + .or_else(|| metadata.as_ref().and_then(auth_config_expires_at)); + Some(CachedOAuthEntry { + provider_type: provider_type.to_string(), + auth_header_name: name, + auth_header_value: value, + expires_at_unix_secs, + metadata, + source_fingerprint: Some(generic_oauth_transport_source_fingerprint(transport)), + }) + } + async fn refresh( &self, executor: &dyn LocalOAuthHttpExecutor, @@ -311,20 +377,33 @@ impl LocalOAuthRefreshAdapter for GenericOAuthRefreshAdapter { fn generic_oauth_transport_source_fingerprint( transport: &GatewayProviderTransportSnapshot, ) -> String { - let provider_type = transport.provider.provider_type.trim().to_ascii_lowercase(); - let auth_type = transport.key.auth_type.trim().to_ascii_lowercase(); let auth_config = transport .key .decrypted_auth_config .as_deref() .unwrap_or_default(); - let api_key = transport.key.decrypted_api_key.as_str(); + generic_oauth_credential_fingerprint( + transport.provider.provider_type.as_str(), + transport.key.auth_type.as_str(), + auth_config, + transport.key.decrypted_api_key.as_str(), + ) +} + +fn generic_oauth_credential_fingerprint( + provider_type: &str, + auth_type: &str, + auth_config: &str, + access_token: &str, +) -> String { + let provider_type = provider_type.trim().to_ascii_lowercase(); + let auth_type = auth_type.trim().to_ascii_lowercase(); let mut digest = Sha256::new(); for field in [ provider_type.as_bytes(), auth_type.as_bytes(), auth_config.as_bytes(), - api_key.as_bytes(), + access_token.as_bytes(), ] { digest.update((field.len() as u64).to_be_bytes()); digest.update(field); @@ -340,6 +419,91 @@ fn generic_oauth_cached_entry_matches_transport( entry.source_fingerprint.as_deref() == Some(transport_fingerprint.as_str()) } +fn generic_oauth_entry_belongs_to_transport( + transport: &GatewayProviderTransportSnapshot, + entry: &CachedOAuthEntry, +) -> bool { + entry + .provider_type + .eq_ignore_ascii_case(transport.provider.provider_type.as_str()) + && generic_oauth_cached_entry_matches_transport(transport, entry) +} + +fn generic_oauth_successor_entry_matches_transport( + transport: &GatewayProviderTransportSnapshot, + entry: &CachedOAuthEntry, +) -> bool { + generic_oauth_entry_belongs_to_transport(transport, entry) + && !expires_at_requires_refresh(entry.expires_at_unix_secs) + && resolved_entry_header(entry).is_some() + && entry.metadata.is_some() +} + +fn generic_oauth_refresh_fingerprint( + transport: &GatewayProviderTransportSnapshot, + entry: Option<&CachedOAuthEntry>, +) -> Option { + if !supports_local_generic_oauth_request_auth_resolution(transport) { + return None; + } + entry + .filter(|entry| generic_oauth_successor_entry_matches_transport(transport, entry)) + .and_then(|entry| { + let metadata = serde_json::to_string(entry.metadata.as_ref()?).ok()?; + let access_token = bearer_access_token(entry.auth_header_value.as_str())?; + Some(generic_oauth_credential_fingerprint( + transport.provider.provider_type.as_str(), + transport.key.auth_type.as_str(), + metadata.as_str(), + access_token, + )) + }) + .or_else(|| Some(generic_oauth_transport_source_fingerprint(transport))) +} + +fn resolved_entry_header(entry: &CachedOAuthEntry) -> Option { + let name = entry.auth_header_name.trim(); + let value = entry.auth_header_value.trim(); + if name.is_empty() || value.is_empty() { + return None; + } + Some(LocalResolvedOAuthRequestAuth::Header { + name: name.to_ascii_lowercase(), + value: value.to_string(), + }) +} + +fn bearer_access_token(authorization: &str) -> Option<&str> { + let mut parts = authorization.split_ascii_whitespace(); + let scheme = parts.next()?; + let token = parts.next()?; + (scheme.eq_ignore_ascii_case("bearer") && parts.next().is_none()).then_some(token) +} + +fn synchronize_authorization_overrides(auth_config: &mut Value, authorization: &str) { + let Some(object) = auth_config.as_object_mut() else { + return; + }; + for (key, value) in object.iter_mut() { + match key.trim().to_ascii_lowercase().as_str() { + "headers" | "extra_headers" | "extraheaders" => { + let Some(headers) = value.as_object_mut() else { + continue; + }; + for (header_name, header_value) in headers.iter_mut() { + if header_name.trim().eq_ignore_ascii_case(AUTH_HEADER_NAME) { + *header_value = Value::String(authorization.to_string()); + } + } + } + "transport" | "request" => { + synchronize_authorization_overrides(value, authorization); + } + _ => {} + } + } +} + fn generic_provider_type(provider_type: &str) -> Option<&'static str> { let normalized = provider_type.trim(); GENERIC_PROVIDER_OAUTH_TEMPLATES @@ -355,6 +519,13 @@ fn refresh_token_from_auth_config(auth_config: &Value) -> Option { .and_then(non_empty_string) } +fn access_token_from_auth_config(auth_config: &Value) -> Option { + let object = auth_config.as_object()?; + ["access_token", "accessToken"] + .iter() + .find_map(|field| object.get(*field).and_then(non_empty_string)) +} + fn auth_config_expires_at(auth_config: &Value) -> Option { auth_config .as_object() @@ -423,10 +594,16 @@ fn current_access_token( #[cfg(test)] mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::Arc; + + use async_trait::async_trait; use serde_json::json; use super::super::oauth_refresh::{ - CachedOAuthEntry, LocalOAuthRefreshAdapter, LocalResolvedOAuthRequestAuth, + CachedOAuthEntry, LocalOAuthHttpExecutor, LocalOAuthHttpRequest, LocalOAuthHttpResponse, + LocalOAuthRefreshAdapter, LocalOAuthRefreshCoordinator, LocalOAuthRefreshError, + LocalResolvedOAuthRequestAuth, }; use super::super::snapshot::{ GatewayProviderTransportEndpoint, GatewayProviderTransportKey, @@ -434,9 +611,35 @@ mod tests { }; use super::{ current_access_token, generic_oauth_transport_source_fingerprint, - GenericOAuthRefreshAdapter, + resolve_local_generic_oauth_transport_authorization, GenericOAuthRefreshAdapter, }; + #[derive(Debug)] + struct StaticTokenExecutor { + hits: Arc, + } + + #[async_trait] + impl LocalOAuthHttpExecutor for StaticTokenExecutor { + async fn execute( + &self, + _provider_type: &'static str, + _transport: &GatewayProviderTransportSnapshot, + _request: &LocalOAuthHttpRequest, + ) -> Result { + self.hits.fetch_add(1, Ordering::SeqCst); + Ok(LocalOAuthHttpResponse { + status_code: 200, + body_text: json!({ + "access_token": "fresh-access-token", + "expires_in": 3600, + "token_type": "Bearer" + }) + .to_string(), + }) + } + } + fn sample_transport() -> GatewayProviderTransportSnapshot { GatewayProviderTransportSnapshot { provider: GatewayProviderTransportProvider { @@ -519,6 +722,58 @@ mod tests { ); } + #[test] + fn transport_authorization_uses_one_effective_bearer_generation() { + let mut transport = sample_transport(); + transport.key.decrypted_api_key = "api-access-token".to_string(); + transport.key.decrypted_auth_config = Some( + json!({ + "accessToken": "legacy-access-token", + "request": { + "extraHeaders": { + "Authorization": "Bearer nested-override-token" + } + } + }) + .to_string(), + ); + assert_eq!( + resolve_local_generic_oauth_transport_authorization(&transport).as_deref(), + Some("Bearer nested-override-token") + ); + + transport.key.decrypted_auth_config = + Some(json!({"accessToken": "legacy-access-token"}).to_string()); + assert_eq!( + resolve_local_generic_oauth_transport_authorization(&transport).as_deref(), + Some("Bearer api-access-token") + ); + + transport.key.decrypted_api_key = "__placeholder__".to_string(); + assert_eq!( + resolve_local_generic_oauth_transport_authorization(&transport).as_deref(), + Some("Bearer legacy-access-token") + ); + } + + #[test] + fn non_bearer_authorization_override_does_not_fall_back_to_stale_token() { + let mut transport = sample_transport(); + transport.key.decrypted_api_key = "api-access-token".to_string(); + transport.key.decrypted_auth_config = Some( + json!({ + "access_token": "legacy-access-token", + "headers": {"Authorization": "Basic imported-session"} + }) + .to_string(), + ); + + assert!(resolve_local_generic_oauth_transport_authorization(&transport).is_none()); + assert!(GenericOAuthRefreshAdapter::default() + .resolve_without_refresh(&transport) + .is_none()); + } + #[test] fn auth_config_authorization_header_overrides_cached_oauth_entry() { let adapter = GenericOAuthRefreshAdapter::default(); @@ -636,4 +891,95 @@ mod tests { Some("refreshed-access-a") ); } + + #[tokio::test] + async fn fenced_force_reuses_successor_when_refresh_token_does_not_rotate() { + let mut transport = sample_transport(); + transport.key.decrypted_api_key = "stale-access-token".to_string(); + transport.key.expires_at_unix_secs = Some(u64::MAX); + transport.key.decrypted_auth_config = Some( + json!({ + "provider_type": "codex", + "refresh_token": "stable-refresh-token", + "expires_at": u64::MAX, + "headers": {"Authorization": "Bearer stale-top-level"}, + "request": { + "extraHeaders": {"authorization": "Bearer stale-request"} + }, + "transport": { + "extra_headers": {"AUTHORIZATION": "Bearer stale-transport"} + } + }) + .to_string(), + ); + let hits = Arc::new(AtomicUsize::new(0)); + let coordinator = LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![Arc::new( + GenericOAuthRefreshAdapter::default() + .with_token_url_for_tests("codex", "https://oauth.example/token"), + )]); + let executor = StaticTokenExecutor { + hits: Arc::clone(&hits), + }; + let expected = coordinator + .refresh_fingerprint_for_transport(&transport) + .expect("refreshable transport should have a generation fence"); + + let first = coordinator + .force_refresh_with_result_fenced( + &executor, + &transport, + None, + None, + Some(expected.as_str()), + ) + .await + .expect("first refresh should succeed") + .expect("first refresh should resolve"); + let first_entry = first + .refreshed_entry + .as_ref() + .expect("first refresh should return a cache entry"); + assert_eq!(first_entry.auth_header_value, "Bearer fresh-access-token"); + let metadata = first_entry + .metadata + .as_ref() + .expect("generic refresh should preserve auth metadata"); + assert_eq!(metadata["refresh_token"], "stable-refresh-token"); + assert_eq!( + metadata["headers"]["Authorization"], + "Bearer fresh-access-token" + ); + assert_eq!( + metadata["request"]["extraHeaders"]["authorization"], + "Bearer fresh-access-token" + ); + assert_eq!( + metadata["transport"]["extra_headers"]["AUTHORIZATION"], + "Bearer fresh-access-token" + ); + coordinator + .store_cached_entry(&transport.key.id, first_entry.clone()) + .await; + + let follower = coordinator + .force_refresh_with_result_fenced( + &executor, + &transport, + None, + None, + Some(expected.as_str()), + ) + .await + .expect("follower should reuse the completed refresh") + .expect("follower should resolve"); + assert!(follower.reused_refresh); + assert_eq!( + follower + .refreshed_entry + .as_ref() + .map(|entry| entry.auth_header_value.as_str()), + Some("Bearer fresh-access-token") + ); + assert_eq!(hits.load(Ordering::SeqCst), 1); + } } diff --git a/crates/aether-provider/transport/src/headers.rs b/crates/aether-provider/transport/src/headers.rs index e3d105b34..a88e27469 100644 --- a/crates/aether-provider/transport/src/headers.rs +++ b/crates/aether-provider/transport/src/headers.rs @@ -2,8 +2,37 @@ use std::collections::BTreeMap; use aether_contracts::USAGE_SERVER_NOW_UNIX_MS_HEADER; +const UPSTREAM_CREDENTIAL_HEADER_NAMES: &[&str] = &[ + "authorization", + "proxy-authorization", + "api-key", + "x-api-key", + "x-goog-api-key", + "cookie", + "cookie2", + "set-cookie", +]; + +pub(crate) fn upstream_credential_header_names() -> &'static [&'static str] { + UPSTREAM_CREDENTIAL_HEADER_NAMES +} + +pub(crate) fn is_upstream_credential_header(name: &str) -> bool { + let name = name.trim(); + UPSTREAM_CREDENTIAL_HEADER_NAMES + .iter() + .any(|candidate| name.eq_ignore_ascii_case(candidate)) +} + +pub(crate) fn is_aether_internal_header(name: &str) -> bool { + name.trim().to_ascii_lowercase().starts_with("x-aether-") +} + pub fn should_skip_request_header(name: &str) -> bool { let normalized = name.to_ascii_lowercase(); + if is_aether_internal_header(&normalized) { + return true; + } matches!( normalized.as_str(), "connection" @@ -15,11 +44,7 @@ pub fn should_skip_request_header(name: &str) -> bool { | "trailer" | "transfer-encoding" | "upgrade" - | "x-aether-execution-path" - | "x-aether-dependency-reason" - | "x-aether-execution-loop-guard" - | "x-aether-control-execute-fallback" - | "x-aether-rate-limit-preflight" + | "set-cookie" | USAGE_SERVER_NOW_UNIX_MS_HEADER ) } @@ -35,12 +60,10 @@ pub fn should_skip_upstream_passthrough_header(name: &str) -> bool { if lower.starts_with("x-stainless-") || lower.starts_with("anthropic-") { return true; } - matches!( - lower.as_str(), - "authorization" - | "x-api-key" - | "x-goog-api-key" - | "host" + is_upstream_credential_header(&lower) + || matches!( + lower.as_str(), + "host" | "content-length" | "transfer-encoding" | "connection" @@ -55,29 +78,29 @@ pub fn should_skip_upstream_passthrough_header(name: &str) -> bool { // Claude CLI client identifier; re-injected by the Claude Code adapter // when the upstream is Anthropic, filtered for everybody else. | "x-app" - ) || should_skip_request_header(name) + ) + || should_skip_request_header(name) } pub(crate) fn should_skip_upstream_complete_passthrough_header(name: &str) -> bool { let lower = name.to_ascii_lowercase(); - matches!( - lower.as_str(), - "authorization" - | "x-api-key" - | "x-goog-api-key" - | "host" - | "content-length" - | "transfer-encoding" - | "connection" - | "content-encoding" - | "x-real-ip" - | "x-real-proto" - | "x-forwarded-for" - | "x-forwarded-proto" - | "x-forwarded-scheme" - | "x-forwarded-host" - | "x-forwarded-port" - ) || should_skip_request_header(name) + is_upstream_credential_header(&lower) + || matches!( + lower.as_str(), + "host" + | "content-length" + | "transfer-encoding" + | "connection" + | "content-encoding" + | "x-real-ip" + | "x-real-proto" + | "x-forwarded-for" + | "x-forwarded-proto" + | "x-forwarded-scheme" + | "x-forwarded-host" + | "x-forwarded-port" + ) + || should_skip_request_header(name) } pub fn normalize_upstream_accept_encoding(value: &str) -> Option { @@ -170,9 +193,9 @@ pub fn force_identity_accept_encoding(headers: &mut BTreeMap) { #[cfg(test)] mod tests { use super::{ - force_identity_accept_encoding, normalize_upstream_accept_encoding, - should_skip_request_header, should_skip_upstream_complete_passthrough_header, - should_skip_upstream_passthrough_header, + force_identity_accept_encoding, is_upstream_credential_header, + normalize_upstream_accept_encoding, should_skip_request_header, + should_skip_upstream_complete_passthrough_header, should_skip_upstream_passthrough_header, }; use aether_contracts::USAGE_SERVER_NOW_UNIX_MS_HEADER; use std::collections::BTreeMap; @@ -273,6 +296,47 @@ mod tests { } } + #[test] + fn strips_all_client_credential_carriers_from_passthrough() { + for header in [ + "authorization", + "proxy-authorization", + "api-key", + "x-api-key", + "x-goog-api-key", + "cookie", + "cookie2", + "set-cookie", + "Authorization", + "COOKIE", + ] { + assert!(is_upstream_credential_header(header), "credential {header}"); + assert!( + should_skip_upstream_passthrough_header(header), + "normal passthrough should strip {header}" + ); + assert!( + should_skip_upstream_complete_passthrough_header(header), + "complete passthrough should strip {header}" + ); + } + } + + #[test] + fn strips_all_aether_owned_headers_from_provider_requests() { + for header in [ + "x-aether-gateway", + "x-aether-auth-user-id", + "x-aether-auth-api-key-id", + "x-aether-auth-balance-remaining", + "X-Aether-Tunnel-Forwarded-By", + ] { + assert!(should_skip_request_header(header)); + assert!(should_skip_upstream_passthrough_header(header)); + assert!(should_skip_upstream_complete_passthrough_header(header)); + } + } + #[test] fn strips_usage_server_time_header_from_provider_requests() { for h in [ diff --git a/crates/aether-provider/transport/src/kiro/refresh.rs b/crates/aether-provider/transport/src/kiro/refresh.rs index 2fa9cb6c6..df10213d8 100644 --- a/crates/aether-provider/transport/src/kiro/refresh.rs +++ b/crates/aether-provider/transport/src/kiro/refresh.rs @@ -1,5 +1,6 @@ use aether_oauth::provider::providers::KiroProviderOAuthAdapter as CoreKiroProviderOAuthAdapter; use async_trait::async_trait; +use sha2::{Digest, Sha256}; use super::super::oauth_refresh::{ oauth_error_to_local_refresh_error, provider_oauth_transport_context_from_snapshot, @@ -51,11 +52,15 @@ impl KiroOAuthRefreshAdapter { .map_err(|error| oauth_error_to_local_refresh_error(PROVIDER_TYPE, error)) } - fn auth_config_from_entry(entry: &CachedOAuthEntry) -> Option { + fn auth_config_from_entry( + transport: &GatewayProviderTransportSnapshot, + entry: &CachedOAuthEntry, + ) -> Option { entry .metadata .as_ref() .filter(|_| entry.provider_type.eq_ignore_ascii_case(PROVIDER_TYPE)) + .filter(|_| kiro_cached_entry_matches_transport(transport, entry)) .and_then(KiroAuthConfig::from_json_value) } @@ -64,12 +69,17 @@ impl KiroOAuthRefreshAdapter { transport: &GatewayProviderTransportSnapshot, entry: Option<&CachedOAuthEntry>, ) -> Option { - entry.and_then(Self::auth_config_from_entry).or_else(|| { - KiroAuthConfig::from_raw_json(transport.key.decrypted_auth_config.as_deref()) - }) + entry + .and_then(|entry| Self::auth_config_from_entry(transport, entry)) + .or_else(|| { + KiroAuthConfig::from_raw_json(transport.key.decrypted_auth_config.as_deref()) + }) } - fn build_cached_entry(auth_config: &KiroAuthConfig) -> Option { + fn build_cached_entry( + transport: &GatewayProviderTransportSnapshot, + auth_config: &KiroAuthConfig, + ) -> Option { let request_auth = build_kiro_request_auth_from_config(auth_config.clone(), None)?; Some(CachedOAuthEntry { provider_type: PROVIDER_TYPE.to_string(), @@ -77,7 +87,21 @@ impl KiroOAuthRefreshAdapter { auth_header_value: request_auth.value, expires_at_unix_secs: auth_config.expires_at, metadata: Some(auth_config.to_json_value()), - source_fingerprint: None, + source_fingerprint: Some(kiro_transport_credential_fingerprint(transport)), + }) + } + + fn build_cached_entry_from_transport( + transport: &GatewayProviderTransportSnapshot, + ) -> Option { + let request_auth = resolve_local_kiro_request_auth(transport)?; + Some(CachedOAuthEntry { + provider_type: PROVIDER_TYPE.to_string(), + auth_header_name: request_auth.name.to_string(), + auth_header_value: request_auth.value, + expires_at_unix_secs: request_auth.auth_config.expires_at, + metadata: Some(request_auth.auth_config.to_json_value()), + source_fingerprint: Some(kiro_transport_credential_fingerprint(transport)), }) } @@ -101,10 +125,10 @@ impl LocalOAuthRefreshAdapter for KiroOAuthRefreshAdapter { fn resolve_cached( &self, - _transport: &GatewayProviderTransportSnapshot, + transport: &GatewayProviderTransportSnapshot, entry: &CachedOAuthEntry, ) -> Option { - let auth_config = Self::auth_config_from_entry(entry)?; + let auth_config = Self::auth_config_from_entry(transport, entry)?; let request_auth = build_kiro_request_auth_from_config(auth_config, None)?; Some(LocalResolvedOAuthRequestAuth::Kiro(request_auth)) } @@ -128,6 +152,24 @@ impl LocalOAuthRefreshAdapter for KiroOAuthRefreshAdapter { && self.refreshable_auth_config(transport, entry).is_some() } + fn refresh_fingerprint( + &self, + transport: &GatewayProviderTransportSnapshot, + entry: Option<&CachedOAuthEntry>, + ) -> Option { + self.supports(transport).then(|| { + kiro_successor_refresh_fingerprint(transport, entry) + .unwrap_or_else(|| kiro_transport_credential_fingerprint(transport)) + }) + } + + fn cached_entry_from_transport( + &self, + transport: &GatewayProviderTransportSnapshot, + ) -> Option { + Self::build_cached_entry_from_transport(transport) + } + async fn refresh( &self, executor: &dyn LocalOAuthHttpExecutor, @@ -140,10 +182,75 @@ impl LocalOAuthRefreshAdapter for KiroOAuthRefreshAdapter { let refreshed = self .refresh_auth_config(executor, transport, &auth_config) .await?; - Ok(Self::build_cached_entry(&refreshed)) + Ok(Self::build_cached_entry(transport, &refreshed)) } } +fn kiro_transport_credential_fingerprint(transport: &GatewayProviderTransportSnapshot) -> String { + kiro_credential_fingerprint( + transport.provider.provider_type.as_str(), + transport.key.auth_type.as_str(), + transport + .key + .decrypted_auth_config + .as_deref() + .unwrap_or_default(), + transport.key.decrypted_api_key.as_str(), + ) +} + +fn kiro_successor_refresh_fingerprint( + transport: &GatewayProviderTransportSnapshot, + entry: Option<&CachedOAuthEntry>, +) -> Option { + let entry = entry.filter(|entry| kiro_cached_entry_matches_transport(transport, entry))?; + let metadata = serde_json::to_string(entry.metadata.as_ref()?).ok()?; + let access_token = bearer_access_token(entry.auth_header_value.as_str())?; + Some(kiro_credential_fingerprint( + transport.provider.provider_type.as_str(), + transport.key.auth_type.as_str(), + metadata.as_str(), + access_token, + )) +} + +fn kiro_cached_entry_matches_transport( + transport: &GatewayProviderTransportSnapshot, + entry: &CachedOAuthEntry, +) -> bool { + entry.provider_type.eq_ignore_ascii_case(PROVIDER_TYPE) + && entry.source_fingerprint.as_deref() + == Some(kiro_transport_credential_fingerprint(transport).as_str()) +} + +fn kiro_credential_fingerprint( + provider_type: &str, + auth_type: &str, + auth_config: &str, + access_token: &str, +) -> String { + let provider_type = provider_type.trim().to_ascii_lowercase(); + let auth_type = auth_type.trim().to_ascii_lowercase(); + let mut digest = Sha256::new(); + for field in [ + provider_type.as_bytes(), + auth_type.as_bytes(), + auth_config.as_bytes(), + access_token.as_bytes(), + ] { + digest.update((field.len() as u64).to_be_bytes()); + digest.update(field); + } + format!("{:x}", digest.finalize()) +} + +fn bearer_access_token(authorization: &str) -> Option<&str> { + let mut parts = authorization.split_ascii_whitespace(); + let scheme = parts.next()?; + let token = parts.next()?; + (scheme.eq_ignore_ascii_case("bearer") && parts.next().is_none()).then_some(token) +} + #[cfg(test)] mod tests { use std::sync::{Arc, Mutex}; @@ -155,7 +262,7 @@ mod tests { GatewayProviderTransportEndpoint, GatewayProviderTransportKey, GatewayProviderTransportProvider, GatewayProviderTransportSnapshot, }; - use super::{KiroOAuthRefreshAdapter, IDC_AMZ_USER_AGENT}; + use super::{KiroAuthConfig, KiroOAuthRefreshAdapter, IDC_AMZ_USER_AGENT}; use axum::body::to_bytes; use axum::extract::Request; use axum::response::IntoResponse; @@ -244,6 +351,72 @@ mod tests { (format!("http://{addr}"), handle) } + #[test] + fn cached_entry_is_bound_to_kiro_credential_generation() { + let stable_refresh_token = "r".repeat(120); + let source_transport = sample_transport( + &json!({ + "refresh_token": stable_refresh_token, + "machine_id": "123e4567-e89b-12d3-a456-426614174000", + "kiro_version": "1.2.3" + }) + .to_string(), + ); + let refreshed_config = KiroAuthConfig::from_raw_json(Some( + &json!({ + "refresh_token": "s".repeat(120), + "access_token": "fresh-kiro-access-token", + "expires_at": u64::MAX, + "machine_id": "123e4567-e89b-12d3-a456-426614174000", + "kiro_version": "1.2.3" + }) + .to_string(), + )) + .expect("refreshed Kiro config should parse"); + let entry = + KiroOAuthRefreshAdapter::build_cached_entry(&source_transport, &refreshed_config) + .expect("refreshed Kiro entry should build"); + let adapter = KiroOAuthRefreshAdapter::default(); + + assert!(adapter.resolve_cached(&source_transport, &entry).is_some()); + assert_ne!( + adapter.refresh_fingerprint(&source_transport, None), + adapter.refresh_fingerprint(&source_transport, Some(&entry)) + ); + + let replacement_transport = sample_transport( + &json!({ + "refresh_token": "admin-refresh-token", + "access_token": "admin-access-token", + "expires_at": u64::MAX, + "machine_id": "123e4567-e89b-12d3-a456-426614174001", + "kiro_version": "1.2.3" + }) + .to_string(), + ); + assert!(adapter + .resolve_cached(&replacement_transport, &entry) + .is_none()); + let selected_refresh_config = adapter + .base_auth_config(&replacement_transport, Some(&entry)) + .expect("replacement transport should provide refresh config"); + assert_eq!( + selected_refresh_config.refresh_token.as_deref(), + Some("admin-refresh-token") + ); + assert_eq!( + selected_refresh_config.access_token.as_deref(), + Some("admin-access-token") + ); + + let replacement_entry = + LocalOAuthRefreshAdapter::cached_entry_from_transport(&adapter, &replacement_transport) + .expect("persisted replacement should reconstruct a generation-bound entry"); + assert!(adapter + .resolve_cached(&replacement_transport, &replacement_entry) + .is_some()); + } + #[tokio::test] async fn refreshes_social_token_via_adapter() { let seen_request = Arc::new(Mutex::new(None::)); diff --git a/crates/aether-provider/transport/src/lib.rs b/crates/aether-provider/transport/src/lib.rs index 6fcda7bb5..187d18e7e 100644 --- a/crates/aether-provider/transport/src/lib.rs +++ b/crates/aether-provider/transport/src/lib.rs @@ -1,4 +1,5 @@ mod agent_identity; +mod anthropic_compat; pub mod antigravity; pub mod auth; mod auth_config; @@ -46,6 +47,10 @@ pub use agent_identity::{ CODEX_AGENT_IDENTITY_AUTH_MODE, CODEX_AGENT_IDENTITY_CACHED_ENTRY_PROVIDER_TYPE, CODEX_AGENT_IDENTITY_PROVIDER_TYPE, CODEX_AGENT_IDENTITY_TASK_REGISTRATION_REQUEST_ID, }; +pub use anthropic_compat::{ + resolve_anthropic_compatibility_profile, validate_anthropic_compatibility_profile_config, + AnthropicCompatibilityProfile, AnthropicCompatibilityProfileConfigError, +}; pub use auth::{build_passthrough_headers, ensure_upstream_auth_header}; pub use auth_config::apply_local_auth_config_header_overrides; pub use cache::{provider_transport_snapshot_looks_refreshed, ProviderTransportSnapshotCacheKey}; @@ -76,6 +81,7 @@ pub use gemini_files::{ GeminiFilesRequestBodyError, GeminiFilesRequestBodyParts, }; pub use generic_oauth::{ + resolve_local_generic_oauth_transport_authorization, supports_local_generic_oauth_request_auth_resolution, GenericOAuthRefreshAdapter, }; pub use grok::{ @@ -122,7 +128,7 @@ pub use request_url::{ build_kiro_cross_format_upstream_url, build_local_openai_chat_upstream_url, build_local_openai_responses_upstream_url, build_transport_request_url, build_transport_request_url_for_request_body, gemini_embedding_request_body_uses_batch, - TransportRequestUrlParams, + transport_supports_api_operation, TransportRequestUrlParams, }; pub use rules::{ apply_local_body_rules, apply_local_body_rules_with_request_headers, apply_local_header_rules, @@ -134,6 +140,8 @@ pub use same_format_provider::{ build_same_format_provider_headers, build_same_format_provider_request_body, build_same_format_provider_request_body_with_compatibility_report, build_same_format_provider_upstream_url, classify_same_format_provider_request_behavior, + classify_same_format_provider_request_behavior_for_operation, + enforce_same_format_provider_api_operation_body_policy, resolve_same_format_provider_direct_auth, same_format_provider_transport_supported, same_format_provider_transport_unsupported_reason, same_format_provider_transport_unsupported_reason_for_trace, diff --git a/crates/aether-provider/transport/src/network.rs b/crates/aether-provider/transport/src/network.rs index 725105706..70224c4cf 100644 --- a/crates/aether-provider/transport/src/network.rs +++ b/crates/aether-provider/transport/src/network.rs @@ -6,6 +6,7 @@ use async_trait::async_trait; use serde_json::{json, Map, Value}; use tracing::warn; +use crate::claude_code::current_claude_code_transport_identity_profile; use crate::grok::grok_browser_resolved_transport_profile_from_auth_config; use super::snapshot::GatewayProviderTransportSnapshot; @@ -152,9 +153,38 @@ pub fn resolve_transport_profile_id( pub fn resolve_transport_profile( transport: &GatewayProviderTransportSnapshot, ) -> Option { - resolve_transport_profile_from_fingerprint(transport.key.fingerprint.as_ref()).or_else(|| { - resolve_transport_profile_from_provider_config(transport.provider.config.as_ref()) - .or_else(|| resolve_grok_browser_transport_profile(transport)) + let configured = resolve_transport_profile_from_fingerprint(transport.key.fingerprint.as_ref()) + .or_else(|| { + resolve_transport_profile_from_provider_config(transport.provider.config.as_ref()) + }); + if configured.is_some() || transport_profile_is_configured(transport) { + return configured; + } + + resolve_claude_code_transport_profile(transport) + .or_else(|| resolve_grok_browser_transport_profile(transport)) +} + +fn resolve_claude_code_transport_profile( + transport: &GatewayProviderTransportSnapshot, +) -> Option { + if !transport + .provider + .provider_type + .trim() + .eq_ignore_ascii_case("claude_code") + { + return None; + } + + let identity_profile = *current_claude_code_transport_identity_profile(); + Some(ResolvedTransportProfile { + profile_id: identity_profile.transport_profile_id().to_string(), + backend: TRANSPORT_BACKEND_REQWEST_RUSTLS.to_string(), + http_mode: TRANSPORT_HTTP_MODE_AUTO.to_string(), + pool_scope: TRANSPORT_POOL_SCOPE_KEY.to_string(), + header_fingerprint: None, + extra: None, }) } @@ -595,6 +625,64 @@ mod tests { assert_eq!(profile.backend, "reqwest_rustls"); } + #[test] + fn resolves_typed_claude_code_transport_profile_when_unconfigured() { + let mut transport = sample_transport(); + transport.provider.provider_type = "claude_code".to_string(); + transport.key.fingerprint = None; + transport.provider.config = None; + + let profile = resolve_transport_profile(&transport).expect("typed Claude Code profile"); + + assert_eq!(profile.profile_id, "claude_code_nodejs"); + assert_eq!(profile.backend, "reqwest_rustls"); + assert_eq!(profile.http_mode, "auto"); + assert_eq!(profile.pool_scope, "key"); + assert!(profile.header_fingerprint.is_none()); + assert!(profile.extra.is_none()); + assert!(!transport_profile_is_configured(&transport)); + } + + #[test] + fn explicit_claude_code_transport_profiles_precede_typed_default() { + let mut transport = sample_transport(); + transport.provider.provider_type = "claude_code".to_string(); + transport.provider.config = Some(json!({ + "fingerprint": {"transport_profile": "provider_claude_profile"} + })); + transport.key.fingerprint = Some(json!({ + "transport_profile": "key_claude_profile" + })); + + assert_eq!( + resolve_transport_profile(&transport) + .expect("key profile") + .profile_id, + "key_claude_profile" + ); + + transport.key.fingerprint = None; + assert_eq!( + resolve_transport_profile(&transport) + .expect("provider profile") + .profile_id, + "provider_claude_profile" + ); + } + + #[test] + fn invalid_explicit_claude_code_transport_profile_blocks_typed_default() { + let mut transport = sample_transport(); + transport.provider.provider_type = "claude_code".to_string(); + transport.provider.config = None; + transport.key.fingerprint = Some(json!({ + "transport_profile": {"backend": "reqwest_rustls"} + })); + + assert!(transport_profile_is_configured(&transport)); + assert!(resolve_transport_profile(&transport).is_none()); + } + #[test] fn maps_string_transport_profile_to_resolved_profile() { let profile = resolve_transport_profile(&sample_transport()).expect("profile"); diff --git a/crates/aether-provider/transport/src/oauth_refresh/mod.rs b/crates/aether-provider/transport/src/oauth_refresh/mod.rs index ee5aada75..07c9c5b45 100644 --- a/crates/aether-provider/transport/src/oauth_refresh/mod.rs +++ b/crates/aether-provider/transport/src/oauth_refresh/mod.rs @@ -12,7 +12,7 @@ use aether_runtime_state::{RuntimeLockLease, RuntimeState}; use async_trait::async_trait; use serde_json::Value; use thiserror::Error; -use tokio::sync::Mutex; +use tokio::sync::{Mutex, OwnedMutexGuard}; use super::agent_identity::{is_codex_agent_identity_transport, CodexAgentIdentityRefreshAdapter}; use super::generic_oauth::supports_local_generic_oauth_request_auth_resolution; @@ -47,6 +47,35 @@ pub struct LocalOAuthResolution { /// Held until the caller persists `refreshed_entry`. The lease TTL remains /// the cancellation fallback if the caller is dropped. pub distributed_lease: Option, + /// Keeps memory-only refreshes singleflight until the caller validates the + /// credential fence and publishes or discards `refreshed_entry`. + #[doc(hidden)] + pub local_refresh_guard: Option, +} + +#[derive(Clone)] +pub struct LocalOAuthRefreshCommitGuard { + guard: Arc>, +} + +impl LocalOAuthRefreshCommitGuard { + fn new(guard: OwnedMutexGuard<()>) -> Self { + Self { + guard: Arc::new(guard), + } + } +} + +impl fmt::Debug for LocalOAuthRefreshCommitGuard { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str("LocalOAuthRefreshCommitGuard") + } +} + +impl PartialEq for LocalOAuthRefreshCommitGuard { + fn eq(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.guard, &other.guard) + } } #[derive(Debug, Clone, PartialEq)] @@ -369,6 +398,12 @@ pub trait LocalOAuthRefreshAdapter: Send + Sync { false } + /// Whether another gateway instance can observe this refresh after the + /// caller persists its result into the provider transport record. + fn shares_refresh_through_transport_persistence(&self) -> bool { + true + } + async fn refresh( &self, executor: &dyn LocalOAuthHttpExecutor, @@ -388,6 +423,7 @@ pub struct LocalOAuthRefreshCoordinator { struct RefreshBackoffState { failures: u32, retry_after: Instant, + refresh_fingerprint: Option, } impl fmt::Debug for LocalOAuthRefreshCoordinator { @@ -424,6 +460,19 @@ impl LocalOAuthRefreshCoordinator { } } + /// Captures the refresh generation represented by this transport snapshot. + /// The coordinator cache is intentionally excluded: callers use this value + /// as the fence for the credential generation that produced their request. + pub fn refresh_fingerprint_for_transport( + &self, + transport: &GatewayProviderTransportSnapshot, + ) -> Option { + self.adapters + .iter() + .find(|adapter| adapter.supports(transport)) + .and_then(|adapter| adapter.refresh_fingerprint(transport, None)) + } + async fn lock_for_key(&self, key_id: &str) -> Arc> { let mut key_locks = self.key_locks.lock().await; key_locks @@ -530,6 +579,11 @@ impl LocalOAuthRefreshCoordinator { } else { self.cached_entry(key_id).await }; + let shares_refresh_through_transport_persistence = + adapter.shares_refresh_through_transport_persistence(); + let pre_lock_local_refresh_fingerprint = force_refresh + .then(|| adapter.refresh_fingerprint(transport, cached_entry.as_ref())) + .flatten(); if !force_refresh { if let Some(auth) = cached_entry .as_ref() @@ -548,7 +602,7 @@ impl LocalOAuthRefreshCoordinator { return Ok(None); } - if force_refresh { + if force_refresh && shares_refresh_through_transport_persistence { if let Some(resolution) = Self::resolve_if_refresh_fence_advanced( adapter.as_ref(), transport, @@ -558,25 +612,46 @@ impl LocalOAuthRefreshCoordinator { return Ok(Some(resolution)); } } - if let Some(error) = self.backoff_error(key_id, adapter.provider_type()).await { + let refresh_fingerprint = adapter.refresh_fingerprint(transport, cached_entry.as_ref()); + if let Some(error) = self + .backoff_error( + key_id, + adapter.provider_type(), + refresh_fingerprint.as_deref(), + ) + .await + { return Err(error); } let key_lock = self.lock_for_key(key_id).await; - let _key_guard = key_lock.lock().await; + let key_guard = key_lock.lock_owned().await; let cached_entry = self.cached_entry(key_id).await; if force_refresh { + let winner_fingerprint = if shares_refresh_through_transport_persistence { + expected_refresh_fingerprint + } else { + pre_lock_local_refresh_fingerprint.as_deref() + }; if let Some(resolution) = Self::resolve_if_refresh_fence_advanced( adapter.as_ref(), transport, cached_entry.as_ref(), - expected_refresh_fingerprint, + winner_fingerprint, ) { return Ok(Some(resolution)); } } - if let Some(error) = self.backoff_error(key_id, adapter.provider_type()).await { + let refresh_fingerprint = adapter.refresh_fingerprint(transport, cached_entry.as_ref()); + if let Some(error) = self + .backoff_error( + key_id, + adapter.provider_type(), + refresh_fingerprint.as_deref(), + ) + .await + { return Err(error); } if !force_refresh { @@ -594,40 +669,48 @@ impl LocalOAuthRefreshCoordinator { } } - let distributed_lease = match (distributed_lock, distributed_owner) { - (Some(lock), Some(owner)) if !owner.trim().is_empty() => { - match lock - .lock_try_acquire( - &format!("provider_oauth_refresh_lock:{key_id}"), - owner, - std::time::Duration::from_millis(Self::DISTRIBUTED_REFRESH_LOCK_TTL_MS), - ) - .await - { - Ok(Some(lease)) => Some(lease), - Ok(None) => return Ok(Some(LocalOAuthResolution::refresh_in_flight())), - Err(err) => { - tracing::warn!( - key_id = %key_id, - provider_type = adapter.provider_type(), - error = ?err, - "gateway local oauth refresh distributed lock unavailable" - ); - if adapter.requires_distributed_refresh_lock() { - let error = LocalOAuthRefreshError::TransportMessage { - provider_type: adapter.provider_type(), - message: "distributed refresh lock is unavailable".to_string(), - }; - if adapter.should_backoff_after_error(&error) { - self.record_refresh_failure(key_id).await; + let distributed_lease = if !shares_refresh_through_transport_persistence { + None + } else { + match (distributed_lock, distributed_owner) { + (Some(lock), Some(owner)) if !owner.trim().is_empty() => { + match lock + .lock_try_acquire( + &format!("provider_oauth_refresh_lock:{key_id}"), + owner, + std::time::Duration::from_millis(Self::DISTRIBUTED_REFRESH_LOCK_TTL_MS), + ) + .await + { + Ok(Some(lease)) => Some(lease), + Ok(None) => return Ok(Some(LocalOAuthResolution::refresh_in_flight())), + Err(err) => { + tracing::warn!( + key_id = %key_id, + provider_type = adapter.provider_type(), + error = ?err, + "gateway local oauth refresh distributed lock unavailable" + ); + if adapter.requires_distributed_refresh_lock() { + let error = LocalOAuthRefreshError::TransportMessage { + provider_type: adapter.provider_type(), + message: "distributed refresh lock is unavailable".to_string(), + }; + if adapter.should_backoff_after_error(&error) { + self.record_refresh_failure( + key_id, + refresh_fingerprint.as_deref(), + ) + .await; + } + return Err(error); } - return Err(error); + None } - None } } + _ => None, } - _ => None, }; // Forced refresh still needs the latest rotated refresh_token as input. @@ -653,7 +736,8 @@ impl LocalOAuthRefreshCoordinator { } Err(error) => { if adapter.should_backoff_after_error(&error) { - self.record_refresh_failure(key_id).await; + self.record_refresh_failure(key_id, refresh_fingerprint.as_deref()) + .await; } Self::release_distributed_lease( distributed_lock, @@ -665,14 +749,9 @@ impl LocalOAuthRefreshCoordinator { return Err(error); } }; - // In production the distributed lease is held through the gateway's - // DB CAS. Do not publish a provisional task before that CAS succeeds; - // otherwise a waiter could consume an assertion that loses the CAS. - // Lock-free/test callers retain the historical in-memory behavior. - if distributed_lease.is_none() { - self.insert_cached_entry(key_id, refreshed_entry.clone()) - .await; - } + // Cache publication belongs to the caller after durable persistence. + // The result still carries the entry so lock-free callers can inspect + // it or explicitly commit it with `store_cached_entry`. let Some(auth) = adapter.resolve_refreshed(transport, &refreshed_entry) else { Self::release_distributed_lease( distributed_lock, @@ -683,10 +762,16 @@ impl LocalOAuthRefreshCoordinator { .await; return Ok(None); }; + let local_refresh_guard = if shares_refresh_through_transport_persistence { + None + } else { + Some(LocalOAuthRefreshCommitGuard::new(key_guard)) + }; Ok(Some(LocalOAuthResolution::refreshed( auth, refreshed_entry, distributed_lease, + local_refresh_guard, ))) } @@ -728,8 +813,16 @@ impl LocalOAuthRefreshCoordinator { &self, key_id: &str, provider_type: &'static str, + refresh_fingerprint: Option<&str>, ) -> Option { - let backoff = self.refresh_backoff.lock().await; + let mut backoff = self.refresh_backoff.lock().await; + if backoff + .get(key_id) + .is_some_and(|state| state.refresh_fingerprint.as_deref() != refresh_fingerprint) + { + backoff.remove(key_id); + return None; + } let state = backoff.get(key_id)?; let remaining = state.retry_after.checked_duration_since(Instant::now())?; Some(LocalOAuthRefreshError::InvalidResponse { @@ -742,14 +835,19 @@ impl LocalOAuthRefreshCoordinator { }) } - async fn record_refresh_failure(&self, key_id: &str) { + async fn record_refresh_failure(&self, key_id: &str, refresh_fingerprint: Option<&str>) { let mut backoff = self.refresh_backoff.lock().await; let state = backoff .entry(key_id.to_string()) .or_insert(RefreshBackoffState { failures: 0, retry_after: Instant::now(), + refresh_fingerprint: refresh_fingerprint.map(ToOwned::to_owned), }); + if state.refresh_fingerprint.as_deref() != refresh_fingerprint { + state.failures = 0; + state.refresh_fingerprint = refresh_fingerprint.map(ToOwned::to_owned); + } state.failures = state.failures.saturating_add(1); let exponent = state.failures.saturating_sub(1).min(4); let delay = Duration::from_millis(500u64.saturating_mul(1u64 << exponent)); @@ -800,6 +898,7 @@ impl LocalOAuthResolution { refresh_in_flight: false, reused_refresh: false, distributed_lease: None, + local_refresh_guard: None, } } @@ -807,6 +906,7 @@ impl LocalOAuthResolution { auth: LocalResolvedOAuthRequestAuth, refreshed_entry: CachedOAuthEntry, distributed_lease: Option, + local_refresh_guard: Option, ) -> Self { Self { auth: Some(auth), @@ -814,6 +914,7 @@ impl LocalOAuthResolution { refresh_in_flight: false, reused_refresh: false, distributed_lease, + local_refresh_guard, } } @@ -829,6 +930,7 @@ impl LocalOAuthResolution { refresh_in_flight: false, reused_refresh: true, distributed_lease: None, + local_refresh_guard: None, } } @@ -839,6 +941,7 @@ impl LocalOAuthResolution { refresh_in_flight: true, reused_refresh: false, distributed_lease: None, + local_refresh_guard: None, } } } @@ -855,6 +958,7 @@ pub fn supports_local_oauth_request_auth_resolution( #[cfg(test)] mod tests { use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + use std::time::Duration; use super::super::snapshot::{ GatewayProviderTransportEndpoint, GatewayProviderTransportKey, @@ -878,6 +982,13 @@ mod tests { struct FencedTestAdapter { refresh_hits: Arc, fail_refresh: Arc, + generation: Arc, + } + + #[derive(Debug)] + struct MemoryOnlyFencedTestAdapter { + refresh_hits: Arc, + fingerprint_hits: Arc, } #[async_trait] @@ -978,7 +1089,12 @@ mod tests { ) -> Option { entry .and_then(|entry| entry.source_fingerprint.clone()) - .or_else(|| Some("generation-1".to_string())) + .or_else(|| { + Some(format!( + "generation-{}", + self.generation.load(Ordering::SeqCst) + )) + }) } fn should_backoff_after_error(&self, _error: &LocalOAuthRefreshError) -> bool { @@ -1004,7 +1120,75 @@ mod tests { auth_header_value: "stale-winner-cache-value".to_string(), expires_at_unix_secs: None, metadata: None, - source_fingerprint: Some("generation-2".to_string()), + source_fingerprint: Some(format!( + "generation-{}", + self.generation.load(Ordering::SeqCst).saturating_add(1) + )), + })) + } + } + + #[async_trait] + impl LocalOAuthRefreshAdapter for MemoryOnlyFencedTestAdapter { + fn provider_type(&self) -> &'static str { + "test-oauth" + } + + fn resolve_cached( + &self, + _transport: &GatewayProviderTransportSnapshot, + entry: &CachedOAuthEntry, + ) -> Option { + Some(LocalResolvedOAuthRequestAuth::Header { + name: entry.auth_header_name.clone(), + value: entry.auth_header_value.clone(), + }) + } + + fn resolve_without_refresh( + &self, + _transport: &GatewayProviderTransportSnapshot, + ) -> Option { + None + } + + fn should_refresh( + &self, + _transport: &GatewayProviderTransportSnapshot, + _entry: Option<&CachedOAuthEntry>, + ) -> bool { + true + } + + fn refresh_fingerprint( + &self, + _transport: &GatewayProviderTransportSnapshot, + entry: Option<&CachedOAuthEntry>, + ) -> Option { + self.fingerprint_hits.fetch_add(1, Ordering::SeqCst); + entry + .and_then(|entry| entry.source_fingerprint.clone()) + .or_else(|| Some("transport-generation".to_string())) + } + + fn shares_refresh_through_transport_persistence(&self) -> bool { + false + } + + async fn refresh( + &self, + _executor: &dyn LocalOAuthHttpExecutor, + _transport: &GatewayProviderTransportSnapshot, + _entry: Option<&CachedOAuthEntry>, + ) -> Result, LocalOAuthRefreshError> { + let hit = self.refresh_hits.fetch_add(1, Ordering::SeqCst) + 1; + Ok(Some(CachedOAuthEntry { + provider_type: "test-oauth".to_string(), + auth_header_name: "authorization".to_string(), + auth_header_value: format!("Bearer refreshed-token-{hit}"), + expires_at_unix_secs: Some(4_102_444_800), + metadata: None, + source_fingerprint: Some(format!("local-generation-{hit}")), })) } } @@ -1082,8 +1266,12 @@ mod tests { .resolve_with_result(&executor, &transport, None, None) .await .expect("first resolve should succeed"); + assert!(coordinator + .cached_entry(transport.key.id.as_str()) + .await + .is_none()); coordinator - .insert_cached_entry( + .store_cached_entry( transport.key.id.as_str(), first .as_ref() @@ -1115,6 +1303,7 @@ mod tests { refresh_in_flight: false, reused_refresh: false, distributed_lease: None, + local_refresh_guard: None, }) ); assert_eq!( @@ -1128,6 +1317,7 @@ mod tests { refresh_in_flight: false, reused_refresh: false, distributed_lease: None, + local_refresh_guard: None, }) ); } @@ -1149,7 +1339,7 @@ mod tests { .await .expect("initial resolve should succeed"); coordinator - .insert_cached_entry( + .store_cached_entry( transport.key.id.as_str(), first .as_ref() @@ -1168,6 +1358,98 @@ mod tests { assert_eq!(refresh_with_entry_hits.load(Ordering::SeqCst), 1); } + #[tokio::test] + async fn memory_only_force_refresh_does_not_reuse_preexisting_cache_as_winner() { + let refresh_hits = Arc::new(AtomicUsize::new(0)); + let fingerprint_hits = Arc::new(AtomicUsize::new(0)); + let coordinator = Arc::new(LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![ + Arc::new(MemoryOnlyFencedTestAdapter { + refresh_hits: Arc::clone(&refresh_hits), + fingerprint_hits: Arc::clone(&fingerprint_hits), + }), + ])); + let transport = sample_transport(); + let executor = ReqwestLocalOAuthHttpExecutor::new(reqwest::Client::new()); + coordinator + .store_cached_entry( + transport.key.id.as_str(), + CachedOAuthEntry { + provider_type: "test-oauth".to_string(), + auth_header_name: "authorization".to_string(), + auth_header_value: "Bearer rejected-token".to_string(), + expires_at_unix_secs: Some(4_102_444_800), + metadata: None, + source_fingerprint: Some("preexisting-local-generation".to_string()), + }, + ) + .await; + + let mut forced = coordinator + .force_refresh_with_result_fenced( + &executor, + &transport, + None, + None, + Some("transport-generation"), + ) + .await + .expect("memory-only force refresh should succeed") + .expect("memory-only force refresh should resolve"); + + assert_eq!(refresh_hits.load(Ordering::SeqCst), 1); + assert!(!forced.reused_refresh); + assert!(forced.local_refresh_guard.is_some()); + assert_eq!( + forced.auth, + Some(LocalResolvedOAuthRequestAuth::Header { + name: "authorization".to_string(), + value: "Bearer refreshed-token-1".to_string(), + }) + ); + + fingerprint_hits.store(0, Ordering::SeqCst); + let follower_coordinator = Arc::clone(&coordinator); + let follower_transport = transport.clone(); + let follower = tokio::spawn(async move { + follower_coordinator + .force_refresh_with_result_fenced( + &ReqwestLocalOAuthHttpExecutor::new(reqwest::Client::new()), + &follower_transport, + None, + None, + Some("transport-generation"), + ) + .await + .expect("memory-only follower should succeed") + .expect("memory-only follower should resolve") + }); + tokio::time::timeout(Duration::from_secs(1), async { + while fingerprint_hits.load(Ordering::SeqCst) == 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("follower should capture its pre-lock fingerprint"); + + coordinator + .store_cached_entry( + transport.key.id.as_str(), + forced + .refreshed_entry + .clone() + .expect("leader should provide the memory-only entry"), + ) + .await; + forced.local_refresh_guard.take(); + let follower = tokio::time::timeout(Duration::from_secs(1), follower) + .await + .expect("follower should unblock after cache publication") + .expect("follower task should join"); + + assert!(follower.reused_refresh); + assert_eq!(refresh_hits.load(Ordering::SeqCst), 1); + } + #[tokio::test] async fn fenced_force_refresh_reuses_the_winner() { let refresh_hits = Arc::new(AtomicUsize::new(0)); @@ -1175,6 +1457,7 @@ mod tests { FencedTestAdapter { refresh_hits: Arc::clone(&refresh_hits), fail_refresh: Arc::new(AtomicBool::new(false)), + generation: Arc::new(AtomicUsize::new(1)), }, )]); let transport = sample_transport(); @@ -1191,6 +1474,15 @@ mod tests { .await .expect("first refresh should succeed") .expect("first refresh should resolve"); + coordinator + .store_cached_entry( + transport.key.id.as_str(), + first + .refreshed_entry + .clone() + .expect("first refresh should return an entry to persist"), + ) + .await; let waiter = coordinator .force_refresh_with_result_fenced( &executor, @@ -1221,10 +1513,12 @@ mod tests { async fn refresh_failure_enters_bounded_negative_backoff() { let refresh_hits = Arc::new(AtomicUsize::new(0)); let fail_refresh = Arc::new(AtomicBool::new(true)); + let generation = Arc::new(AtomicUsize::new(1)); let coordinator = LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![Arc::new( FencedTestAdapter { refresh_hits: Arc::clone(&refresh_hits), fail_refresh: Arc::clone(&fail_refresh), + generation: Arc::clone(&generation), }, )]); let transport = sample_transport(); @@ -1241,11 +1535,11 @@ mod tests { assert!(second.to_string().contains("temporarily backed off")); assert_eq!(refresh_hits.load(Ordering::SeqCst), 1); fail_refresh.store(false, Ordering::SeqCst); - coordinator.invalidate_cached_entry("key-1").await; + generation.store(2, Ordering::SeqCst); assert!(coordinator .force_refresh_with_result(&executor, &transport, None, None) .await - .expect("replacement should refresh immediately") + .expect("new credential generation should bypass old backoff") .is_some()); assert_eq!(refresh_hits.load(Ordering::SeqCst), 2); } diff --git a/crates/aether-provider/transport/src/provider_types.rs b/crates/aether-provider/transport/src/provider_types.rs index f8efa4eaf..0a85b477d 100644 --- a/crates/aether-provider/transport/src/provider_types.rs +++ b/crates/aether-provider/transport/src/provider_types.rs @@ -363,7 +363,7 @@ const GEMINI_CLI_FIXED_PROVIDER_TEMPLATE: FixedProviderTemplate = FixedProviderT const VERTEX_AI_FIXED_PROVIDER_TEMPLATE: FixedProviderTemplate = FixedProviderTemplate { provider_type: "vertex_ai", - version: 1, + version: 2, base_url: "https://aiplatform.googleapis.com", endpoints: &[ FixedProviderEndpointTemplate { @@ -378,12 +378,6 @@ const VERTEX_AI_FIXED_PROVIDER_TEMPLATE: FixedProviderTemplate = FixedProviderTe custom_path: None, config_defaults: EMPTY_ENDPOINT_CONFIG_DEFAULTS, }, - FixedProviderEndpointTemplate { - item_key: "claude:messages", - api_format: "claude:messages", - custom_path: None, - config_defaults: EMPTY_ENDPOINT_CONFIG_DEFAULTS, - }, ], runtime_policy: VERTEX_AI_RUNTIME_POLICY, }; @@ -936,26 +930,27 @@ mod tests { } #[test] - fn vertex_fixed_provider_template_includes_gemini_embedding_endpoint() { + fn vertex_fixed_provider_template_exposes_only_implemented_gemini_endpoints() { let template = fixed_provider_template("vertex_ai").expect("vertex_ai template should exist"); + assert_eq!(template.version, 2); assert_eq!( template .endpoints .iter() .map(|item| item.api_format) .collect::>(), - vec![ - "gemini:generate_content", - "gemini:embedding", - "claude:messages", - ] + vec!["gemini:generate_content", "gemini:embedding"] ); assert!( fixed_provider_endpoint_template_by_api_format("vertex_ai", "gemini:embedding") .is_some() ); + assert!( + fixed_provider_endpoint_template_by_api_format("vertex_ai", "claude:messages") + .is_none() + ); } } diff --git a/crates/aether-provider/transport/src/request_url/mod.rs b/crates/aether-provider/transport/src/request_url/mod.rs index ab77d7e4b..e045f7cd4 100644 --- a/crates/aether-provider/transport/src/request_url/mod.rs +++ b/crates/aether-provider/transport/src/request_url/mod.rs @@ -1,6 +1,7 @@ use std::collections::BTreeMap; use std::sync::OnceLock; +use aether_ai_formats::ApiOperation; use regex::Regex; use serde_json::Value; use url::form_urlencoded; @@ -15,9 +16,11 @@ use crate::gemini_cli::{ }; use crate::snapshot::GatewayProviderTransportSnapshot; use crate::url::{ + build_claude_count_tokens_url as build_default_claude_count_tokens_url, build_claude_messages_url, build_gemini_content_url, build_openai_chat_url, build_openai_responses_url, build_openai_search_url, build_passthrough_path_url, - normalize_gemini_content_action_path, + normalize_gemini_content_action_path, strip_gateway_credential_query_parameters, + GATEWAY_CREDENTIAL_QUERY_KEYS, }; use crate::vertex::{ build_vertex_api_key_gemini_content_url, build_vertex_api_key_gemini_embedding_url, @@ -32,6 +35,7 @@ pub struct TransportRequestUrlParams<'a> { pub upstream_is_stream: bool, pub request_query: Option<&'a str>, pub kiro_api_region: Option<&'a str>, + pub api_operation: Option, } pub fn build_transport_request_url( @@ -70,23 +74,62 @@ fn build_transport_request_url_inner( let provider_api_format = params.provider_api_format.trim().to_ascii_lowercase(); let normalized_provider_api_format = aether_ai_formats::normalize_api_format_alias(&provider_api_format); + let sanitized_claude_request_query = (normalized_provider_api_format == "claude:messages") + .then(|| strip_gateway_credential_query_parameters(params.request_query)) + .flatten(); + let params = if normalized_provider_api_format == "claude:messages" { + TransportRequestUrlParams { + request_query: sanitized_claude_request_query.as_deref(), + ..params + } + } else { + params + }; + let is_claude_count_tokens = normalized_provider_api_format == "claude:messages" + && params.api_operation == Some(ApiOperation::ClaudeCountTokens); + if !transport_supports_api_operation( + transport, + normalized_provider_api_format.as_str(), + params.api_operation, + ) { + return None; + } + if is_claude_count_tokens { + if let Some(url) = build_configured_claude_count_tokens_url(transport, params.request_query) + { + return Some(url); + } + } if let Some(url) = build_transport_hook_url(transport, params) { return Some(url); } - let custom_path = transport + let custom_path_template = transport .endpoint .custom_path .as_deref() .map(str::trim) - .filter(|value| !value.is_empty()) - .map(|path| { - expand_custom_path_template(path, build_path_params(params, gemini_embedding_batch)) - }); + .filter(|value| !value.is_empty()); + let custom_path_handles_operation = + custom_path_template.is_some_and(|path| path.contains("{operation}")); + let custom_path = custom_path_template.map(|path| { + expand_custom_path_template(path, build_path_params(params, gemini_embedding_batch)) + }); if let Some(path) = custom_path.as_deref() { - let blocked_keys = if normalized_provider_api_format.starts_with("gemini:") { - &["key"][..] + let custom_path_is_complete_claude_count_tokens = normalized_provider_api_format + == "claude:messages" + && !custom_path_handles_operation + && path + .split_once('?') + .map(|(path, _)| path) + .unwrap_or(path) + .trim_end_matches('/') + .ends_with("/messages/count_tokens"); + let blocked_keys = if normalized_provider_api_format.starts_with("gemini:") + || normalized_provider_api_format == "claude:messages" + { + GATEWAY_CREDENTIAL_QUERY_KEYS } else { &[][..] }; @@ -97,12 +140,19 @@ fn build_transport_request_url_inner( } else { path.to_string() }; - let url = build_passthrough_path_url( + let mut url = build_passthrough_path_url( &transport.endpoint.base_url, normalized_path.as_str(), params.request_query, blocked_keys, )?; + if is_claude_count_tokens && !custom_path_handles_operation { + url = build_default_claude_count_tokens_url(&url, None); + } else if params.api_operation == Some(ApiOperation::ClaudeMessagesCreate) + && custom_path_is_complete_claude_count_tokens + { + url = build_claude_messages_url(&url, None); + } return Some(maybe_add_gemini_stream_alt_sse( url, &provider_api_format, @@ -139,10 +189,14 @@ fn build_transport_request_url_inner( "openai:rerank" | "jina:rerank" => { build_provider_rerank_v1_url(&transport.endpoint.base_url, params.request_query) } - "claude:messages" => Some(build_claude_messages_url( - &transport.endpoint.base_url, - params.request_query, - )), + "claude:messages" => Some(if is_claude_count_tokens { + build_default_claude_count_tokens_url( + &transport.endpoint.base_url, + params.request_query, + ) + } else { + build_claude_messages_url(&transport.endpoint.base_url, params.request_query) + }), "gemini:generate_content" => build_gemini_content_url( &transport.endpoint.base_url, params.mapped_model?, @@ -186,6 +240,7 @@ pub fn build_local_openai_chat_upstream_url( upstream_is_stream: false, request_query, kiro_api_region: None, + api_operation: None, }, ) } @@ -206,6 +261,7 @@ pub fn build_cross_format_openai_chat_upstream_url( upstream_is_stream, request_query, kiro_api_region: None, + api_operation: None, }, ) } @@ -228,6 +284,7 @@ pub fn build_local_openai_responses_upstream_url( upstream_is_stream: false, request_query, kiro_api_region: None, + api_operation: None, }, ) } @@ -249,6 +306,7 @@ pub fn build_cross_format_openai_responses_upstream_url( upstream_is_stream, request_query, kiro_api_region: None, + api_operation: None, }, ) } @@ -269,6 +327,7 @@ pub fn build_kiro_cross_format_upstream_url( upstream_is_stream, request_query, kiro_api_region: Some(api_region), + api_operation: None, }, ) } @@ -291,10 +350,15 @@ fn build_transport_hook_url( .trim() .eq_ignore_ascii_case("claude_code") { - return Some(build_claude_code_messages_url( - &transport.endpoint.base_url, - params.request_query, - )); + let messages_url = + build_claude_code_messages_url(&transport.endpoint.base_url, params.request_query); + return Some( + if params.api_operation == Some(aether_ai_formats::ApiOperation::ClaudeCountTokens) { + build_default_claude_count_tokens_url(&messages_url, None) + } else { + messages_url + }, + ); } let normalized_provider_api_format = @@ -408,6 +472,9 @@ fn build_path_params( } let provider_api_format = aether_ai_formats::normalize_api_format_alias(params.provider_api_format); + if let Some(operation) = params.api_operation { + path_params.insert("operation", operation.as_str()); + } if provider_api_format == "gemini:generate_content" || provider_api_format == "gemini:embedding" { path_params.insert( @@ -428,6 +495,79 @@ fn build_path_params( path_params } +pub fn transport_supports_api_operation( + transport: &GatewayProviderTransportSnapshot, + provider_api_format: &str, + operation: Option, +) -> bool { + if operation != Some(ApiOperation::ClaudeCountTokens) { + return true; + } + if aether_ai_formats::normalize_api_format_alias(provider_api_format) != "claude:messages" { + return false; + } + + anthropic_count_tokens_supported(transport) +} + +fn anthropic_count_tokens_supported(transport: &GatewayProviderTransportSnapshot) -> bool { + // Private message adapters do not implement Anthropic's token-counting + // operation. A config flag cannot make their request envelopes compatible. + if crate::kiro::is_kiro_provider_transport(transport) + || crate::grok::is_grok_provider_transport(transport) + { + return false; + } + + let Some(operations) = anthropic_transport_config_field(transport, "supported_operations") + else { + return true; + }; + operations.as_array().is_some_and(|operations| { + operations.iter().any(|operation| { + operation + .as_str() + .is_some_and(|value| value.eq_ignore_ascii_case("count_tokens")) + }) + }) +} + +fn build_configured_claude_count_tokens_url( + transport: &GatewayProviderTransportSnapshot, + request_query: Option<&str>, +) -> Option { + let path = anthropic_transport_config_field(transport, "count_tokens_path") + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()); + let path = path?; + let path = path + .starts_with('/') + .then(|| path.to_string()) + .unwrap_or_else(|| format!("/{path}")); + build_passthrough_path_url( + &transport.endpoint.base_url, + path.as_str(), + request_query, + GATEWAY_CREDENTIAL_QUERY_KEYS, + ) +} + +fn anthropic_transport_config_field<'a>( + transport: &'a GatewayProviderTransportSnapshot, + field: &str, +) -> Option<&'a Value> { + let from_config = |config: Option<&'a Value>| { + config + .and_then(Value::as_object) + .and_then(|config| config.get("anthropic")) + .and_then(Value::as_object) + .and_then(|anthropic| anthropic.get(field)) + }; + from_config(transport.endpoint.config.as_ref()) + .or_else(|| from_config(transport.provider.config.as_ref())) +} + fn normalize_gemini_embedding_action_path(path: &str, batch: bool) -> String { if batch { path.replace(":embedContent", ":batchEmbedContents") @@ -571,6 +711,8 @@ fn custom_path_template_regex() -> &'static Regex { #[cfg(test)] mod tests { + use aether_ai_formats::ApiOperation; + use super::{ build_kiro_cross_format_upstream_url, build_transport_request_url, build_transport_request_url_for_request_body, TransportRequestUrlParams, @@ -660,6 +802,7 @@ mod tests { upstream_is_stream: true, request_query: Some("foo=bar"), kiro_api_region: None, + api_operation: None, }, ) .expect("vertex hook url"); @@ -697,6 +840,7 @@ mod tests { upstream_is_stream: false, request_query: Some("foo=bar&beta=1"), kiro_api_region: None, + api_operation: None, }, ) .expect("vertex service account hook url"); @@ -738,6 +882,7 @@ mod tests { upstream_is_stream: false, request_query: Some("foo=bar&beta=1"), kiro_api_region: None, + api_operation: None, }, Some(&provider_request_body), ) @@ -767,6 +912,7 @@ mod tests { upstream_is_stream: false, request_query: Some("key=blocked&beta=true&foo=bar"), kiro_api_region: None, + api_operation: None, }, ) .as_deref(), @@ -792,6 +938,7 @@ mod tests { upstream_is_stream: true, request_query: Some("foo=bar"), kiro_api_region: None, + api_operation: None, }, ) .as_deref(), @@ -839,6 +986,7 @@ mod tests { upstream_is_stream: false, request_query: None, kiro_api_region: None, + api_operation: None, }, Some(&batch_body), ) @@ -866,6 +1014,7 @@ mod tests { upstream_is_stream: false, request_query: Some("tenant=demo"), kiro_api_region: None, + api_operation: None, }, ) .expect("openai responses url"); @@ -890,6 +1039,7 @@ mod tests { upstream_is_stream: false, request_query: Some("tenant=demo"), kiro_api_region: None, + api_operation: None, }, ) .expect("openai search url"); @@ -917,6 +1067,7 @@ mod tests { upstream_is_stream: false, request_query: Some("key=client-key&foo=bar"), kiro_api_region: None, + api_operation: None, }, ) .expect("expanded custom path url"); @@ -944,6 +1095,7 @@ mod tests { upstream_is_stream: true, request_query: Some("key=client-key&foo=bar"), kiro_api_region: None, + api_operation: None, }, ) .expect("stream custom path url"); @@ -968,6 +1120,7 @@ mod tests { upstream_is_stream: false, request_query: Some("foo=bar"), kiro_api_region: None, + api_operation: None, }, ) .expect("sync custom path url"); @@ -992,6 +1145,7 @@ mod tests { upstream_is_stream: true, request_query: None, kiro_api_region: None, + api_operation: None, }, ) .expect("v1 stream custom path url"); @@ -1019,6 +1173,7 @@ mod tests { upstream_is_stream: false, request_query: None, kiro_api_region: None, + api_operation: None, }, ) .expect("fallback custom path url"); @@ -1026,6 +1181,374 @@ mod tests { assert_eq!(url, "https://api.example.com/v1/messages/{model}"); } + #[test] + fn strips_gateway_query_key_from_claude_messages_url() { + let transport = sample_transport( + "custom", + "claude:messages", + "https://api.anthropic.example/v1?region=us", + None, + ); + + let url = build_transport_request_url( + &transport, + TransportRequestUrlParams { + provider_api_format: "claude:messages", + mapped_model: Some("claude-sonnet-4"), + upstream_is_stream: false, + request_query: Some("key=gateway-secret&trace=1"), + kiro_api_region: None, + api_operation: Some(aether_ai_formats::ApiOperation::ClaudeMessagesCreate), + }, + ) + .expect("messages url"); + + assert_eq!( + url, + "https://api.anthropic.example/v1/messages?region=us&trace=1" + ); + assert!(!url.contains("gateway-secret")); + } + + #[test] + fn routes_claude_count_tokens_as_an_operation_on_messages_format() { + let mut transport = sample_transport( + "custom", + "claude:messages", + "https://api.anthropic.example/v1", + None, + ); + transport.endpoint.config = Some(json!({ + "anthropic": { + "supported_operations": ["messages", "count_tokens"] + } + })); + + let url = build_transport_request_url( + &transport, + TransportRequestUrlParams { + provider_api_format: "claude:messages", + mapped_model: Some("claude-sonnet-4"), + upstream_is_stream: false, + request_query: Some("key=gateway-secret&trace=1"), + kiro_api_region: None, + api_operation: Some(aether_ai_formats::ApiOperation::ClaudeCountTokens), + }, + ) + .expect("count_tokens url"); + + assert_eq!( + url, + "https://api.anthropic.example/v1/messages/count_tokens?trace=1" + ); + } + + #[test] + fn count_tokens_default_preserves_custom_anthropic_prefix() { + let transport = sample_transport( + "custom", + "claude:messages", + "https://proxy.example/anthropic?key=base-secret&tenant=base", + None, + ); + + assert_eq!( + build_transport_request_url( + &transport, + TransportRequestUrlParams { + provider_api_format: "claude:messages", + mapped_model: Some("claude-sonnet-4"), + upstream_is_stream: false, + request_query: Some("KEY=client-secret&trace=1"), + kiro_api_region: None, + api_operation: Some(aether_ai_formats::ApiOperation::ClaudeCountTokens), + }, + ) + .as_deref(), + Some( + "https://proxy.example/anthropic/messages/count_tokens?key=base-secret&tenant=base&trace=1" + ) + ); + } + + #[test] + fn count_tokens_uses_operation_aware_custom_path() { + let mut transport = sample_transport( + "custom", + "claude:messages", + "https://proxy.example/anthropic", + Some("/operations/{operation}"), + ); + transport.endpoint.config = Some(json!({ + "anthropic": {"supported_operations": ["messages", "count_tokens"]} + })); + + assert_eq!( + build_transport_request_url( + &transport, + TransportRequestUrlParams { + provider_api_format: "claude:messages", + mapped_model: Some("claude-sonnet-4"), + upstream_is_stream: false, + request_query: Some("key=client-secret&trace=1"), + kiro_api_region: None, + api_operation: Some(aether_ai_formats::ApiOperation::ClaudeCountTokens), + }, + ) + .as_deref(), + Some("https://proxy.example/anthropic/operations/count_tokens?trace=1") + ); + } + + #[test] + fn count_tokens_derives_from_messages_only_custom_path() { + let transport = sample_transport( + "custom", + "claude:messages", + "https://api.anthropic.example", + Some("/custom/v1/messages"), + ); + + assert_eq!( + build_transport_request_url( + &transport, + TransportRequestUrlParams { + provider_api_format: "claude:messages", + mapped_model: Some("claude-sonnet-4"), + upstream_is_stream: false, + request_query: None, + kiro_api_region: None, + api_operation: Some(aether_ai_formats::ApiOperation::ClaudeCountTokens), + }, + ) + .as_deref(), + Some("https://api.anthropic.example/custom/v1/messages/count_tokens") + ); + } + + #[test] + fn count_tokens_keeps_complete_custom_count_tokens_path() { + let transport = sample_transport( + "custom", + "claude:messages", + "https://api.anthropic.example", + Some("/custom/v1/messages/count_tokens?key=path-secret"), + ); + + assert_eq!( + build_transport_request_url( + &transport, + TransportRequestUrlParams { + provider_api_format: "claude:messages", + mapped_model: Some("claude-sonnet-4"), + upstream_is_stream: false, + request_query: Some("key=client-secret&trace=1"), + kiro_api_region: None, + api_operation: Some(aether_ai_formats::ApiOperation::ClaudeCountTokens), + }, + ) + .as_deref(), + Some( + "https://api.anthropic.example/custom/v1/messages/count_tokens?key=path-secret&trace=1" + ) + ); + } + + #[test] + fn messages_create_normalizes_complete_custom_count_tokens_path() { + let transport = sample_transport( + "custom", + "claude:messages", + "https://api.anthropic.example", + Some("/custom/v1/messages/count_tokens?key=path-secret"), + ); + + assert_eq!( + build_transport_request_url( + &transport, + TransportRequestUrlParams { + provider_api_format: "claude:messages", + mapped_model: Some("claude-sonnet-4"), + upstream_is_stream: false, + request_query: Some("key=client-secret&trace=1"), + kiro_api_region: None, + api_operation: Some(aether_ai_formats::ApiOperation::ClaudeMessagesCreate), + }, + ) + .as_deref(), + Some("https://api.anthropic.example/custom/v1/messages?key=path-secret&trace=1") + ); + } + + #[test] + fn private_anthropic_adapters_reject_count_tokens_even_when_config_claims_support() { + for provider_type in ["kiro", "grok"] { + let mut transport = sample_transport( + provider_type, + "claude:messages", + "https://private.example", + None, + ); + transport.endpoint.config = Some(json!({ + "anthropic": { + "supported_operations": ["messages", "count_tokens"], + "count_tokens_path": "/v1/messages/count_tokens" + } + })); + + assert!(build_transport_request_url( + &transport, + TransportRequestUrlParams { + provider_api_format: "claude:messages", + mapped_model: Some("claude-sonnet-4"), + upstream_is_stream: false, + request_query: None, + kiro_api_region: Some("us-east-1"), + api_operation: Some(ApiOperation::ClaudeCountTokens), + }, + ) + .is_none()); + } + } + + #[test] + fn claude_code_custom_root_keeps_messages_and_count_tokens_on_v1_surface() { + let mut transport = sample_transport( + "claude_code", + "claude:messages", + "https://proxy.example?key=base-secret", + None, + ); + transport.endpoint.config = Some(json!({ + "anthropic": {"supported_operations": ["messages", "count_tokens"]} + })); + + assert_eq!( + build_transport_request_url( + &transport, + TransportRequestUrlParams { + provider_api_format: "claude:messages", + mapped_model: Some("claude-sonnet-4"), + upstream_is_stream: false, + request_query: Some("key=client-secret&trace=messages"), + kiro_api_region: None, + api_operation: Some(aether_ai_formats::ApiOperation::ClaudeMessagesCreate), + }, + ) + .as_deref(), + Some("https://proxy.example/v1/messages?key=base-secret&trace=messages") + ); + assert_eq!( + build_transport_request_url( + &transport, + TransportRequestUrlParams { + provider_api_format: "claude:messages", + mapped_model: Some("claude-sonnet-4"), + upstream_is_stream: false, + request_query: Some("key=client-secret&trace=count"), + kiro_api_region: None, + api_operation: Some(aether_ai_formats::ApiOperation::ClaudeCountTokens), + }, + ) + .as_deref(), + Some("https://proxy.example/v1/messages/count_tokens?key=base-secret&trace=count") + ); + } + + #[test] + fn count_tokens_config_fields_fall_back_from_endpoint_to_provider() { + let mut transport = sample_transport( + "custom", + "claude:messages", + "https://api.anthropic.example/v1?key=base-secret", + None, + ); + transport.endpoint.config = Some(json!({ + "anthropic": {"profile": "native_transparent"} + })); + transport.provider.config = Some(json!({ + "anthropic": { + "supported_operations": ["messages", "count_tokens"], + "count_tokens_path": "/v1/messages/provider_count_tokens?key=path-secret" + } + })); + + assert_eq!( + build_transport_request_url( + &transport, + TransportRequestUrlParams { + provider_api_format: "claude:messages", + mapped_model: Some("claude-sonnet-4"), + upstream_is_stream: false, + request_query: Some("key=client-secret&trace=1"), + kiro_api_region: None, + api_operation: Some(aether_ai_formats::ApiOperation::ClaudeCountTokens), + }, + ) + .as_deref(), + Some( + "https://api.anthropic.example/v1/messages/provider_count_tokens?key=path-secret&trace=1" + ) + ); + } + + #[test] + fn rejects_count_tokens_when_provider_capability_excludes_it() { + let mut transport = sample_transport( + "custom", + "claude:messages", + "https://api.anthropic.example/v1", + None, + ); + transport.endpoint.config = Some(json!({ + "anthropic": {"profile": "native_transparent"} + })); + transport.provider.config = Some(json!({ + "anthropic": {"supported_operations": ["messages"]} + })); + + assert!(build_transport_request_url( + &transport, + TransportRequestUrlParams { + provider_api_format: "claude:messages", + mapped_model: Some("claude-sonnet-4"), + upstream_is_stream: false, + request_query: None, + kiro_api_region: None, + api_operation: Some(aether_ai_formats::ApiOperation::ClaudeCountTokens), + }, + ) + .is_none()); + } + + #[test] + fn preserves_configured_query_credentials_and_strips_client_credentials() { + let transport = sample_transport( + "custom", + "claude:messages", + "https://proxy.example/anthropic?key=base-secret&tenant=base", + Some("/messages?key=path-secret&variant=custom"), + ); + + assert_eq!( + build_transport_request_url( + &transport, + TransportRequestUrlParams { + provider_api_format: "claude:messages", + mapped_model: Some("claude-sonnet-4"), + upstream_is_stream: false, + request_query: Some("KEY=client-secret&trace=1"), + kiro_api_region: None, + api_operation: None, + }, + ) + .as_deref(), + Some( + "https://proxy.example/anthropic/messages?key=path-secret&tenant=base&trace=1&variant=custom" + ) + ); + } + #[test] fn kiro_cross_format_helper_uses_region_specific_generate_assistant_url() { let transport = sample_transport( @@ -1040,7 +1563,7 @@ mod tests { "claude-sonnet-4", "claude:messages", true, - Some("conversationId=abc"), + Some("key=gateway-secret&conversationId=abc"), "us-west-2", ) .expect("kiro url"); @@ -1049,6 +1572,7 @@ mod tests { "https://codewhisperer.us-west-2.amazonaws.com/generateAssistantResponse" )); assert!(url.contains("conversationId=abc")); + assert!(!url.contains("gateway-secret")); } #[test] @@ -1093,6 +1617,7 @@ mod tests { upstream_is_stream: false, request_query: Some("tenant=demo"), kiro_api_region: None, + api_operation: None, }, ) .as_deref(), @@ -1107,6 +1632,7 @@ mod tests { upstream_is_stream: false, request_query: None, kiro_api_region: None, + api_operation: None, }, ) .as_deref(), @@ -1121,6 +1647,7 @@ mod tests { upstream_is_stream: false, request_query: Some("key=client-key&foo=bar"), kiro_api_region: None, + api_operation: None, }, ) .as_deref(), @@ -1137,6 +1664,7 @@ mod tests { upstream_is_stream: false, request_query: None, kiro_api_region: None, + api_operation: None, }, ) .as_deref(), @@ -1151,6 +1679,7 @@ mod tests { upstream_is_stream: false, request_query: None, kiro_api_region: None, + api_operation: None, }, ) .as_deref(), @@ -1182,6 +1711,7 @@ mod tests { upstream_is_stream: false, request_query: Some("key=client-key&trace=1"), kiro_api_region: None, + api_operation: None, }, ) .as_deref(), @@ -1196,6 +1726,7 @@ mod tests { upstream_is_stream: true, request_query: Some("key=client-key&trace=2"), kiro_api_region: None, + api_operation: None, }, ) .as_deref(), @@ -1221,6 +1752,7 @@ mod tests { upstream_is_stream: false, request_query: Some("key=client-key"), kiro_api_region: None, + api_operation: None, }, ) .as_deref(), @@ -1246,6 +1778,7 @@ mod tests { upstream_is_stream: true, request_query: Some("key=client-aether-key&trace=1&beta=true"), kiro_api_region: None, + api_operation: None, }, ) .as_deref(), @@ -1279,6 +1812,7 @@ mod tests { upstream_is_stream: false, request_query: Some("trace=1"), kiro_api_region: None, + api_operation: None, }, ) .as_deref(), @@ -1293,6 +1827,7 @@ mod tests { upstream_is_stream: false, request_query: None, kiro_api_region: None, + api_operation: None, }, ) .as_deref(), @@ -1332,6 +1867,7 @@ mod tests { upstream_is_stream: false, request_query: Some("key=client-key&foo=bar"), kiro_api_region: None, + api_operation: None, }, Some(&batch_body), ) @@ -1368,6 +1904,7 @@ mod tests { upstream_is_stream: false, request_query: None, kiro_api_region: None, + api_operation: None, }, Some(&batch_body), ) @@ -1397,6 +1934,7 @@ mod tests { upstream_is_stream: false, request_query: Some("tenant=demo"), kiro_api_region: None, + api_operation: None, }, ) .as_deref(), @@ -1411,6 +1949,7 @@ mod tests { upstream_is_stream: false, request_query: None, kiro_api_region: None, + api_operation: None, }, ) .as_deref(), @@ -1454,6 +1993,7 @@ mod tests { upstream_is_stream: false, request_query: Some("tenant=request&trace=1"), kiro_api_region: None, + api_operation: None, }, ) .as_deref(), @@ -1468,6 +2008,7 @@ mod tests { upstream_is_stream: false, request_query: Some("trace=2"), kiro_api_region: None, + api_operation: None, }, ) .as_deref(), @@ -1482,6 +2023,7 @@ mod tests { upstream_is_stream: false, request_query: Some("key=client-key&trace=3"), kiro_api_region: None, + api_operation: None, }, ) .as_deref(), @@ -1496,6 +2038,7 @@ mod tests { upstream_is_stream: false, request_query: Some("trace=4"), kiro_api_region: None, + api_operation: None, }, ) .as_deref(), @@ -1520,6 +2063,7 @@ mod tests { upstream_is_stream: false, request_query: None, kiro_api_region: None, + api_operation: None, }, ) .is_none()); @@ -1542,6 +2086,7 @@ mod tests { upstream_is_stream: true, request_query: Some("key=client-key&foo=bar"), kiro_api_region: None, + api_operation: None, }, ) .expect("expanded custom embedding path url"); diff --git a/crates/aether-provider/transport/src/same_format_provider/mod.rs b/crates/aether-provider/transport/src/same_format_provider/mod.rs index 8dc7c2190..a504a2c08 100644 --- a/crates/aether-provider/transport/src/same_format_provider/mod.rs +++ b/crates/aether-provider/transport/src/same_format_provider/mod.rs @@ -3,15 +3,22 @@ use std::collections::BTreeMap; use serde::Serialize; use serde_json::Value; +use crate::anthropic_compat::{ + resolve_anthropic_compatibility_profile, AnthropicCompatibilityProfile, +}; use crate::antigravity::is_antigravity_provider_transport; use crate::auth::{ build_complete_passthrough_headers, build_complete_passthrough_headers_with_auth, - resolve_local_gemini_auth, resolve_local_openai_bearer_auth, resolve_local_standard_auth, + replace_upstream_auth_headers, resolve_local_gemini_auth, resolve_local_openai_bearer_auth, + resolve_local_standard_auth, +}; +use crate::claude_code::{ + build_claude_code_passthrough_headers, current_claude_code_transport_identity_profile, + local_claude_code_transport_unsupported_reason_with_network, }; -use crate::claude_code::build_claude_code_passthrough_headers; -use crate::claude_code::local_claude_code_transport_unsupported_reason_with_network; use crate::gemini_cli::is_gemini_cli_provider_transport; use crate::grok::{is_grok_provider_transport, resolve_grok_session_auth}; +use crate::headers::{force_identity_accept_encoding, upstream_credential_header_names}; use crate::kiro::{ build_kiro_provider_headers, build_kiro_provider_request_body, is_kiro_provider_transport, local_kiro_request_transport_unsupported_reason_with_network, KiroAuthConfig, @@ -29,10 +36,8 @@ use crate::vertex::{ is_vertex_service_account_transport_context, is_vertex_transport_context, local_vertex_gemini_transport_unsupported_reason_with_network, }; -use crate::{ - build_transport_request_url_for_request_body, ensure_upstream_auth_header, - TransportRequestUrlParams, -}; + +use crate::{build_transport_request_url_for_request_body, TransportRequestUrlParams}; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum SameFormatProviderFamily { @@ -52,6 +57,8 @@ pub struct SameFormatProviderRequestBehavior { pub is_antigravity: bool, pub is_gemini_cli: bool, pub is_claude_code: bool, + pub is_claude_code_transport: bool, + pub anthropic_compatibility_profile: AnthropicCompatibilityProfile, pub is_vertex: bool, pub is_kiro: bool, pub upstream_is_stream: bool, @@ -106,6 +113,7 @@ pub struct SameFormatProviderUpstreamUrlParams<'a> { pub upstream_is_stream: bool, pub request_query: Option<&'a str>, pub kiro_api_region: Option<&'a str>, + pub api_operation: Option, pub provider_request_body: Option<&'a Value>, } @@ -116,10 +124,10 @@ pub struct SameFormatProviderHeadersInput<'a> { pub original_request_body: &'a Value, pub header_rules: Option<&'a Value>, pub behavior: SameFormatProviderRequestBehavior, + pub api_operation: Option, pub auth_header: Option<&'a str>, pub auth_value: Option<&'a str>, pub extra_headers: &'a BTreeMap, - pub key_fingerprint: Option<&'a Value>, pub kiro_auth_config: Option<&'a KiroAuthConfig>, pub kiro_machine_id: Option<&'a str>, } @@ -127,14 +135,25 @@ pub struct SameFormatProviderHeadersInput<'a> { pub fn classify_same_format_provider_request_behavior( transport: &GatewayProviderTransportSnapshot, params: SameFormatProviderRequestBehaviorParams<'_>, +) -> SameFormatProviderRequestBehavior { + classify_same_format_provider_request_behavior_for_operation(transport, params, None) +} + +pub fn classify_same_format_provider_request_behavior_for_operation( + transport: &GatewayProviderTransportSnapshot, + params: SameFormatProviderRequestBehaviorParams<'_>, + api_operation: Option, ) -> SameFormatProviderRequestBehavior { let is_antigravity = is_antigravity_provider_transport(transport); let is_gemini_cli = is_gemini_cli_provider_transport(transport); - let is_claude_code = transport + let is_claude_code_transport = transport .provider .provider_type .trim() .eq_ignore_ascii_case("claude_code"); + let anthropic_compatibility_profile = + resolve_anthropic_compatibility_profile(transport, params.provider_api_format); + let is_claude_code = anthropic_compatibility_profile.uses_claude_code_compatibility(); let is_vertex = is_vertex_transport_context(transport); let is_kiro = is_kiro_provider_transport(transport); let gemini_cli_requires_upstream_streaming = is_gemini_cli @@ -142,18 +161,23 @@ pub fn classify_same_format_provider_request_behavior( params.provider_api_format, params.require_streaming, ); - let upstream_is_stream = aether_ai_formats::resolve_upstream_is_stream_for_provider( - transport.endpoint.config.as_ref(), - transport.provider.provider_type.as_str(), - params.provider_api_format, - params.require_streaming, - is_kiro || is_antigravity || gemini_cli_requires_upstream_streaming, + let operation_requires_sync = matches!( + api_operation, + Some(aether_ai_formats::ApiOperation::ClaudeCountTokens) ); - let force_body_stream_field = - aether_ai_formats::api_format_uses_body_stream_field(params.provider_api_format) - && aether_ai_formats::endpoint_config_forces_upstream_stream_policy( - transport.endpoint.config.as_ref(), - ); + let upstream_is_stream = !operation_requires_sync + && aether_ai_formats::resolve_upstream_is_stream_for_provider( + transport.endpoint.config.as_ref(), + transport.provider.provider_type.as_str(), + params.provider_api_format, + params.require_streaming, + is_kiro || is_antigravity || gemini_cli_requires_upstream_streaming, + ); + let force_body_stream_field = !operation_requires_sync + && aether_ai_formats::api_format_uses_body_stream_field(params.provider_api_format) + && aether_ai_formats::endpoint_config_forces_upstream_stream_policy( + transport.endpoint.config.as_ref(), + ); let report_kind = if is_kiro && !params.require_streaming { "claude_cli_sync_finalize" } else if (is_gemini_cli || is_antigravity) && !params.require_streaming { @@ -170,6 +194,8 @@ pub fn classify_same_format_provider_request_behavior( is_antigravity, is_gemini_cli, is_claude_code, + is_claude_code_transport, + anthropic_compatibility_profile, is_vertex, is_kiro, upstream_is_stream, @@ -196,6 +222,20 @@ pub fn build_same_format_provider_request_body_with_compatibility_report( }) } +pub fn enforce_same_format_provider_api_operation_body_policy( + body: &mut Value, + api_operation: Option, +) -> bool { + if !matches!( + api_operation, + Some(aether_ai_formats::ApiOperation::ClaudeCountTokens) + ) { + return false; + } + body.as_object_mut() + .is_some_and(|object| object.remove("stream").is_some()) +} + fn build_same_format_provider_request_body_inner( input: SameFormatProviderRequestBodyInput<'_>, mut compatibility_edits: Option<&mut Vec>, @@ -503,6 +543,7 @@ pub fn build_same_format_provider_upstream_url( upstream_is_stream: params.upstream_is_stream, request_query: params.request_query, kiro_api_region: params.kiro_api_region, + api_operation: params.api_operation, }, params.provider_request_body, ) @@ -526,14 +567,13 @@ pub fn build_same_format_provider_headers( let auth_header = input.auth_header.unwrap_or_default(); let auth_value = input.auth_value.unwrap_or_default(); - let mut provider_request_headers = if input.behavior.is_claude_code { + let mut provider_request_headers = if input.behavior.is_claude_code_transport { build_claude_code_passthrough_headers( input.headers, auth_header, auth_value, input.extra_headers, input.behavior.upstream_is_stream, - input.key_fingerprint, ) } else if input.behavior.is_vertex { build_complete_passthrough_headers( @@ -551,11 +591,8 @@ pub fn build_same_format_provider_headers( ) }; - let protected_headers = input - .auth_header - .filter(|value| !value.trim().is_empty()) - .map(|value| vec![value, "content-type"]) - .unwrap_or_else(|| vec!["content-type"]); + let mut protected_headers = upstream_credential_header_names().to_vec(); + protected_headers.push("content-type"); if !apply_local_header_rules_with_request_headers( &mut provider_request_headers, input.header_rules, @@ -567,10 +604,28 @@ pub fn build_same_format_provider_headers( return None; } if let (Some(auth_header), Some(auth_value)) = (input.auth_header, input.auth_value) { - ensure_upstream_auth_header(&mut provider_request_headers, auth_header, auth_value); + replace_upstream_auth_headers(&mut provider_request_headers, auth_header, auth_value); + } else { + replace_upstream_auth_headers(&mut provider_request_headers, "", ""); } - if input.behavior.upstream_is_stream { + let claude_code_profile = *current_claude_code_transport_identity_profile(); + if input.behavior.is_claude_code_transport { + claude_code_profile.apply_fixed_headers( + &mut provider_request_headers, + input.behavior.upstream_is_stream, + ); + } + if input.behavior.is_claude_code_transport || input.behavior.is_claude_code { + claude_code_profile.apply_beta_policy(&mut provider_request_headers, input.api_operation); + } + if matches!( + input.api_operation, + Some(aether_ai_formats::ApiOperation::ClaudeCountTokens) + ) { + provider_request_headers.insert("accept".to_string(), "application/json".to_string()); + } else if input.behavior.upstream_is_stream { provider_request_headers.insert("accept".to_string(), "text/event-stream".to_string()); + force_identity_accept_encoding(&mut provider_request_headers); } Some(provider_request_headers) } @@ -599,7 +654,7 @@ pub fn same_format_provider_transport_unsupported_reason( local_kiro_request_transport_unsupported_reason_with_network(transport) } else if behavior.is_antigravity { None - } else if behavior.is_claude_code { + } else if behavior.is_claude_code_transport { local_claude_code_transport_unsupported_reason_with_network(transport, api_format) } else if behavior.is_vertex { local_vertex_gemini_transport_unsupported_reason_with_network(transport) @@ -639,7 +694,7 @@ pub fn same_format_provider_transport_unsupported_reason_for_trace( }, ); if !behavior.is_antigravity - && !behavior.is_claude_code + && !behavior.is_claude_code_transport && !behavior.is_gemini_cli && !behavior.is_vertex && !behavior.is_kiro @@ -825,6 +880,183 @@ mod tests { assert_eq!(behavior.report_kind, "gemini_cli_sync_finalize"); } + #[test] + fn anthropic_compatibility_profile_controls_legacy_claude_code_edits() { + let mut native = sample_transport("custom"); + native.endpoint.api_format = "claude:messages".to_string(); + let native_behavior = classify_same_format_provider_request_behavior( + &native, + SameFormatProviderRequestBehaviorParams { + require_streaming: false, + provider_api_format: "claude:messages", + report_kind: "claude_chat_sync_success", + }, + ); + assert!(!native_behavior.is_claude_code); + assert!(!native_behavior.is_claude_code_transport); + assert_eq!( + native_behavior.anthropic_compatibility_profile, + AnthropicCompatibilityProfile::NativeTransparent + ); + + native.endpoint.config = Some(json!({ + "anthropic": {"compatibility_profile": "claude_code_legacy"} + })); + let compat_behavior = classify_same_format_provider_request_behavior( + &native, + SameFormatProviderRequestBehaviorParams { + require_streaming: false, + provider_api_format: "claude:messages", + report_kind: "claude_chat_sync_success", + }, + ); + assert!(compat_behavior.is_claude_code); + assert!(!compat_behavior.is_claude_code_transport); + assert_eq!( + compat_behavior.anthropic_compatibility_profile, + AnthropicCompatibilityProfile::ClaudeCodeLegacy + ); + + let request_body = json!({ + "model": "claude-client", + "thinking": {"type": "enabled"}, + "context_management": {"edits": [{"type": "client_strategy"}]}, + "system": [{ + "type": "text", + "text": "x-anthropic-billing-header: cc_version=9.9.9.abc; cc_entrypoint=cli;" + }], + "messages": [{ + "role": "assistant", + "content": [ + {"type": "text", "text": "visible"}, + {"type": "thinking", "thinking": "unsigned", "signature": ""} + ] + }] + }); + let build_body = |behavior: SameFormatProviderRequestBehavior| { + build_same_format_provider_request_body(SameFormatProviderRequestBodyInput { + body_json: &request_body, + mapped_model: "claude-upstream", + client_api_format: "claude:messages", + provider_api_format: "claude:messages", + source_model: Some("claude-client"), + family: SameFormatProviderFamily::Standard, + body_rules: None, + request_headers: None, + upstream_is_stream: false, + force_body_stream_field: false, + kiro_auth_config: None, + is_claude_code: behavior.is_claude_code, + enable_model_directives: false, + }) + .expect("body should build") + }; + assert_eq!( + build_body(native_behavior)["messages"][0]["content"] + .as_array() + .map(Vec::len), + Some(2), + "native Anthropic requests must remain untouched" + ); + assert_eq!( + build_body(native_behavior)["system"][0]["text"], + request_body["system"][0]["text"], + "native transparent requests must not rewrite billing identity" + ); + assert_eq!( + build_body(native_behavior)["context_management"], + request_body["context_management"], + "native transparent requests must not apply Claude Code body gates" + ); + assert_eq!( + build_body(compat_behavior)["messages"][0]["content"] + .as_array() + .map(Vec::len), + Some(1), + "legacy compatibility may sanitize invalid thinking blocks" + ); + assert_eq!( + build_body(compat_behavior)["system"][0]["text"], + "x-anthropic-billing-header: cc_version=2.1.161.abc; cc_entrypoint=cli;" + ); + + let mut legacy = sample_transport("claude_code"); + legacy.endpoint.api_format = "claude:messages".to_string(); + legacy.endpoint.config = Some(json!({ + "anthropic": {"compatibility_profile": "native_transparent"} + })); + let transparent_legacy_transport = classify_same_format_provider_request_behavior( + &legacy, + SameFormatProviderRequestBehaviorParams { + require_streaming: false, + provider_api_format: "claude:messages", + report_kind: "claude_chat_sync_success", + }, + ); + assert!(!transparent_legacy_transport.is_claude_code); + assert!(transparent_legacy_transport.is_claude_code_transport); + + let provider_request_body = json!({"model": "claude-upstream"}); + let empty_headers = http::HeaderMap::new(); + let empty_extra_headers = BTreeMap::new(); + let build_headers = + |behavior: SameFormatProviderRequestBehavior, + api_operation: Option| { + build_same_format_provider_headers(SameFormatProviderHeadersInput { + headers: &empty_headers, + provider_request_body: &provider_request_body, + original_request_body: &request_body, + header_rules: None, + behavior, + api_operation, + auth_header: Some("x-api-key"), + auth_value: Some("upstream-secret"), + extra_headers: &empty_extra_headers, + kiro_auth_config: None, + kiro_machine_id: None, + }) + .expect("headers should build") + }; + assert!( + build_headers(compat_behavior, None).get("x-app").is_none(), + "compatibility profile must not impersonate the Claude Code transport" + ); + assert_eq!( + build_headers(transparent_legacy_transport, None) + .get("x-app") + .map(String::as_str), + Some("cli"), + "Claude Code transport headers must survive a transparent body profile" + ); + assert!( + build_headers( + native_behavior, + Some(aether_ai_formats::ApiOperation::ClaudeCountTokens), + ) + .get("anthropic-beta") + .is_none(), + "native transparent token counting must not inject compatibility betas" + ); + assert!( + build_headers( + compat_behavior, + Some(aether_ai_formats::ApiOperation::ClaudeCountTokens), + )["anthropic-beta"] + .split(',') + .any(|token| token == "token-counting-2024-11-01"), + "legacy compatibility token counting requires the token-counting beta" + ); + assert!( + build_headers( + transparent_legacy_transport, + Some(aether_ai_formats::ApiOperation::ClaudeCountTokens), + )["anthropic-beta"] + .split(',') + .any(|token| token == "token-counting-2024-11-01"), + "Claude Code transport token counting requires the token-counting beta" + ); + } + #[test] fn same_format_behavior_resolves_endpoint_stream_policy() { let mut force_stream = sample_transport("openai"); @@ -908,6 +1140,53 @@ mod tests { assert!(!search_behavior.force_body_stream_field); } + #[test] + fn count_tokens_overrides_endpoint_stream_policy_and_removes_stream_field() { + let mut transport = sample_transport("custom"); + transport.endpoint.api_format = "claude:messages".to_string(); + transport.endpoint.config = Some(json!({ + "upstream_stream_policy": "force_stream" + })); + let behavior = classify_same_format_provider_request_behavior_for_operation( + &transport, + SameFormatProviderRequestBehaviorParams { + require_streaming: false, + provider_api_format: "claude:messages", + report_kind: "claude_count_tokens_sync_success", + }, + Some(aether_ai_formats::ApiOperation::ClaudeCountTokens), + ); + + assert!(!behavior.upstream_is_stream); + assert!(!behavior.force_body_stream_field); + + let mut body = json!({"model": "claude-sonnet-4", "messages": [], "stream": true}); + assert!(enforce_same_format_provider_api_operation_body_policy( + &mut body, + Some(aether_ai_formats::ApiOperation::ClaudeCountTokens), + )); + assert!(body.get("stream").is_none()); + + let headers = build_same_format_provider_headers(SameFormatProviderHeadersInput { + headers: &http::HeaderMap::new(), + provider_request_body: &body, + original_request_body: &body, + header_rules: None, + behavior, + api_operation: Some(aether_ai_formats::ApiOperation::ClaudeCountTokens), + auth_header: Some("x-api-key"), + auth_value: Some("secret"), + extra_headers: &BTreeMap::new(), + kiro_auth_config: None, + kiro_machine_id: None, + }) + .expect("headers should build"); + assert_eq!( + headers.get("accept").map(String::as_str), + Some("application/json") + ); + } + #[test] fn same_format_behavior_preserves_hard_streaming_constraint() { let mut kiro = sample_transport("kiro"); @@ -1857,8 +2136,13 @@ mod tests { fn builds_same_format_headers_with_auth_and_stream_accept() { let provider_request_body = json!({"model": "upstream-model"}); let original_request_body = json!({"model": "client-model"}); + let mut request_headers = http::HeaderMap::new(); + request_headers.insert( + http::header::ACCEPT_ENCODING, + http::HeaderValue::from_static("gzip, br"), + ); let headers = build_same_format_provider_headers(SameFormatProviderHeadersInput { - headers: &http::HeaderMap::new(), + headers: &request_headers, provider_request_body: &provider_request_body, original_request_body: &original_request_body, header_rules: None, @@ -1866,16 +2150,18 @@ mod tests { is_antigravity: false, is_gemini_cli: false, is_claude_code: false, + is_claude_code_transport: false, + anthropic_compatibility_profile: AnthropicCompatibilityProfile::NativeTransparent, is_vertex: false, is_kiro: false, upstream_is_stream: true, force_body_stream_field: false, report_kind: "openai_chat_stream_success", }, + api_operation: None, auth_header: Some("x-api-key"), auth_value: Some("secret"), extra_headers: &BTreeMap::new(), - key_fingerprint: None, kiro_auth_config: None, kiro_machine_id: None, }) @@ -1890,5 +2176,89 @@ mod tests { headers.get("accept").map(String::as_str), Some("text/event-stream") ); + assert_eq!( + headers.get("accept-encoding").map(String::as_str), + Some("identity") + ); + } + + #[test] + fn same_format_headers_cannot_restore_credentials_or_internal_headers() { + let provider_request_body = json!({"model": "upstream-model"}); + let original_request_body = json!({"model": "client-model"}); + let mut request_headers = http::HeaderMap::new(); + for (name, value) in [ + ("authorization", "Bearer client"), + ("api-key", "client-api-key"), + ("x-api-key", "client-x-api-key"), + ("cookie", "session=client"), + ("proxy-authorization", "Basic client"), + ("x-aether-auth-user-id", "user-private"), + ("x-aether-auth-api-key-id", "key-private"), + ("x-aether-auth-balance-remaining", "12.34"), + ("x-aether-gateway", "gateway-internal"), + ] { + request_headers.insert( + http::HeaderName::from_bytes(name.as_bytes()).expect("valid header name"), + http::HeaderValue::from_str(value).expect("valid header value"), + ); + } + let behavior = SameFormatProviderRequestBehavior { + is_antigravity: false, + is_gemini_cli: false, + is_claude_code: false, + is_claude_code_transport: false, + anthropic_compatibility_profile: AnthropicCompatibilityProfile::NativeTransparent, + is_vertex: false, + is_kiro: false, + upstream_is_stream: false, + force_body_stream_field: false, + report_kind: "claude_chat_sync_success", + }; + let header_rules = json!([ + {"action": "set", "key": "cookie", "value": "session=rule"}, + {"action": "set", "key": "authorization", "value": "Bearer rule"}, + {"action": "set", "key": "x-aether-auth-user-id", "value": "user-rule"} + ]); + let extra_headers = BTreeMap::from([ + ("api-key".to_string(), "extra-api-key".to_string()), + ("authorization".to_string(), "Bearer extra".to_string()), + ( + "x-aether-auth-balance-remaining".to_string(), + "99.99".to_string(), + ), + ("anthropic-beta".to_string(), "custom-beta".to_string()), + ]); + + let headers = build_same_format_provider_headers(SameFormatProviderHeadersInput { + headers: &request_headers, + provider_request_body: &provider_request_body, + original_request_body: &original_request_body, + header_rules: Some(&header_rules), + behavior, + api_operation: None, + auth_header: Some("x-api-key"), + auth_value: Some("upstream-secret"), + extra_headers: &extra_headers, + kiro_auth_config: None, + kiro_machine_id: None, + }) + .expect("headers should build"); + + assert_eq!( + headers.get("x-api-key").map(String::as_str), + Some("upstream-secret") + ); + for stripped in ["authorization", "api-key", "cookie", "proxy-authorization"] { + assert!(!headers.contains_key(stripped), "should strip {stripped}"); + } + assert!( + headers.keys().all(|name| !name.starts_with("x-aether-")), + "Aether-owned headers must never leave provider egress: {headers:?}" + ); + assert_eq!( + headers.get("anthropic-beta").map(String::as_str), + Some("custom-beta") + ); } } diff --git a/crates/aether-provider/transport/src/url.rs b/crates/aether-provider/transport/src/url.rs index 07af0239b..9fd9f2a06 100644 --- a/crates/aether-provider/transport/src/url.rs +++ b/crates/aether-provider/transport/src/url.rs @@ -3,6 +3,25 @@ use std::collections::BTreeMap; use url::form_urlencoded; use url::Url; +pub(crate) const GATEWAY_CREDENTIAL_QUERY_KEYS: &[&str] = &["key"]; + +pub(crate) fn strip_gateway_credential_query_parameters(query: Option<&str>) -> Option { + let query = query.map(str::trim).filter(|value| !value.is_empty())?; + let mut serializer = form_urlencoded::Serializer::new(String::new()); + let mut retained = false; + for (key, value) in form_urlencoded::parse(query.as_bytes()) { + if GATEWAY_CREDENTIAL_QUERY_KEYS + .iter() + .any(|blocked| key.eq_ignore_ascii_case(blocked)) + { + continue; + } + serializer.append_pair(&key, &value); + retained = true; + } + retained.then(|| serializer.finish()) +} + pub fn build_openai_chat_url(upstream_base_url: &str, query: Option<&str>) -> String { let (trimmed, base_query) = split_base_url_query(upstream_base_url); let trimmed = trimmed.trim_end_matches('/'); @@ -72,10 +91,61 @@ fn openai_image_base_includes_operation_path(base_url: &str) -> bool { } pub fn build_claude_messages_url(upstream_base_url: &str, query: Option<&str>) -> String { + build_claude_messages_operation_url(upstream_base_url, "", query) +} + +pub(crate) fn build_claude_count_tokens_url( + upstream_base_url: &str, + query: Option<&str>, +) -> String { + build_claude_messages_operation_url(upstream_base_url, "/count_tokens", query) +} + +fn build_claude_messages_operation_url( + upstream_base_url: &str, + operation_suffix: &str, + query: Option<&str>, +) -> String { let (trimmed, base_query) = split_base_url_query(upstream_base_url); let trimmed = trimmed.trim_end_matches('/'); - let mut url = format!("{trimmed}/messages"); - append_merged_query(&mut url, base_query, None, query, &[]); + let parsed_url = Url::parse(trimmed).ok(); + let parsed_path = parsed_url + .as_ref() + .map(|url| url.path().trim_matches('/')) + .unwrap_or_default(); + let base_includes_count_tokens = + parsed_path == "messages/count_tokens" || parsed_path.ends_with("/messages/count_tokens"); + let messages_base = if base_includes_count_tokens { + trimmed + .strip_suffix("/count_tokens") + .unwrap_or(trimmed) + .to_string() + } else if parsed_path.is_empty() + && parsed_url + .as_ref() + .and_then(Url::host_str) + .is_some_and(|host| host.eq_ignore_ascii_case("api.anthropic.com")) + { + format!("{trimmed}/v1/messages") + } else if parsed_path.is_empty() { + format!("{trimmed}/messages") + } else if parsed_path.rsplit('/').next() == Some("messages") { + trimmed.to_string() + } else { + format!("{trimmed}/messages") + }; + let mut url = if base_includes_count_tokens && operation_suffix == "/count_tokens" { + trimmed.to_string() + } else { + format!("{messages_base}{operation_suffix}") + }; + append_merged_query( + &mut url, + base_query, + None, + query, + GATEWAY_CREDENTIAL_QUERY_KEYS, + ); url } @@ -358,9 +428,10 @@ fn append_merged_query( base_query: Option<&str>, path_query: Option<&str>, request_query: Option<&str>, - blocked_keys: &[&str], + blocked_request_keys: &[&str], ) { - let Some(query) = merge_query_layers(base_query, path_query, request_query, blocked_keys) + let Some(query) = + merge_query_layers(base_query, path_query, request_query, blocked_request_keys) else { return; }; @@ -376,9 +447,9 @@ fn merge_query_layers( base_query: Option<&str>, path_query: Option<&str>, request_query: Option<&str>, - blocked_keys: &[&str], + blocked_request_keys: &[&str], ) -> Option { - if blocked_keys.is_empty() + if blocked_request_keys.is_empty() && path_query.is_none() && base_query.is_none() && request_query @@ -392,9 +463,9 @@ fn merge_query_layers( } let mut merged = BTreeMap::new(); - for source in [base_query, path_query, request_query] { - merge_query_string(&mut merged, source, blocked_keys); - } + merge_query_string(&mut merged, base_query, &[]); + merge_query_string(&mut merged, path_query, &[]); + merge_query_string(&mut merged, request_query, blocked_request_keys); if merged.is_empty() { return None; } @@ -429,11 +500,11 @@ fn merge_query_string( #[cfg(test)] mod tests { use super::{ - build_bigmodel_coding_models_url, build_claude_messages_url, build_gemini_content_url, - build_gemini_files_passthrough_url, build_gemini_video_predict_long_running_url, - build_openai_chat_url, build_openai_compatible_models_url, build_openai_image_url, - build_openai_responses_url, build_openai_search_url, build_passthrough_path_url, - normalize_gemini_content_action_path, + build_bigmodel_coding_models_url, build_claude_count_tokens_url, build_claude_messages_url, + build_gemini_content_url, build_gemini_files_passthrough_url, + build_gemini_video_predict_long_running_url, build_openai_chat_url, + build_openai_compatible_models_url, build_openai_image_url, build_openai_responses_url, + build_openai_search_url, build_passthrough_path_url, normalize_gemini_content_action_path, }; #[test] @@ -534,6 +605,54 @@ mod tests { build_claude_messages_url("https://api.anthropic.example", None), "https://api.anthropic.example/messages" ); + assert_eq!( + build_claude_messages_url("https://api.anthropic.com", None), + "https://api.anthropic.com/v1/messages" + ); + assert_eq!( + build_claude_messages_url( + "https://proxy.example.com/anthropic?key=base-secret&tenant=base", + Some("KEY=request-secret&trace=1") + ), + "https://proxy.example.com/anthropic/messages?key=base-secret&tenant=base&trace=1" + ); + } + + #[test] + fn claude_count_tokens_url_preserves_configured_query_and_strips_request_credentials() { + assert_eq!( + build_claude_count_tokens_url( + "https://proxy.example.com/anthropic?key=base-secret&tenant=base", + Some("key=request-secret&trace=1") + ), + "https://proxy.example.com/anthropic/messages/count_tokens?key=base-secret&tenant=base&trace=1" + ); + assert_eq!( + build_claude_count_tokens_url("https://api.anthropic.com", None), + "https://api.anthropic.com/v1/messages/count_tokens" + ); + assert_eq!( + build_claude_count_tokens_url("https://api.anthropic.example", None), + "https://api.anthropic.example/messages/count_tokens" + ); + assert_eq!( + build_claude_count_tokens_url( + "https://proxy.example.com/anthropic/messages", + Some("trace=1") + ), + "https://proxy.example.com/anthropic/messages/count_tokens?trace=1" + ); + assert_eq!( + build_claude_count_tokens_url( + "https://proxy.example.com/v1/messages/count_tokens?key=base-secret", + Some("key=request-secret&trace=1") + ), + "https://proxy.example.com/v1/messages/count_tokens?key=base-secret&trace=1" + ); + assert_eq!( + build_claude_messages_url("https://proxy.example.com/v1/messages/count_tokens", None), + "https://proxy.example.com/v1/messages" + ); } #[test] diff --git a/crates/aether-provider/transport/src/vertex/auth.rs b/crates/aether-provider/transport/src/vertex/auth.rs index 3ac6fe206..e6fcb9b28 100644 --- a/crates/aether-provider/transport/src/vertex/auth.rs +++ b/crates/aether-provider/transport/src/vertex/auth.rs @@ -9,7 +9,7 @@ use rsa::pkcs8::DecodePrivateKey; use rsa::signature::{SignatureEncoding, Signer}; use rsa::RsaPrivateKey; use serde_json::{json, Value}; -use sha2::Sha256; +use sha2::{Digest, Sha256}; use url::form_urlencoded; use super::super::oauth_refresh::{ @@ -153,7 +153,7 @@ impl LocalOAuthRefreshAdapter for VertexServiceAccountRefreshAdapter { fn resolve_cached( &self, - _transport: &GatewayProviderTransportSnapshot, + transport: &GatewayProviderTransportSnapshot, entry: &CachedOAuthEntry, ) -> Option { if !entry @@ -162,6 +162,9 @@ impl LocalOAuthRefreshAdapter for VertexServiceAccountRefreshAdapter { { return None; } + if !vertex_service_account_cached_entry_matches_transport(transport, entry) { + return None; + } if service_account_token_expires_soon(entry.expires_at_unix_secs) { return None; } @@ -194,6 +197,34 @@ impl LocalOAuthRefreshAdapter for VertexServiceAccountRefreshAdapter { .is_none() } + fn refresh_fingerprint( + &self, + transport: &GatewayProviderTransportSnapshot, + entry: Option<&CachedOAuthEntry>, + ) -> Option { + let source_fingerprint = vertex_service_account_credential_fingerprint(transport)?; + Some( + entry + .filter(|entry| { + vertex_service_account_cached_entry_matches_transport(transport, entry) + }) + .map(|entry| { + let mut digest = Sha256::new(); + digest.update(source_fingerprint.as_bytes()); + digest.update([0]); + digest.update(entry.auth_header_value.as_bytes()); + digest.update([0]); + digest.update(entry.expires_at_unix_secs.unwrap_or_default().to_be_bytes()); + format!("{:x}", digest.finalize()) + }) + .unwrap_or(source_fingerprint), + ) + } + + fn shares_refresh_through_transport_persistence(&self) -> bool { + false + } + async fn refresh( &self, executor: &dyn LocalOAuthHttpExecutor, @@ -259,11 +290,48 @@ impl LocalOAuthRefreshAdapter for VertexServiceAccountRefreshAdapter { "project_id": auth_config.project_id, "client_email": auth_config.client_email, })), - source_fingerprint: None, + source_fingerprint: vertex_service_account_credential_fingerprint(transport), })) } } +fn vertex_service_account_credential_fingerprint( + transport: &GatewayProviderTransportSnapshot, +) -> Option { + supports_local_vertex_service_account_auth_resolution(transport).then(|| { + let provider_type = transport.provider.provider_type.trim().to_ascii_lowercase(); + let auth_type = transport.key.auth_type.trim().to_ascii_lowercase(); + let auth_config = transport + .key + .decrypted_auth_config + .as_deref() + .unwrap_or_default(); + let mut digest = Sha256::new(); + for field in [ + provider_type.as_bytes(), + auth_type.as_bytes(), + auth_config.as_bytes(), + transport.key.decrypted_api_key.as_bytes(), + ] { + digest.update((field.len() as u64).to_be_bytes()); + digest.update(field); + } + format!("{:x}", digest.finalize()) + }) +} + +fn vertex_service_account_cached_entry_matches_transport( + transport: &GatewayProviderTransportSnapshot, + entry: &CachedOAuthEntry, +) -> bool { + entry + .provider_type + .eq_ignore_ascii_case(VERTEX_SERVICE_ACCOUNT_PROVIDER_TYPE) + && vertex_service_account_credential_fingerprint(transport) + .as_deref() + .is_some_and(|fingerprint| entry.source_fingerprint.as_deref() == Some(fingerprint)) +} + pub fn build_vertex_service_account_assertion( auth_config: &VertexServiceAccountAuthConfig, now_unix_secs: u64, @@ -324,6 +392,9 @@ fn body_excerpt(value: &str) -> String { #[cfg(test)] mod tests { + use super::super::super::oauth_refresh::{ + CachedOAuthEntry, LocalOAuthRefreshAdapter, LocalResolvedOAuthRequestAuth, + }; use super::super::super::snapshot::{ GatewayProviderTransportEndpoint, GatewayProviderTransportKey, GatewayProviderTransportProvider, GatewayProviderTransportSnapshot, @@ -335,7 +406,10 @@ mod tests { use super::{ decode_vertex_service_account_private_key, parse_vertex_service_account_auth_config, resolve_local_vertex_api_key_query_auth, - supports_local_vertex_service_account_auth_resolution, VERTEX_API_KEY_QUERY_PARAM, + supports_local_vertex_service_account_auth_resolution, + vertex_service_account_credential_fingerprint, VertexServiceAccountRefreshAdapter, + VERTEX_API_KEY_QUERY_PARAM, VERTEX_SERVICE_ACCOUNT_AUTH_HEADER, + VERTEX_SERVICE_ACCOUNT_PROVIDER_TYPE, }; fn sample_transport() -> GatewayProviderTransportSnapshot { @@ -395,6 +469,59 @@ mod tests { } } + fn sample_service_account_transport(private_key: &str) -> GatewayProviderTransportSnapshot { + let mut transport = sample_transport(); + transport.key.auth_type = "service_account".to_string(); + transport.key.decrypted_api_key = "__placeholder__".to_string(); + transport.key.decrypted_auth_config = Some( + serde_json::json!({ + "client_email": "svc@example.iam.gserviceaccount.com", + "private_key": private_key, + "project_id": "demo-project" + }) + .to_string(), + ); + transport + } + + #[test] + fn cached_token_is_bound_to_vertex_service_account_generation() { + let source_transport = sample_service_account_transport("SOURCE-PRIVATE-KEY"); + let source_fingerprint = vertex_service_account_credential_fingerprint(&source_transport) + .expect("source service account should have a fingerprint"); + let entry = CachedOAuthEntry { + provider_type: VERTEX_SERVICE_ACCOUNT_PROVIDER_TYPE.to_string(), + auth_header_name: VERTEX_SERVICE_ACCOUNT_AUTH_HEADER.to_string(), + auth_header_value: "Bearer source-access-token".to_string(), + expires_at_unix_secs: Some(u64::MAX), + metadata: None, + source_fingerprint: Some(source_fingerprint), + }; + let adapter = VertexServiceAccountRefreshAdapter; + + assert!(!adapter.shares_refresh_through_transport_persistence()); + assert_eq!( + adapter.resolve_cached(&source_transport, &entry), + Some(LocalResolvedOAuthRequestAuth::Header { + name: VERTEX_SERVICE_ACCOUNT_AUTH_HEADER.to_string(), + value: "Bearer source-access-token".to_string(), + }) + ); + assert_ne!( + adapter.refresh_fingerprint(&source_transport, None), + adapter.refresh_fingerprint(&source_transport, Some(&entry)) + ); + + let replacement_transport = sample_service_account_transport("ADMIN-PRIVATE-KEY"); + assert!(adapter + .resolve_cached(&replacement_transport, &entry) + .is_none()); + assert_eq!( + adapter.refresh_fingerprint(&replacement_transport, Some(&entry)), + vertex_service_account_credential_fingerprint(&replacement_transport) + ); + } + #[test] fn resolves_query_auth_for_vertex_api_key_subset() { let auth = resolve_local_vertex_api_key_query_auth(&sample_transport()) diff --git a/crates/aether-routing-core/src/mutations.rs b/crates/aether-routing-core/src/mutations.rs index 7262a324c..de530ea75 100644 --- a/crates/aether-routing-core/src/mutations.rs +++ b/crates/aether-routing-core/src/mutations.rs @@ -8,15 +8,17 @@ use crate::actions::{RoutingHeaderPatch, RoutingJsonPatchOperation}; const RESERVED_HEADERS: &[&str] = &[ "authorization", + "proxy-authorization", "x-api-key", "api-key", + "x-goog-api-key", "cookie", + "cookie2", "set-cookie", - "x-aether-trace-id", - "x-aether-internal", - "x-aether-scheduler-group", ]; +const AETHER_INTERNAL_HEADER_PREFIX: &str = "x-aether-"; + #[derive(Debug, Error, Clone, PartialEq, Eq)] pub enum MutationError { #[error("json patch path must be an absolute JSON pointer: {0}")] @@ -96,7 +98,7 @@ pub fn validate_header_patch(patch: &[RoutingHeaderPatch]) -> Result<(), Mutatio { return Err(MutationError::InvalidHeaderName(item.name().to_string())); } - if reserved.contains(name.as_str()) { + if reserved.contains(name.as_str()) || name.starts_with(AETHER_INTERNAL_HEADER_PREFIX) { return Err(MutationError::ReservedHeader(item.name().to_string())); } } @@ -224,13 +226,27 @@ mod tests { } #[test] - fn rejects_reserved_headers() { - assert_eq!( - validate_header_patch(&[RoutingHeaderPatch::Set { - name: "authorization".to_string(), - value: "secret".to_string() - }]), - Err(MutationError::ReservedHeader("authorization".to_string())) - ); + fn rejects_upstream_credentials_and_aether_internal_headers() { + for name in [ + "authorization", + "proxy-authorization", + "api-key", + "x-api-key", + "x-goog-api-key", + "cookie", + "cookie2", + "set-cookie", + "x-aether-auth-user-id", + "X-Aether-Future-Control", + ] { + assert_eq!( + validate_header_patch(&[RoutingHeaderPatch::Set { + name: name.to_string(), + value: "secret".to_string() + }]), + Err(MutationError::ReservedHeader(name.to_string())), + "header should be reserved: {name}" + ); + } } } diff --git a/frontend/src/features/providers/components/endpoint-default-paths.ts b/frontend/src/features/providers/components/endpoint-default-paths.ts index f49f3ef9f..ca0f7402c 100644 --- a/frontend/src/features/providers/components/endpoint-default-paths.ts +++ b/frontend/src/features/providers/components/endpoint-default-paths.ts @@ -187,9 +187,6 @@ export function getDefaultEndpointPath(params: { if (normalizedApiFormat === 'gemini:embedding') { return '/v1/projects/{project_id}/locations/{region}/publishers/google/models/{model}:predict' } - if (normalizedApiFormat === 'claude:messages') { - return '/v1/projects/{project_id}/locations/{region}/publishers/anthropic/models/{model}:{action}' - } } const format = params.apiFormats.find(f => f.value === normalizedApiFormat)