From 531cf110250579af2013baccea5bece4765a3442 Mon Sep 17 00:00:00 2001 From: elky Date: Mon, 27 Jul 2026 09:36:31 +0800 Subject: [PATCH] feat(gateway): harden provider request execution Preserve exact request payloads and model client surface and API operation explicitly. Add Anthropic compatibility profiles, bounded stream commitment, and scoped OAuth retry behavior across provider transports. --- Cargo.lock | 1 + apps/aether-gateway/src/ai_serving/api.rs | 40 +- .../src/ai_serving/finalize/tests_stream.rs | 1 + .../src/ai_serving/finalize/tests_sync.rs | 9 + apps/aether-gateway/src/ai_serving/mod.rs | 6 +- .../src/ai_serving/planner/antigravity.rs | 1 + .../planner/candidate_materialization.rs | 258 +- .../src/ai_serving/planner/common.rs | 2 +- .../planner/decision/control_plan.rs | 25 +- .../src/ai_serving/planner/decision_input.rs | 68 +- .../src/ai_serving/planner/gemini_cli.rs | 1 + .../passthrough/provider/family/candidates.rs | 8 +- .../passthrough/provider/family/payload.rs | 294 ++ .../passthrough/provider/family/request.rs | 129 +- .../provider/family/request/policy.rs | 4 +- .../provider/family/request/prepare.rs | 1 + .../planner/passthrough/provider/plans.rs | 20 + .../passthrough/provider/request/url.rs | 1 + .../src/ai_serving/planner/route.rs | 64 +- .../ai_serving/planner/specialized/files.rs | 20 + .../ai_serving/planner/specialized/image.rs | 20 + .../ai_serving/planner/specialized/video.rs | 10 + .../planner/standard/family/build.rs | 20 + .../planner/standard/family/payload.rs | 2 + .../planner/standard/gemini/plan_builders.rs | 20 +- .../src/ai_serving/planner/standard/mod.rs | 1 + .../standard/openai/chat/decision/request.rs | 2 + .../standard/openai/chat/plans/stream.rs | 14 + .../standard/openai/chat/plans/sync.rs | 10 + .../standard/openai/responses/plans.rs | 20 + .../planner/standard/plan_builders.rs | 102 +- .../aether-gateway/src/ai_serving/pure/mod.rs | 51 +- .../src/ai_serving/transport.rs | 37 +- apps/aether-gateway/src/api/ai/registry.rs | 32 +- apps/aether-gateway/src/api/response.rs | 211 +- .../src/control/auth/credentials.rs | 14 + apps/aether-gateway/src/control/auth/mod.rs | 2 + apps/aether-gateway/src/control/auth/types.rs | 22 +- apps/aether-gateway/src/control/mod.rs | 2 +- apps/aether-gateway/src/control/route/ai.rs | 65 +- apps/aether-gateway/src/control/route/mod.rs | 66 +- apps/aether-gateway/src/control/tests/ai.rs | 80 +- .../src/dispatch/pool_scheduler.rs | 4 + .../src/execution_runtime/grok.rs | 8 +- .../src/execution_runtime/mod.rs | 81 +- .../src/execution_runtime/oauth_retry.rs | 388 ++- .../src/execution_runtime/server.rs | 2 + .../execution_runtime/stream/commit_policy.rs | 527 +++ .../src/execution_runtime/stream/execution.rs | 3076 +++++++++++++++-- .../stream/execution_failures.rs | 39 +- .../src/execution_runtime/stream/mod.rs | 5 +- .../src/execution_runtime/stream_pump.rs | 163 +- .../src/execution_runtime/submission.rs | 80 + .../src/execution_runtime/sync/execution.rs | 311 +- .../sync/execution/policy.rs | 53 +- .../src/execution_runtime/sync/mod.rs | 1 + .../src/execution_runtime/tests.rs | 2 +- .../src/execution_runtime/transport.rs | 557 ++- .../src/executor/candidate_loop.rs | 258 +- apps/aether-gateway/src/executor/mod.rs | 3 +- .../src/executor/orchestration.rs | 16 +- apps/aether-gateway/src/executor/outcome.rs | 72 +- .../src/executor/stream_path.rs | 26 +- apps/aether-gateway/src/executor/sync_path.rs | 27 +- apps/aether-gateway/src/fallback_metrics.rs | 3 + .../provider/oauth/dispatch/complete/key.rs | 1 + .../admin/provider/oauth/quota/shared.rs | 1 + .../handlers/admin/provider/oauth/runtime.rs | 49 +- .../admin/provider/query/models/model_test.rs | 2 + .../admin/provider/write/normalize.rs | 20 +- .../admin/provider/write/provider/create.rs | 2 + .../admin/provider/write/provider/update.rs | 4 + .../admin/request/provider/builders.rs | 8 + .../handlers/admin/request/system/import.rs | 51 + apps/aether-gateway/src/handlers/proxy/mod.rs | 29 +- .../src/handlers/public/ai_public.rs | 222 +- .../public/support/test_connection/route.rs | 1 + .../runtime/fixed_provider_reconciliation.rs | 127 +- .../src/orchestration/classifier.rs | 357 +- .../src/orchestration/effects.rs | 518 ++- apps/aether-gateway/src/orchestration/mod.rs | 18 +- .../src/orchestration/oauth_error.rs | 113 + .../src/orchestration/recovery.rs | 81 +- apps/aether-gateway/src/rate_limit.rs | 3 + apps/aether-gateway/src/state/oauth.rs | 488 ++- .../src/tests/ai_execute/fallback.rs | 8 +- .../ai_execute/finalize_local_cli/direct.rs | 1 + .../finalize_local_provider/claude.rs | 1 + .../src/tests/ai_execute/stream_provider.rs | 27 +- .../ai_execute/sync/claude/claude_code.rs | 10 +- .../src/tests/ai_execute/sync/claude/kiro.rs | 2 + .../ai_execute/sync/claude/local_chat.rs | 199 +- .../tests/ai_execute/sync/claude/local_cli.rs | 6 +- apps/aether-gateway/src/tests/concurrency.rs | 3 + .../tests/control/admin/endpoints/quota.rs | 2 +- .../src/tests/control/admin/security.rs | 39 + .../src/tests/control/admin/system_import.rs | 138 + .../src/tests/control/proxy/local_denials.rs | 107 +- apps/aether-gateway/src/tests/frontdoor/ai.rs | 130 - apps/aether-gateway/src/tests/proxy.rs | 80 + apps/aether-gateway/src/tests/usage/local.rs | 10 +- crates/aether-ai/formats/src/api.rs | 9 +- crates/aether-ai/formats/src/contracts/mod.rs | 57 +- .../formats/src/contracts/plan_kinds.rs | 1 + .../formats/src/contracts/report_kinds.rs | 1 + .../src/contracts/request_dimensions.rs | 59 + .../formats/src/formats/shared/error_body.rs | 66 +- .../formats/src/formats/shared/passthrough.rs | 38 +- .../formats/src/formats/shared/routing.rs | 96 +- .../src/formats/shared/stream_rewrite.rs | 71 +- .../src/formats/shared/sync_products.rs | 142 +- .../src/formats/shared/sync_to_stream.rs | 8 +- crates/aether-ai/formats/src/lib.rs | 2 + crates/aether-ai/serving/Cargo.toml | 1 + crates/aether-ai/serving/src/attempt_loop.rs | 246 +- crates/aether-ai/serving/src/dto.rs | 22 + .../aether-ai/serving/src/execution_path.rs | 96 +- crates/aether-ai/serving/src/lib.rs | 7 +- .../aether-ai/serving/src/payload_fidelity.rs | 77 + crates/aether-contracts/src/lib.rs | 7 +- crates/aether-contracts/src/plan.rs | 43 + .../adapters/mysql/src/candidate_selection.rs | 38 +- .../adapters/mysql/src/provider_catalog.rs | 31 +- .../postgres/src/candidate_selection.rs | 101 +- .../adapters/postgres/src/provider_catalog.rs | 46 + .../sqlite/src/candidate_selection.rs | 71 +- .../adapters/sqlite/src/provider_catalog.rs | 60 +- .../src/repository/provider_catalog/mod.rs | 8 +- .../src/repository/provider_catalog/types.rs | 19 +- .../repository/candidate_selection/memory.rs | 59 +- .../src/repository/provider_catalog/memory.rs | 135 +- .../src/repository/provider_catalog/mod.rs | 5 +- .../transport/src/anthropic_compat.rs | 376 ++ crates/aether-provider/transport/src/auth.rs | 85 +- .../transport/src/claude_code/fingerprint.rs | 373 +- .../transport/src/claude_code/mod.rs | 11 +- .../transport/src/claude_code/profile.rs | 322 ++ .../transport/src/claude_code/request.rs | 310 +- .../transport/src/claude_code/url.rs | 28 +- .../transport/src/generic_oauth/mod.rs | 428 ++- .../aether-provider/transport/src/headers.rs | 130 +- .../transport/src/kiro/refresh.rs | 193 +- crates/aether-provider/transport/src/lib.rs | 10 +- .../aether-provider/transport/src/network.rs | 94 +- .../transport/src/oauth_refresh/mod.rs | 398 ++- .../transport/src/provider_types.rs | 21 +- .../transport/src/request_url/mod.rs | 581 +++- .../transport/src/same_format_provider/mod.rs | 436 ++- crates/aether-provider/transport/src/url.rs | 147 +- .../transport/src/vertex/auth.rs | 135 +- crates/aether-routing-core/src/mutations.rs | 40 +- .../components/endpoint-default-paths.ts | 3 - 152 files changed, 13984 insertions(+), 2075 deletions(-) create mode 100644 apps/aether-gateway/src/execution_runtime/stream/commit_policy.rs create mode 100644 apps/aether-gateway/src/orchestration/oauth_error.rs create mode 100644 crates/aether-ai/formats/src/contracts/request_dimensions.rs create mode 100644 crates/aether-ai/serving/src/payload_fidelity.rs create mode 100644 crates/aether-provider/transport/src/anthropic_compat.rs create mode 100644 crates/aether-provider/transport/src/claude_code/profile.rs 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)