diff --git a/.github/workflows/rust-ci.yml b/.github/workflows/rust-ci.yml index 67b127ccb..c498c8b23 100644 --- a/.github/workflows/rust-ci.yml +++ b/.github/workflows/rust-ci.yml @@ -383,6 +383,13 @@ jobs: AETHER_TEST_POSTGRES_URL: postgres://aether:aether@127.0.0.1:5432/aether_test run: cargo test -p aether-data postgres_migrations_create_core_config_tables_when_url_is_set --lib -- --nocapture + - name: Run Postgres provider metadata migration smoke test + env: + RUSTC_WRAPPER: sccache + SCCACHE_GHA_ENABLED: "true" + AETHER_TEST_POSTGRES_URL: postgres://aether:aether@127.0.0.1:5432/aether_test + run: cargo test -p aether-data postgres_provider_upstream_metadata_migration_preserves_json_when_url_is_set --lib -- --nocapture + - name: Run Postgres core export smoke test env: RUSTC_WRAPPER: sccache diff --git a/apps/aether-gateway/src/ai_serving/api.rs b/apps/aether-gateway/src/ai_serving/api.rs index c3b0f0e39..07d622d70 100644 --- a/apps/aether-gateway/src/ai_serving/api.rs +++ b/apps/aether-gateway/src/ai_serving/api.rs @@ -68,6 +68,7 @@ pub(crate) use aether_ai_formats::api::{ OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND, }; pub(crate) use aether_ai_formats::protocol::stream::CanonicalUsage as StreamingCanonicalUsage; +pub(crate) use aether_ai_formats::CODEX_RESPONSES_LITE_HEADER; pub(crate) fn parse_direct_request_body( parts: &http::request::Parts, 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 911459ed6..f3ebe2516 100644 --- a/apps/aether-gateway/src/ai_serving/finalize/tests_sync.rs +++ b/apps/aether-gateway/src/ai_serving/finalize/tests_sync.rs @@ -31,6 +31,7 @@ fn test_decision() -> GatewayControlDecision { auth_context: None, admin_principal: None, local_auth_rejection: None, + model_directive_policy: Default::default(), } } @@ -1166,7 +1167,7 @@ fn local_finalize_handles_openai_responses_compact_cross_format_sync_response() assert_eq!(report.report_kind, "openai_responses_compact_sync_success"); assert_eq!( report.client_body_json.expect("client body should exist")["object"], - "response" + "response.compaction" ); } @@ -1222,7 +1223,7 @@ fn local_finalize_handles_openai_responses_compact_cross_format_function_call_re .background_report .expect("compact tool-call should downgrade to success report"); let client_body = report.client_body_json.expect("client body should exist"); - assert_eq!(client_body["object"], "response"); + assert_eq!(client_body["object"], "response.compaction"); assert_eq!(client_body["output"][1]["type"], "function_call"); } @@ -1841,6 +1842,7 @@ fn local_finalize_handles_claude_chat_cross_format_sync_response_from_openai_cha auth_context: None, admin_principal: None, local_auth_rejection: None, + model_directive_policy: Default::default(), }, &payload, ) @@ -1908,6 +1910,7 @@ fn local_finalize_handles_gemini_cli_cross_format_sync_response_from_claude_cli( auth_context: None, admin_principal: None, local_auth_rejection: None, + model_directive_policy: Default::default(), }, &payload, ) diff --git a/apps/aether-gateway/src/ai_serving/mod.rs b/apps/aether-gateway/src/ai_serving/mod.rs index ce55cb81e..835e27f88 100644 --- a/apps/aether-gateway/src/ai_serving/mod.rs +++ b/apps/aether-gateway/src/ai_serving/mod.rs @@ -13,7 +13,9 @@ use crate::{usage::GatewaySyncReportRequest, AppState, GatewayError}; pub(crate) use self::adaptation::{ maybe_build_provider_private_stream_normalizer, ProviderPrivateStreamNormalizer, }; -pub(crate) use self::api::gemini_generate_content_response_has_visible_output; +pub(crate) use self::api::{ + gemini_generate_content_response_has_visible_output, CODEX_RESPONSES_LITE_HEADER, +}; pub(crate) use self::finalize::common::LocalCoreSyncFinalizeOutcome; pub(crate) use self::finalize::internal::{ maybe_bridge_standard_sync_json_to_stream, maybe_build_stream_response_rewriter, @@ -48,16 +50,17 @@ pub(crate) use self::planner::{ build_standard_family_stream_plan_and_reports, build_standard_family_sync_attempt_source, build_standard_family_sync_plan_and_reports, build_standard_stream_plan_from_decision, build_standard_sync_plan_from_decision, candidate_auth_channel_skip_reason, - extract_pool_sticky_session_token, maybe_build_stream_decision_payload, - maybe_build_stream_plan_payload, maybe_build_sync_decision_payload, - maybe_build_sync_plan_payload, planner_is_matching_stream_request, provider_key_pool_score_id, - provider_key_pool_score_scope, read_candidate_transport_snapshot, - record_local_runtime_candidate_skip_reason, - set_local_openai_chat_execution_exhausted_diagnostic, - set_local_openai_image_execution_exhausted_diagnostic, CandidateFailureDiagnostic, - CandidateFailureDiagnosticKind, EligibleLocalExecutionCandidate, GatewayAuthApiKeySnapshot, - GatewayProviderTransportSnapshot, LocalExecutionAttemptSource, LocalExecutionCandidateKind, - LocalResolvedOAuthRequestAuth, PlannerAppState, SkippedLocalExecutionCandidate, + codex_model_capabilities_for_transport, extract_pool_sticky_session_token, + maybe_build_stream_decision_payload, maybe_build_stream_plan_payload, + maybe_build_sync_decision_payload, maybe_build_sync_plan_payload, + planner_is_matching_stream_request, provider_key_pool_score_id, provider_key_pool_score_scope, + read_candidate_transport_snapshot, record_local_runtime_candidate_skip_reason, + resolve_upstream_is_stream_for_provider, set_local_openai_chat_execution_exhausted_diagnostic, + set_local_openai_image_execution_exhausted_diagnostic, validate_final_openai_provider_request, + CandidateFailureDiagnostic, CandidateFailureDiagnosticKind, EligibleLocalExecutionCandidate, + GatewayAuthApiKeySnapshot, GatewayProviderTransportSnapshot, LocalExecutionAttemptSource, + LocalExecutionCandidateKind, LocalResolvedOAuthRequestAuth, PlannerAppState, + SkippedLocalExecutionCandidate, }; pub(crate) use self::pure::*; pub(crate) use self::transport::{ 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 bdeaf7736..64fce324d 100644 --- a/apps/aether-gateway/src/ai_serving/planner/candidate_materialization.rs +++ b/apps/aether-gateway/src/ai_serving/planner/candidate_materialization.rs @@ -702,6 +702,7 @@ pub(crate) async fn build_lazy_requested_model_execution_candidate_attempt_sourc G, >( state: PlannerAppState<'a>, + model_directive_policy: &crate::system_features::ModelDirectivePolicySnapshot, trace_id: &str, client_api_format: &str, requested_model: &str, @@ -730,6 +731,7 @@ where let record_runtime_miss_diagnostic = persistence_policy.skipped.record_runtime_miss_diagnostic; let page_cursor = LocalCandidatePreselectionPageCursor::new( state, + model_directive_policy, client_api_format, requested_model, require_streaming, @@ -1151,6 +1153,9 @@ async fn resolve_priority_candidate_page_with_cache( .page_cursor .resolved_page_cache_use_api_format_alias_match(), cursor.client_session_affinity.as_ref(), + cursor + .page_cursor + .resolved_page_cache_model_directive_policy_hash(), cursor.resolution_mode, ); let page_candidates_for_fallback = page_candidates.clone(); @@ -1949,6 +1954,7 @@ mod tests { global_model_id: "global-model-1".to_string(), global_model_name: "gpt-5".to_string(), selected_provider_model_name: "gpt-5".to_string(), + supports_streaming: true, mapping_matched_model: None, } } @@ -2177,8 +2183,11 @@ mod tests { async fn resolved_candidate_page_cache_requires_fixed_order_or_explicit_affinity() { let app = AppState::new().expect("state should build"); let auth_snapshot = sample_auth_snapshot(); + let model_directive_policy = + crate::system_features::ModelDirectivePolicySnapshot::default(); let mut page_cursor = LocalCandidatePreselectionPageCursor::new( PlannerAppState::new(&app), + &model_directive_policy, "openai:chat", "gpt-5", true, @@ -2233,6 +2242,7 @@ mod tests { let mut page_cursor = LocalCandidatePreselectionPageCursor::new( PlannerAppState::new(&app), + &model_directive_policy, "openai:chat", "gpt-5", true, @@ -2269,6 +2279,7 @@ mod tests { ); let mut page_cursor = LocalCandidatePreselectionPageCursor::new( PlannerAppState::new(&fixed_order_app), + &model_directive_policy, "openai:chat", "gpt-5", true, diff --git a/apps/aether-gateway/src/ai_serving/planner/candidate_metadata.rs b/apps/aether-gateway/src/ai_serving/planner/candidate_metadata.rs index 4b5e97fcd..699941e8a 100644 --- a/apps/aether-gateway/src/ai_serving/planner/candidate_metadata.rs +++ b/apps/aether-gateway/src/ai_serving/planner/candidate_metadata.rs @@ -134,6 +134,7 @@ mod tests { global_model_id: "global-1".to_string(), global_model_name: "gpt-5.4".to_string(), selected_provider_model_name: "gpt-5.4".to_string(), + supports_streaming: true, mapping_matched_model: None, } } diff --git a/apps/aether-gateway/src/ai_serving/planner/candidate_preparation.rs b/apps/aether-gateway/src/ai_serving/planner/candidate_preparation.rs index 6c1124296..0cf791fc3 100644 --- a/apps/aether-gateway/src/ai_serving/planner/candidate_preparation.rs +++ b/apps/aether-gateway/src/ai_serving/planner/candidate_preparation.rs @@ -169,6 +169,7 @@ mod tests { global_model_id: "global-model-1".to_string(), global_model_name: "gpt-test".to_string(), selected_provider_model_name: "gpt-test-upstream".to_string(), + supports_streaming: true, mapping_matched_model: None, } } diff --git a/apps/aether-gateway/src/ai_serving/planner/candidate_ranking.rs b/apps/aether-gateway/src/ai_serving/planner/candidate_ranking.rs index 854ad130b..daaf1a7c1 100644 --- a/apps/aether-gateway/src/ai_serving/planner/candidate_ranking.rs +++ b/apps/aether-gateway/src/ai_serving/planner/candidate_ranking.rs @@ -348,6 +348,7 @@ mod tests { global_model_id: "global-model-1".to_string(), global_model_name: "gpt-4.1".to_string(), selected_provider_model_name: "gpt-4.1".to_string(), + supports_streaming: true, mapping_matched_model: None, } } @@ -624,6 +625,7 @@ mod tests { global_model_id: "global-model-1".to_string(), global_model_name: "gpt-4.1".to_string(), selected_provider_model_name: "gpt-4.1".to_string(), + supports_streaming: true, mapping_matched_model: None, } } diff --git a/apps/aether-gateway/src/ai_serving/planner/candidate_resolution.rs b/apps/aether-gateway/src/ai_serving/planner/candidate_resolution.rs index df1b27317..fdf7b00cf 100644 --- a/apps/aether-gateway/src/ai_serving/planner/candidate_resolution.rs +++ b/apps/aether-gateway/src/ai_serving/planner/candidate_resolution.rs @@ -643,6 +643,7 @@ mod tests { global_model_id: "global-model-1".to_string(), global_model_name: "claude-sonnet".to_string(), selected_provider_model_name: "claude-sonnet".to_string(), + supports_streaming: true, mapping_matched_model: None, } } diff --git a/apps/aether-gateway/src/ai_serving/planner/candidate_source.rs b/apps/aether-gateway/src/ai_serving/planner/candidate_source.rs index de876982e..1ff3fc292 100644 --- a/apps/aether-gateway/src/ai_serving/planner/candidate_source.rs +++ b/apps/aether-gateway/src/ai_serving/planner/candidate_source.rs @@ -64,10 +64,25 @@ struct GatewayLocalCandidatePreselectionPort<'a> { use_api_format_alias_match: bool, key_mode: LocalCandidatePreselectionKeyMode, candidate_api_formats: Vec, - model_directive_enabled_api_formats: BTreeSet, + model_directive_routing_models: BTreeMap, ranking_seed: u64, } +impl GatewayLocalCandidatePreselectionPort<'_> { + fn model_directive_base_model(&self, candidate_api_format: &str) -> Option<&str> { + self.model_directive_routing_models + .get(&crate::ai_serving::normalize_api_format_alias( + candidate_api_format, + )) + .map(String::as_str) + } + + fn routing_model(&self, candidate_api_format: &str) -> &str { + self.model_directive_base_model(candidate_api_format) + .unwrap_or(self.requested_model) + } +} + #[async_trait] impl AiCandidatePreselectionPort for GatewayLocalCandidatePreselectionPort<'_> { type Candidate = SchedulerMinimalCandidateSelectionCandidate; @@ -99,12 +114,13 @@ impl AiCandidatePreselectionPort for GatewayLocalCandidatePreselectionPort<'_> { .state .list_selectable_candidates_with_skip_reasons( candidate_api_format, - self.requested_model, + self.routing_model(candidate_api_format), self.require_streaming, self.required_capabilities, auth_snapshot, self.client_session_affinity, self.ranking_seed, + false, ) .await?; @@ -123,16 +139,13 @@ impl AiCandidatePreselectionPort for GatewayLocalCandidatePreselectionPort<'_> { candidate_api_format: &str, matches_client_format: bool, ) -> bool { - let enable_model_directives = self.model_directive_enabled_api_formats.contains( - &crate::ai_serving::normalize_api_format_alias(candidate_api_format), - ); routing_policy_allows_provider(self.routing_policy, candidate) && (matches_client_format || auth_snapshot_allows_cross_format_candidate( self.auth_snapshot, self.requested_model, + self.model_directive_base_model(candidate_api_format), candidate, - enable_model_directives, )) } @@ -142,16 +155,13 @@ impl AiCandidatePreselectionPort for GatewayLocalCandidatePreselectionPort<'_> { candidate_api_format: &str, matches_client_format: bool, ) -> bool { - let enable_model_directives = self.model_directive_enabled_api_formats.contains( - &crate::ai_serving::normalize_api_format_alias(candidate_api_format), - ); routing_policy_allows_provider(self.routing_policy, &skipped_candidate.candidate) && (matches_client_format || auth_snapshot_allows_cross_format_candidate( self.auth_snapshot, self.requested_model, + self.model_directive_base_model(candidate_api_format), &skipped_candidate.candidate, - enable_model_directives, )) } @@ -164,9 +174,27 @@ impl AiCandidatePreselectionPort for GatewayLocalCandidatePreselectionPort<'_> { } } +fn resolve_model_directive_routing_models( + policy: &crate::system_features::ModelDirectivePolicySnapshot, + candidate_api_formats: &[String], + requested_model: &str, +) -> BTreeMap { + candidate_api_formats + .iter() + .filter_map(|api_format| { + let api_format = crate::ai_serving::normalize_api_format_alias(api_format); + let resolution = policy.resolve_reasoning(&api_format, Some(requested_model)); + resolution + .base_model() + .map(|base_model| (api_format, base_model.to_string())) + }) + .collect() +} + #[allow(clippy::too_many_arguments)] pub(crate) async fn preselect_local_execution_candidates_with_serving( state: PlannerAppState<'_>, + model_directive_policy: &crate::system_features::ModelDirectivePolicySnapshot, client_api_format: &str, requested_model: &str, require_streaming: bool, @@ -190,6 +218,7 @@ pub(crate) async fn preselect_local_execution_candidates_with_serving( .collect::>(); preselect_local_execution_candidates_for_api_formats_with_serving( state, + model_directive_policy, client_api_format, requested_model, require_streaming, @@ -207,6 +236,7 @@ pub(crate) async fn preselect_local_execution_candidates_with_serving( #[allow(clippy::too_many_arguments)] pub(crate) async fn preselect_local_execution_candidates_for_api_formats_with_serving( state: PlannerAppState<'_>, + model_directive_policy: &crate::system_features::ModelDirectivePolicySnapshot, client_api_format: &str, requested_model: &str, require_streaming: bool, @@ -224,19 +254,11 @@ pub(crate) async fn preselect_local_execution_candidates_for_api_formats_with_se >, GatewayError, > { - let mut model_directive_enabled_api_formats = BTreeSet::new(); - for api_format in &candidate_api_formats { - if crate::system_features::reasoning_model_directive_enabled_for_api_format_and_model( - state.app(), - api_format, - Some(requested_model), - ) - .await - { - model_directive_enabled_api_formats - .insert(crate::ai_serving::normalize_api_format_alias(api_format)); - } - } + let model_directive_routing_models = resolve_model_directive_routing_models( + model_directive_policy, + &candidate_api_formats, + requested_model, + ); let port = GatewayLocalCandidatePreselectionPort { state, client_api_format, @@ -249,7 +271,7 @@ pub(crate) async fn preselect_local_execution_candidates_for_api_formats_with_se use_api_format_alias_match, key_mode, candidate_api_formats, - model_directive_enabled_api_formats, + model_directive_routing_models, ranking_seed: request_distribution_seed(), }; @@ -271,7 +293,8 @@ pub(crate) struct LocalCandidatePreselectionPageCursor<'a> { key_mode: LocalCandidatePreselectionKeyMode, allow_priority_page_cache: bool, candidate_api_formats: Vec, - model_directive_enabled_api_formats: BTreeSet, + model_directive_routing_models: BTreeMap, + model_directive_policy_cache_key: String, ordering_config: SchedulerOrderingConfig, ranking_seed: u64, priority_page_emitted: bool, @@ -294,9 +317,23 @@ pub(crate) struct LocalCandidatePreselectionPageCursor<'a> { } impl<'a> LocalCandidatePreselectionPageCursor<'a> { + fn model_directive_base_model(&self, candidate_api_format: &str) -> Option<&str> { + self.model_directive_routing_models + .get(&crate::ai_serving::normalize_api_format_alias( + candidate_api_format, + )) + .map(String::as_str) + } + + fn routing_model(&self, candidate_api_format: &str) -> &str { + self.model_directive_base_model(candidate_api_format) + .unwrap_or(&self.requested_model) + } + #[allow(clippy::too_many_arguments)] pub(crate) async fn new( state: PlannerAppState<'a>, + model_directive_policy: &crate::system_features::ModelDirectivePolicySnapshot, client_api_format: &str, requested_model: &str, require_streaming: bool, @@ -315,19 +352,11 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> { .into_iter() .map(str::to_string) .collect::>(); - let mut model_directive_enabled_api_formats = BTreeSet::new(); - for api_format in &candidate_api_formats { - if crate::system_features::reasoning_model_directive_enabled_for_api_format_and_model( - state.app(), - api_format, - Some(requested_model), - ) - .await - { - model_directive_enabled_api_formats - .insert(crate::ai_serving::normalize_api_format_alias(api_format)); - } - } + let model_directive_routing_models = resolve_model_directive_routing_models( + model_directive_policy, + &candidate_api_formats, + requested_model, + ); let ordering_config = super::candidate_ranking::scheduler_ordering_config_for_routing_policy( @@ -351,7 +380,8 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> { key_mode, allow_priority_page_cache, candidate_api_formats, - model_directive_enabled_api_formats, + model_directive_routing_models, + model_directive_policy_cache_key: model_directive_policy.cache_key().to_string(), ordering_config, ranking_seed: request_distribution_seed(), priority_page_emitted: false, @@ -426,6 +456,10 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> { self.use_api_format_alias_match } + pub(crate) fn resolved_page_cache_model_directive_policy_hash(&self) -> &str { + &self.model_directive_policy_cache_key + } + pub(crate) fn should_cache_current_priority_resolved_page(&self) -> bool { if !(self.priority_page_emitted && self.format_index == 0 @@ -491,6 +525,7 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> { self.key_mode.cache_key_name(), self.use_api_format_alias_match, self.client_session_affinity.as_ref(), + &self.model_directive_policy_cache_key, ); let cache = self.state.app().candidate_page_cache.clone(); let ttl = candidate_page_cache_ttl_from_env(); @@ -752,11 +787,8 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> { if normalized_api_format.is_empty() { return Ok(None); } - let enable_model_directives = self.model_directive_enabled_api_formats.contains( - &crate::ai_serving::normalize_api_format_alias(candidate_api_format), - ); - let requested_names = - requested_model_candidate_names(&self.requested_model, enable_model_directives); + let routing_model = self.routing_model(candidate_api_format).to_string(); + let requested_names = requested_model_candidate_names(&routing_model, false); let scanned = *self .scanned_rows_by_format .get(&normalized_api_format) @@ -772,11 +804,7 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> { .or_insert(0); let Some(requested_name) = requested_names.get(requested_name_index) else { return self - .next_fallback_page_for_api_format( - candidate_api_format, - &normalized_api_format, - enable_model_directives, - ) + .next_fallback_page_for_api_format(candidate_api_format, &normalized_api_format) .await; }; if requested_name.trim().is_empty() { @@ -803,9 +831,9 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> { .read_requested_model_rows_fast_path_page_cached( &normalized_api_format, requested_name, + &routing_model, offset, limit, - enable_model_directives, ) .await?; self.scanned_rows_by_format.insert( @@ -824,7 +852,6 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> { .next_fallback_page_for_api_format( candidate_api_format, &normalized_api_format, - enable_model_directives, ) .await; } @@ -835,7 +862,6 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> { .build_page_outcome_from_rows( candidate_api_format, &normalized_api_format, - enable_model_directives, page.rows, ) .await? @@ -849,17 +875,17 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> { &self, normalized_api_format: &str, requested_name: &str, + routing_model: &str, offset: u32, limit: u32, - enable_model_directives: bool, ) -> Result { let key = CandidateRowPageCacheKey::new( normalized_api_format, - &self.requested_model, + routing_model, requested_name, offset, limit, - enable_model_directives, + false, ); let cache = self.state.app().candidate_row_page_cache.clone(); let ttl = candidate_page_cache_ttl_from_env(); @@ -873,11 +899,11 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> { let page = read_requested_model_rows_fast_path_page( self.state.app().data.as_ref(), normalized_api_format, - &self.requested_model, + routing_model, requested_name, offset, limit, - enable_model_directives, + false, ) .await .map_err(|err| GatewayError::Internal(err.to_string()))?; @@ -913,7 +939,6 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> { &mut self, candidate_api_format: &str, normalized_api_format: &str, - enable_model_directives: bool, ) -> Result< Option< AiCandidatePreselectionOutcome< @@ -930,6 +955,7 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> { return Ok(None); } + let routing_model = self.routing_model(candidate_api_format).to_string(); let rows = self .state .app() @@ -941,27 +967,21 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> { .filter(|row| { row_supports_requested_model_with_model_directives( row, - &self.requested_model, + &routing_model, normalized_api_format, - enable_model_directives, + false, ) }) .collect::>(); - self.build_page_outcome_from_rows( - candidate_api_format, - normalized_api_format, - enable_model_directives, - rows, - ) - .await + self.build_page_outcome_from_rows(candidate_api_format, normalized_api_format, rows) + .await } async fn build_page_outcome_from_rows( &mut self, candidate_api_format: &str, normalized_api_format: &str, - enable_model_directives: bool, rows: Vec, ) -> Result< Option< @@ -984,15 +1004,16 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> { if rows.is_empty() { return Ok(None); } + let routing_model = self.routing_model(candidate_api_format).to_string(); let resolved_global_model_name = if let Some(value) = self.resolved_global_model_names.get(normalized_api_format) { value.clone() } else { let Some(value) = resolve_requested_global_model_name_with_model_directives( &rows, - &self.requested_model, + &routing_model, normalized_api_format, - enable_model_directives, + false, ) else { return Ok(None); }; @@ -1016,22 +1037,18 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> { EnumerateMinimalCandidateSelectionInput { rows, normalized_api_format, - requested_model_name: &self.requested_model, + requested_model_name: &routing_model, resolved_global_model_name: resolved_global_model_name.as_str(), require_streaming: self.require_streaming, required_capabilities: self.required_capabilities.as_ref(), auth_constraints: auth_constraints.as_ref(), }, - enable_model_directives, + false, ) .map_err(|err| GatewayError::Internal(err.to_string()))?; let mut candidates = Vec::new(); for candidate in enumerated_candidates { - if !self.candidate_allowed_for_page( - &candidate, - candidate_api_format, - enable_model_directives, - ) { + if !self.candidate_allowed_for_page(&candidate, candidate_api_format) { continue; } if !self @@ -1065,11 +1082,7 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> { .into_iter() .map(skipped_local_execution_candidate_from_scheduler_skip) .filter(|skipped_candidate| { - self.skipped_candidate_allowed_for_page( - skipped_candidate, - candidate_api_format, - enable_model_directives, - ) + self.skipped_candidate_allowed_for_page(skipped_candidate, candidate_api_format) }) .collect::>(); @@ -1083,7 +1096,6 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> { &self, candidate: &SchedulerMinimalCandidateSelectionCandidate, candidate_api_format: &str, - enable_model_directives: bool, ) -> bool { routing_policy_allows_provider(self.routing_policy.as_ref(), candidate) && (matches_client_api_format( @@ -1093,8 +1105,8 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> { ) || auth_snapshot_allows_cross_format_candidate( &self.auth_snapshot, &self.requested_model, + self.model_directive_base_model(candidate_api_format), candidate, - enable_model_directives, )) } @@ -1102,7 +1114,6 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> { &self, skipped_candidate: &SkippedLocalExecutionCandidate, candidate_api_format: &str, - enable_model_directives: bool, ) -> bool { routing_policy_allows_provider(self.routing_policy.as_ref(), &skipped_candidate.candidate) && (matches_client_api_format( @@ -1112,8 +1123,8 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> { ) || auth_snapshot_allows_cross_format_candidate( &self.auth_snapshot, &self.requested_model, + self.model_directive_base_model(candidate_api_format), &skipped_candidate.candidate, - enable_model_directives, )) } } @@ -1199,8 +1210,8 @@ fn matches_client_api_format( pub(crate) fn auth_snapshot_allows_cross_format_candidate( auth_snapshot: &GatewayAuthApiKeySnapshot, requested_model: &str, + requested_base_model: Option<&str>, candidate: &SchedulerMinimalCandidateSelectionCandidate, - enable_model_directives: bool, ) -> bool { if let Some(allowed_providers) = auth_snapshot.effective_allowed_providers() { let provider_allowed = allowed_providers.iter().any(|value| { @@ -1217,15 +1228,10 @@ pub(crate) fn auth_snapshot_allows_cross_format_candidate( } if let Some(allowed_models) = auth_snapshot.effective_allowed_models() { - let requested_base_model = enable_model_directives - .then(|| crate::ai_serving::model_directive_base_model(requested_model)) - .flatten(); let model_allowed = allowed_models.iter().any(|value| { value == requested_model || value == &candidate.global_model_name - || requested_base_model - .as_ref() - .is_some_and(|base_model| value == base_model) + || requested_base_model.is_some_and(|base_model| value == base_model) }); if !model_allowed { return false; @@ -1303,8 +1309,11 @@ mod tests { .expect("gateway state should build") .with_data_state_for_tests(data_state); let auth_snapshot = unrestricted_auth_snapshot(); + let model_directive_policy = + crate::system_features::ModelDirectivePolicySnapshot::load(&app).await; let mut cursor = LocalCandidatePreselectionPageCursor::new( PlannerAppState::new(&app), + &model_directive_policy, "openai:chat", "gpt-5", true, @@ -1545,8 +1554,11 @@ mod tests { .expect("gateway state should build") .with_data_state_for_tests(data_state); let auth_snapshot = unrestricted_auth_snapshot(); + let model_directive_policy = + crate::system_features::ModelDirectivePolicySnapshot::load(&app).await; let mut cursor = LocalCandidatePreselectionPageCursor::new( PlannerAppState::new(&app), + &model_directive_policy, "claude:messages", "gpt-5.5-xhigh", false, @@ -1578,6 +1590,77 @@ mod tests { ); } + #[tokio::test] + async fn custom_policy_suffix_uses_the_same_base_model_for_candidate_selection() { + let mut row = openai_responses_mapping_row(); + row.global_model_name = "deployment-alias".to_string(); + row.global_model_mappings = None; + row.model_provider_model_name = "gpt-5.6-sol".to_string(); + let repository: Arc = + Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed([row])); + let data_state = + GatewayDataState::with_minimal_candidate_selection_reader_for_tests(repository) + .with_system_config_values_for_tests([ + ( + crate::system_features::ENABLE_MODEL_DIRECTIVES_CONFIG_KEY.to_string(), + serde_json::json!(true), + ), + ( + crate::system_features::MODEL_DIRECTIVES_CONFIG_KEY.to_string(), + serde_json::json!({ + "reasoning_effort": { + "api_formats": { + "openai:responses": { + "suffixes": ["VendorFuture"], + "mappings": { + "VendorFuture": { + "reasoning": { "context": "all_turns" } + } + } + } + } + } + }), + ), + ]); + let app = AppState::new() + .expect("gateway state should build") + .with_data_state_for_tests(data_state); + let auth_snapshot = unrestricted_auth_snapshot(); + let model_directive_policy = + crate::system_features::ModelDirectivePolicySnapshot::load(&app).await; + let mut cursor = LocalCandidatePreselectionPageCursor::new( + PlannerAppState::new(&app), + &model_directive_policy, + "openai:responses", + "deployment-alias-VendorFuture", + false, + None, + &auth_snapshot, + None, + None, + None, + true, + LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModelAndApiFormat, + true, + None, + ) + .await; + + let page = cursor + .next_page() + .await + .expect("preselection should succeed") + .expect("custom directive base model should resolve a candidate"); + + assert_eq!(page.candidates.len(), 1); + assert_eq!(page.candidates[0].global_model_name, "deployment-alias"); + assert_eq!( + page.candidates[0].selected_provider_model_name, + "gpt-5.6-sol" + ); + } + #[tokio::test] async fn claude_request_uses_cross_format_key_when_same_provider_messages_key_lacks_model() { let repository: Arc = @@ -1605,8 +1688,11 @@ mod tests { .expect("gateway state should build") .with_data_state_for_tests(data_state); let auth_snapshot = unrestricted_auth_snapshot(); + let model_directive_policy = + crate::system_features::ModelDirectivePolicySnapshot::load(&app).await; let mut cursor = LocalCandidatePreselectionPageCursor::new( PlannerAppState::new(&app), + &model_directive_policy, "claude:messages", "deepseek-v4-pro", false, @@ -1680,8 +1766,11 @@ mod tests { .expect("gateway state should build") .with_data_state_for_tests(data_state); let auth_snapshot = unrestricted_auth_snapshot(); + let model_directive_policy = + crate::system_features::ModelDirectivePolicySnapshot::load(&app).await; let mut cursor = LocalCandidatePreselectionPageCursor::new( PlannerAppState::new(&app), + &model_directive_policy, "claude:messages", "gpt-5", false, diff --git a/apps/aether-gateway/src/ai_serving/planner/common.rs b/apps/aether-gateway/src/ai_serving/planner/common.rs index d05791de5..33565a583 100644 --- a/apps/aether-gateway/src/ai_serving/planner/common.rs +++ b/apps/aether-gateway/src/ai_serving/planner/common.rs @@ -1,11 +1,12 @@ use axum::body::Bytes; +use crate::ai_serving::is_json_request; use crate::ai_serving::{ endpoint_config_forces_upstream_stream_policy as endpoint_config_forces_upstream_stream_policy_impl, enforce_request_body_stream_field as enforce_request_body_stream_field_impl, force_upstream_streaming_for_provider as force_upstream_streaming_for_provider_impl, - is_json_request, parse_direct_request_body as parse_direct_request_body_impl, - resolve_upstream_is_stream_from_endpoint_config as resolve_upstream_is_stream_from_endpoint_config_impl, + parse_direct_request_body as parse_direct_request_body_impl, + resolve_format_upstream_is_stream_for_provider as resolve_upstream_is_stream_for_provider_impl, }; pub(crate) use crate::ai_serving::{ CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND, @@ -20,9 +21,10 @@ pub(crate) use crate::ai_serving::{ 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_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, + 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(crate) use aether_ai_serving::AiRequestedModelFamily as RequestedModelFamily; @@ -54,10 +56,10 @@ pub(crate) fn resolve_upstream_is_stream_for_provider( client_is_stream: bool, hard_requires_streaming: bool, ) -> bool { - let hard_requires_streaming = hard_requires_streaming - || force_upstream_streaming_for_provider(provider_type, provider_api_format); - resolve_upstream_is_stream_from_endpoint_config_impl( + resolve_upstream_is_stream_for_provider_impl( endpoint_config, + provider_type, + provider_api_format, client_is_stream, hard_requires_streaming, ) @@ -178,6 +180,27 @@ mod tests { true, false, )); + assert!(!resolve_upstream_is_stream_for_provider( + Some(&json!({"upstream_stream_policy": "force_stream"})), + "codex", + "openai:image", + true, + true, + )); + assert!(!resolve_upstream_is_stream_for_provider( + Some(&json!({"upstream_stream_policy": "force_stream"})), + "codex", + "openai:responses:compact", + true, + true, + )); + assert!(!resolve_upstream_is_stream_for_provider( + Some(&json!({"upstream_stream_policy": "force_stream"})), + "custom", + "openai:responses:compact", + true, + true, + )); } #[test] 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 f681f8558..276df03c0 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 @@ -8,9 +8,10 @@ use crate::ai_serving::planner::common::{ 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_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, + 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, @@ -100,7 +101,9 @@ fn build_sync_plan_payload_from_decision( OPENAI_RESPONSES_SYNC_PLAN_KIND => { build_openai_responses_sync_plan_from_decision(parts, body_json, payload, false)? } - OPENAI_IMAGE_SYNC_PLAN_KIND => build_passthrough_sync_plan_from_decision(parts, payload)?, + OPENAI_IMAGE_SYNC_PLAN_KIND | OPENAI_SEARCH_SYNC_PLAN_KIND => { + build_passthrough_sync_plan_from_decision(parts, payload)? + } OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND => { build_openai_responses_sync_plan_from_decision(parts, body_json, payload, true)? } 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 a65dc9b2a..5f5d1d7f8 100644 --- a/apps/aether-gateway/src/ai_serving/planner/decision_input.rs +++ b/apps/aether-gateway/src/ai_serving/planner/decision_input.rs @@ -14,8 +14,11 @@ use serde_json::{json, Value}; use tracing::warn; use crate::ai_serving::planner::common::extract_standard_requested_model; -use crate::ai_serving::{ExecutionRuntimeAuthContext, GatewayAuthApiKeySnapshot, PlannerAppState}; -use crate::client_session_affinity::client_session_affinity_from_request; +use crate::ai_serving::{ + ExecutionRuntimeAuthContext, GatewayAuthApiKeySnapshot, GatewayProviderTransportSnapshot, + PlannerAppState, CODEX_RESPONSES_LITE_HEADER, +}; +use crate::client_session_affinity::client_session_affinity_from_api_request; use crate::clock::current_unix_secs; use crate::routing::{ apply_routing_mutation_plan, build_routing_trace_seed, resolve_gateway_routing_policy, @@ -27,12 +30,15 @@ use crate::stage_metrics::observe_gateway_stage_ms; use crate::{AiExecutionDecision, AppState, GatewayError}; const ROUTING_GROUP_SELECTION_CACHE_TTL: Duration = Duration::from_secs(30); +const CODEX_ACCOUNT_ID_HEADER: &str = "chatgpt-account-id"; +const CODEX_FEDRAMP_HEADER: &str = "x-openai-fedramp"; #[derive(Debug, Clone)] pub(crate) struct ResolvedLocalDecisionAuthInput { pub(crate) auth_context: ExecutionRuntimeAuthContext, pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot, pub(crate) required_capabilities: Option, + pub(crate) model_directive_policy: crate::system_features::ModelDirectivePolicySnapshot, } #[derive(Debug, Clone)] @@ -46,6 +52,7 @@ pub(crate) struct LocalRequestedModelDecisionInput { pub(crate) routing_policy: Option, pub(crate) routing_trace_seed: Option, pub(crate) routing_context: Option, + pub(crate) model_directive_policy: crate::system_features::ModelDirectivePolicySnapshot, } #[derive(Debug, Clone)] @@ -86,14 +93,52 @@ impl LocalRequestedModelDecisionInput { pub(crate) fn apply_provider_request_routing_policy_to_decision( input: &LocalRequestedModelDecisionInput, decision: &mut AiExecutionDecision, + transport: Option<&GatewayProviderTransportSnapshot>, ) -> Result<(), GatewayError> { + let provider_api_format = decision + .provider_api_format + .clone() + .or_else(|| { + input + .routing_context + .as_ref() + .map(|context| context.client_api_format.clone()) + }) + .unwrap_or_default(); + let provider_type = decision.provider_type.clone().unwrap_or_default(); + let terminal_provider_model = decision + .provider_request_body + .as_ref() + .and_then(|body| body.get("model")) + .and_then(Value::as_str) + .or(decision.mapped_model.as_deref()) + .or(decision.model_name.as_deref()) + .unwrap_or(input.requested_model.as_str()); + let model_capabilities = transport.and_then(|transport| { + crate::ai_serving::codex_model_capabilities_for_transport( + transport, + provider_api_format.as_str(), + terminal_provider_model, + input.requested_model.as_str(), + ) + }); + crate::ai_serving::apply_codex_openai_responses_lite_header_with_capabilities( + &mut decision.provider_request_headers, + provider_type.as_str(), + provider_api_format.as_str(), + terminal_provider_model, + input.requested_model.as_str(), + model_capabilities.as_ref(), + ); + let Some(context) = input.routing_context.as_ref() else { return Ok(()); }; - let provider_api_format = decision - .provider_api_format - .as_deref() - .unwrap_or(context.client_api_format.as_str()); + let provider_body_rules = decision + .report_context + .as_ref() + .and_then(|context| context.get("body_rules")) + .cloned(); let resolved_model = decision .mapped_model .as_deref() @@ -104,6 +149,21 @@ pub(crate) fn apply_provider_request_routing_policy_to_decision( .clone() .unwrap_or(serde_json::Value::Null); let mut provider_headers = btree_headers_to_header_map(&decision.provider_request_headers)?; + let mut protected_codex_header_names = vec![CODEX_ACCOUNT_ID_HEADER, CODEX_FEDRAMP_HEADER]; + if provider_type.eq_ignore_ascii_case("codex") + && crate::ai_serving::is_openai_responses_family_format(provider_api_format.as_str()) + { + protected_codex_header_names.extend([ + "x-client-request-id", + "accept", + "content-encoding", + CODEX_RESPONSES_LITE_HEADER, + ]); + } + let protected_codex_headers = protected_codex_header_names + .into_iter() + .map(|name| (name, provider_headers.get(name).cloned())) + .collect::>(); let provider_headers_json = headers_to_routing_value(&provider_headers); let policy = resolve_gateway_routing_policy(GatewayRoutingPolicyInput { group_id: context.group_id.as_deref(), @@ -112,7 +172,7 @@ pub(crate) fn apply_provider_request_routing_policy_to_decision( selection_source: context.selection_source.as_str(), requested_model: input.requested_model.as_str(), resolved_model, - api_format: provider_api_format, + api_format: provider_api_format.as_str(), user_id: Some(input.auth_context.user_id.as_str()), api_key_id: Some(input.auth_context.api_key_id.as_str()), headers: &provider_headers_json, @@ -134,6 +194,80 @@ pub(crate) fn apply_provider_request_routing_policy_to_decision( &mut provider_headers, &policy.mutation_plan, )?; + for (name, value) in protected_codex_headers { + provider_headers.remove(name); + if let Some(value) = value { + provider_headers.insert(HeaderName::from_static(name), value); + } + } + if original_provider_request_body.is_some() { + let provider_model = provider_request_body + .get("model") + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .or(decision.mapped_model.as_deref()) + .or(decision.model_name.as_deref()) + .unwrap_or(input.requested_model.as_str()) + .to_string(); + let model_capabilities = transport.and_then(|transport| { + crate::ai_serving::codex_model_capabilities_for_transport( + transport, + provider_api_format.as_str(), + provider_model.as_str(), + input.requested_model.as_str(), + ) + }); + crate::ai_serving::finalize_openai_provider_request_with_codex_model_capabilities( + &mut provider_request_body, + crate::ai_serving::OpenAiProviderRequestFinalization { + source_api_format: context.client_api_format.as_str(), + provider_api_format: provider_api_format.as_str(), + provider_type: provider_type.as_str(), + provider_model: provider_model.as_str(), + source_model: input.requested_model.as_str(), + body_rules: provider_body_rules.as_ref(), + upstream_is_stream: decision.upstream_is_stream, + require_body_stream_field: original_provider_request_body + .as_ref() + .is_some_and(|body| body.get("stream").is_some()), + }, + model_capabilities.as_ref(), + ) + .map_err(|violation| GatewayError::Client { + status: StatusCode::BAD_REQUEST, + message: format!("routing provider_request violates provider contract: {violation:?}"), + })?; + } + let provider_model = provider_request_body + .get("model") + .and_then(Value::as_str) + .or(decision.mapped_model.as_deref()) + .or(decision.model_name.as_deref()) + .unwrap_or(input.requested_model.as_str()); + let mut provider_request_headers = header_map_to_btree_headers(&provider_headers); + let model_capabilities = transport.and_then(|transport| { + crate::ai_serving::codex_model_capabilities_for_transport( + transport, + provider_api_format.as_str(), + provider_model, + input.requested_model.as_str(), + ) + }); + crate::ai_serving::apply_codex_openai_responses_lite_header_with_capabilities( + &mut provider_request_headers, + provider_type.as_str(), + provider_api_format.as_str(), + provider_model, + input.requested_model.as_str(), + model_capabilities.as_ref(), + ); + crate::ai_serving::apply_codex_openai_compact_terminal_headers( + &mut provider_request_headers, + provider_type.as_str(), + provider_api_format.as_str(), + ); + provider_headers = btree_headers_to_header_map(&provider_request_headers)?; decision.provider_request_headers = header_map_to_btree_headers(&provider_headers); if original_provider_request_body.is_some() { decision.provider_request_body = Some(provider_request_body); @@ -145,6 +279,8 @@ pub(crate) fn apply_provider_request_routing_policy_to_decision( struct GatewayAuthenticatedDecisionInputPort<'a> { state: PlannerAppState<'a>, now_unix_secs: u64, + model_directive_policy: &'a crate::system_features::ModelDirectivePolicySnapshot, + model_directive_base_model: Option, } #[async_trait] @@ -181,6 +317,7 @@ impl AiAuthenticatedDecisionInputPort for GatewayAuthenticatedDecisionInputPort< &auth_context.api_key_id, requested_model, explicit_required_capabilities, + self.model_directive_base_model.as_deref(), ) .await) } @@ -195,6 +332,7 @@ impl AiAuthenticatedDecisionInputPort for GatewayAuthenticatedDecisionInputPort< auth_context, auth_snapshot, required_capabilities, + model_directive_policy: self.model_directive_policy.clone(), } } } @@ -213,6 +351,7 @@ pub(crate) fn build_local_requested_model_decision_input( routing_policy: None, routing_trace_seed: None, routing_context: None, + model_directive_policy: resolved_input.model_directive_policy, } } @@ -310,8 +449,11 @@ pub(crate) async fn attach_routing_policy_to_local_requested_model_input( let Some((group_id, group_version, group_config_json, selection_source)) = selected_group else { - input.client_session_affinity = - client_session_affinity_from_request(&parts.headers, Some(body_json)); + input.client_session_affinity = client_session_affinity_from_api_request( + client_api_format, + &parts.headers, + Some(body_json), + ); input.routing_policy = None; input.routing_trace_seed = None; input.routing_context = None; @@ -373,19 +515,26 @@ pub(crate) async fn attach_routing_policy_to_local_requested_model_input( } } if requested_model_changed { + let model_directive_resolution = input + .model_directive_policy + .resolve_reasoning(client_api_format, Some(input.requested_model.as_str())); input.required_capabilities = PlannerAppState::new(state) .resolve_request_candidate_required_capabilities( &input.auth_context.user_id, &input.auth_context.api_key_id, Some(input.requested_model.as_str()), input.required_capabilities.as_ref(), + model_directive_resolution.base_model(), ) .await; } let effective_headers_json = headers_to_routing_value(&effective_headers); - input.client_session_affinity = - client_session_affinity_from_request(&effective_headers, Some(&effective_body_json)); + input.client_session_affinity = client_session_affinity_from_api_request( + client_api_format, + &effective_headers, + Some(&effective_body_json), + ); let final_policy_resolve_started_at = std::time::Instant::now(); let mut final_policy = resolve_gateway_routing_policy(GatewayRoutingPolicyInput { group_id: group_id.as_deref(), @@ -452,8 +601,11 @@ fn try_attach_static_default_routing_policy_to_input( static_policy_resolve_started_at.elapsed().as_millis() as u64, ); - input.client_session_affinity = - client_session_affinity_from_request(&parts.headers, Some(body_json)); + input.client_session_affinity = client_session_affinity_from_api_request( + client_api_format, + &parts.headers, + Some(body_json), + ); input.routing_trace_seed = Some(build_routing_trace_seed(&policy, client_api_format)); input.routing_policy = Some(policy); input.routing_context = None; @@ -475,11 +627,22 @@ pub(crate) async fn resolve_local_authenticated_decision_input( state: &AppState, auth_context: ExecutionRuntimeAuthContext, requested_model: Option<&str>, + requested_model_api_format: Option<&str>, explicit_required_capabilities: Option<&serde_json::Value>, + model_directive_policy: &crate::system_features::ModelDirectivePolicySnapshot, ) -> Result, GatewayError> { + let model_directive_base_model = match (requested_model, requested_model_api_format) { + (Some(model), Some(api_format)) => model_directive_policy + .resolve_reasoning(api_format, Some(model)) + .base_model() + .map(str::to_owned), + _ => None, + }; let port = GatewayAuthenticatedDecisionInputPort { state: PlannerAppState::new(state), now_unix_secs: current_unix_secs(), + model_directive_policy, + model_directive_base_model, }; run_ai_authenticated_decision_input( @@ -730,6 +893,10 @@ fn ensure_report_context_routing_trace( #[cfg(test)] mod tests { use super::*; + use aether_provider_transport::snapshot::{ + GatewayProviderTransportEndpoint, GatewayProviderTransportKey, + GatewayProviderTransportProvider, + }; fn sample_auth_context() -> ExecutionRuntimeAuthContext { ExecutionRuntimeAuthContext { @@ -782,6 +949,7 @@ mod tests { client_session_affinity: None, routing_policy: None, routing_trace_seed: None, + model_directive_policy: Default::default(), routing_context: Some(LocalRoutingRequestContext { group_id: Some("group-1".to_string()), group_version: Some(3), @@ -830,6 +998,7 @@ mod tests { request_id: Some("trace-1".to_string()), candidate_id: Some("candidate-1".to_string()), provider_name: Some("provider".to_string()), + provider_type: Some("openai".to_string()), provider_id: Some("provider-1".to_string()), endpoint_id: Some("endpoint-1".to_string()), key_id: Some("key-1".to_string()), @@ -869,9 +1038,80 @@ mod tests { } } - fn set_provider_request_rules(input: &mut LocalRequestedModelDecisionInput, actions: Value) { + fn sample_codex_transport_with_card() -> GatewayProviderTransportSnapshot { + let card = json!({ + "id": "gpt-future-agent", + "slug": "gpt-future-agent", + "use_responses_lite": true, + "supports_reasoning_summary_parameter": true, + "default_reasoning_level": "low", + "default_reasoning_summary": "none", + "supported_reasoning_levels": [{"effort": "low"}, {"effort": "high"}] + }); + GatewayProviderTransportSnapshot { + provider: GatewayProviderTransportProvider { + id: "provider-codex".to_string(), + name: "Codex".to_string(), + provider_type: "codex".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-codex".to_string(), + provider_id: "provider-codex".to_string(), + api_format: "openai:responses:compact".to_string(), + api_family: Some("openai".to_string()), + endpoint_kind: Some("compact".to_string()), + is_active: true, + base_url: "https://chatgpt.com/backend-api/codex".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-codex".to_string(), + provider_id: "provider-codex".to_string(), + name: "Codex key".to_string(), + auth_type: "oauth".to_string(), + is_active: true, + api_formats: Some(vec!["openai:responses:compact".to_string()]), + auth_type_by_format: None, + allow_auth_channel_mismatch_formats: None, + allowed_models: Some(vec!["gpt-future-agent".to_string()]), + capabilities: None, + rate_multipliers: None, + global_priority_by_format: None, + expires_at_unix_secs: None, + proxy: None, + fingerprint: None, + upstream_metadata: Some(crate::ai_serving::build_codex_model_catalog_metadata(&[ + card, + ])), + decrypted_api_key: "access-token".to_string(), + decrypted_auth_config: None, + }, + } + } + + fn set_provider_request_rules( + input: &mut LocalRequestedModelDecisionInput, + allowed_models: &[&str], + actions: Value, + ) { let config = json!({ - "allowed_models": ["gpt-5"], + "allowed_models": allowed_models, "rules": [{ "id": "provider-patch", "priority": 1, @@ -904,6 +1144,7 @@ mod tests { client_session_affinity: None, routing_policy: None, routing_trace_seed: None, + model_directive_policy: Default::default(), routing_context: Some(LocalRoutingRequestContext { group_id: Some("stale".to_string()), group_version: Some(1), @@ -971,6 +1212,7 @@ mod tests { routing_policy: None, routing_trace_seed: None, routing_context: None, + model_directive_policy: Default::default(), }; let group_config_json = json!({ "rules": [{ @@ -1006,7 +1248,7 @@ mod tests { let input = sample_decision_input(); let mut decision = sample_decision(); - apply_provider_request_routing_policy_to_decision(&input, &mut decision) + apply_provider_request_routing_policy_to_decision(&input, &mut decision, None) .expect("provider routing mutation should apply"); assert_eq!( @@ -1035,6 +1277,179 @@ mod tests { ); } + #[test] + fn codex_compact_contract_is_terminal_after_routing_mutations() { + let mut input = sample_decision_input(); + input + .routing_context + .as_mut() + .expect("routing context") + .client_api_format = "openai:responses:compact".to_string(); + set_provider_request_rules( + &mut input, + &["gpt-5"], + json!([ + { + "type": "json_patch_body", + "patch": [ + {"op": "add", "path": "/store", "value": true}, + {"op": "add", "path": "/top_logprobs", "value": 5}, + {"op": "add", "path": "/custom_extension", "value": true}, + {"op": "replace", "path": "/input", "value": "routed compact input"}, + {"op": "replace", "path": "/tools", "value": [{ + "type": "function", + "name": "lookup", + "cache_control": {"type": "ephemeral"} + }]} + ] + }, + { + "type": "patch_headers", + "patch": [ + {"op": "set", "name": "chatgpt-account-id", "value": "spoofed"}, + {"op": "set", "name": "x-openai-fedramp", "value": "false"}, + {"op": "set", "name": "x-client-request-id", "value": "spoofed"}, + {"op": "set", "name": "accept", "value": "text/event-stream"}, + {"op": "set", "name": "content-encoding", "value": "zstd"} + ] + } + ]), + ); + let mut decision = sample_decision(); + decision.provider_type = Some("codex".to_string()); + decision.provider_api_format = Some("openai:responses:compact".to_string()); + decision.client_api_format = Some("openai:responses:compact".to_string()); + decision.provider_request_body = Some(json!({ + "model": "gpt-5", + "input": [], + "tools": [{"type": "function", "name": "lookup"}] + })); + decision + .provider_request_headers + .insert(CODEX_ACCOUNT_ID_HEADER.to_string(), "account-1".to_string()); + decision + .provider_request_headers + .insert(CODEX_FEDRAMP_HEADER.to_string(), "true".to_string()); + + apply_provider_request_routing_policy_to_decision(&input, &mut decision, None) + .expect("terminal contract should accept the projected request"); + + let body = decision.provider_request_body.as_ref().expect("body"); + assert_eq!(body["parallel_tool_calls"], false); + assert_eq!(body["input"][0]["type"], "message"); + assert_eq!( + body["input"][0]["content"][0]["text"], + "routed compact input" + ); + assert!(body["tools"][0].get("cache_control").is_none()); + for field in ["store", "top_logprobs", "custom_extension"] { + assert!( + body.get(field).is_none(), + "unexpected Compact field: {field}" + ); + } + assert_eq!( + decision + .provider_request_headers + .get(CODEX_ACCOUNT_ID_HEADER), + Some(&"account-1".to_string()) + ); + assert_eq!( + decision.provider_request_headers.get(CODEX_FEDRAMP_HEADER), + Some(&"true".to_string()) + ); + for header in ["x-client-request-id", "accept", "content-encoding"] { + assert!( + !decision.provider_request_headers.contains_key(header), + "unexpected Compact header: {header}" + ); + } + } + + #[test] + fn codex_responses_lite_contract_is_terminal_after_routing_mutations() { + let mut input = sample_decision_input(); + input.requested_model = "gpt-future-agent".to_string(); + input + .routing_context + .as_mut() + .expect("routing context") + .client_api_format = "openai:responses:compact".to_string(); + set_provider_request_rules( + &mut input, + &["gpt-future-agent"], + json!([ + { + "type": "json_patch_body", + "patch": [ + {"op": "replace", "path": "/input", "value": "routed compact input"}, + {"op": "add", "path": "/instructions", "value": "Routed instructions"}, + {"op": "replace", "path": "/tools", "value": [{ + "type": "function", + "name": "lookup", + "parameters": {}, + "cache_control": {"type": "ephemeral"} + }]}, + {"op": "add", "path": "/parallel_tool_calls", "value": true}, + {"op": "add", "path": "/reasoning", "value": { + "effort": "high", + "context": "current_turn" + }} + ] + }, + { + "type": "patch_headers", + "patch": [{ + "op": "set", + "name": "x-openai-internal-codex-responses-lite", + "value": "false" + }] + } + ]), + ); + let mut decision = sample_decision(); + decision.provider_type = Some("codex".to_string()); + decision.provider_api_format = Some("openai:responses:compact".to_string()); + decision.client_api_format = Some("openai:responses:compact".to_string()); + decision.mapped_model = Some("gpt-future-agent".to_string()); + decision.provider_request_body = Some(json!({ + "model": "gpt-future-agent", + "input": [], + "tools": [] + })); + let transport = sample_codex_transport_with_card(); + + apply_provider_request_routing_policy_to_decision(&input, &mut decision, Some(&transport)) + .expect("terminal Lite contract should accept the projected request"); + + let body = decision.provider_request_body.as_ref().expect("body"); + assert_eq!(body["input"][0]["type"], "additional_tools"); + assert_eq!(body["input"][0]["tools"][0]["name"], "lookup"); + assert!(body["input"][0]["tools"][0].get("cache_control").is_none()); + assert_eq!(body["input"][1]["role"], "developer"); + assert_eq!( + body["input"][1]["content"][0]["text"], + "Routed instructions" + ); + assert_eq!(body["input"][2]["role"], "user"); + assert_eq!( + body["input"][2]["content"][0]["text"], + "routed compact input" + ); + assert!(body.get("tools").is_none()); + assert!(body.get("instructions").is_none()); + assert_eq!(body["parallel_tool_calls"], false); + assert_eq!(body["reasoning"]["effort"], "high"); + assert_eq!(body["reasoning"]["context"], "all_turns"); + assert_eq!( + decision + .provider_request_headers + .get(CODEX_RESPONSES_LITE_HEADER) + .map(String::as_str), + Some("true") + ); + } + #[test] fn provider_request_routing_policy_rejects_body_patch_without_json_body() { let input = sample_decision_input(); @@ -1042,7 +1457,7 @@ mod tests { decision.provider_request_body = None; decision.provider_request_body_base64 = Some("AA==".to_string()); - let error = apply_provider_request_routing_policy_to_decision(&input, &mut decision) + let error = apply_provider_request_routing_policy_to_decision(&input, &mut decision, None) .expect_err("provider body patch should reject binary upstream bodies"); match error { @@ -1067,6 +1482,7 @@ mod tests { let mut input = sample_decision_input(); set_provider_request_rules( &mut input, + &["gpt-5"], json!([{ "type": "patch_headers", "patch": [{ @@ -1080,7 +1496,7 @@ mod tests { decision.provider_request_body = None; decision.provider_request_body_base64 = Some("AA==".to_string()); - apply_provider_request_routing_policy_to_decision(&input, &mut decision) + apply_provider_request_routing_policy_to_decision(&input, &mut decision, None) .expect("header-only provider routing mutation should apply without JSON body"); assert_eq!(decision.provider_request_body, None); @@ -1112,7 +1528,7 @@ mod tests { "priority_slot": 3 })); - apply_provider_request_routing_policy_to_decision(&input, &mut decision) + apply_provider_request_routing_policy_to_decision(&input, &mut decision, None) .expect("provider routing mutation should seed pool trace"); let routing_trace = &decision.report_context.as_ref().unwrap()["routing_trace"]; diff --git a/apps/aether-gateway/src/ai_serving/planner/mod.rs b/apps/aether-gateway/src/ai_serving/planner/mod.rs index 29d86f149..d3161464c 100644 --- a/apps/aether-gateway/src/ai_serving/planner/mod.rs +++ b/apps/aether-gateway/src/ai_serving/planner/mod.rs @@ -34,6 +34,7 @@ pub(crate) use self::candidate_resolution::{ candidate_auth_channel_skip_reason, read_candidate_transport_snapshot, EligibleLocalExecutionCandidate, LocalExecutionCandidateKind, SkippedLocalExecutionCandidate, }; +pub(crate) use self::common::resolve_upstream_is_stream_for_provider; pub(crate) use self::passthrough::{ build_local_same_format_stream_attempt_source, build_local_same_format_stream_plan_and_reports, build_local_same_format_sync_attempt_source, build_local_same_format_sync_plan_and_reports, @@ -48,7 +49,7 @@ pub(crate) use self::plan_builders::{ pub(crate) use self::pool_scores::{ build_provider_key_pool_score_upsert, provider_key_pool_score_id, provider_key_pool_score_scope, }; -pub(crate) use self::request_gzip::resolve_transport_request_gzip_policy; +pub(crate) use self::request_gzip::resolve_transport_request_encoding_policy; pub(crate) use self::route::is_matching_stream_request as planner_is_matching_stream_request; pub(crate) use self::runtime_miss::{ apply_local_runtime_candidate_terminal_reason, record_local_runtime_candidate_skip_reason, @@ -79,7 +80,8 @@ pub(crate) use self::standard::{ build_local_stream_plan_and_reports as build_standard_family_stream_plan_and_reports, build_local_sync_attempt_source as build_standard_family_sync_attempt_source, build_local_sync_plan_and_reports as build_standard_family_sync_plan_and_reports, - set_local_openai_chat_execution_exhausted_diagnostic, + codex_model_capabilities_for_transport, set_local_openai_chat_execution_exhausted_diagnostic, + validate_final_openai_provider_request, }; pub(crate) use self::state::{ GatewayAuthApiKeySnapshot, GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth, diff --git a/apps/aether-gateway/src/ai_serving/planner/passthrough/plan_builders.rs b/apps/aether-gateway/src/ai_serving/planner/passthrough/plan_builders.rs index 4f9c76c39..1a55a2250 100644 --- a/apps/aether-gateway/src/ai_serving/planner/passthrough/plan_builders.rs +++ b/apps/aether-gateway/src/ai_serving/planner/passthrough/plan_builders.rs @@ -71,6 +71,7 @@ pub(crate) fn build_passthrough_stream_plan_from_decision( .content_type .take() .or_else(|| provider_request_headers.get("content-type").cloned()); + let stream = payload.upstream_is_stream; let plan = build_ai_execution_plan_from_decision( &mut payload, AiExecutionPlanFromDecisionParts { @@ -84,7 +85,7 @@ pub(crate) fn build_passthrough_stream_plan_from_decision( body_bytes_b64: None, body_ref: None, }, - stream: true, + stream, }, ); 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 d0dfbf872..fb9d7a102 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 @@ -24,7 +24,7 @@ use crate::ai_serving::{ ai_local_execution_contract_for_formats, extract_pool_sticky_session_token, resolve_local_decision_execution_runtime_auth_context, GatewayControlDecision, PlannerAppState, }; -use crate::client_session_affinity::client_session_affinity_from_parts; +use crate::client_session_affinity::client_session_affinity_from_api_request; use crate::clock::current_unix_secs; use crate::{AppState, GatewayError}; @@ -60,7 +60,9 @@ pub(crate) async fn resolve_local_same_format_provider_decision_input( state, auth_context, Some(requested_model.as_str()), + decision.auth_endpoint_signature.as_deref(), None, + &decision.model_directive_policy, ) .await { @@ -79,7 +81,11 @@ 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_session_affinity = client_session_affinity_from_parts(parts, Some(body_json)); + input.client_session_affinity = client_session_affinity_from_api_request( + spec_metadata.api_format, + &parts.headers, + Some(body_json), + ); if let Err(err) = attach_routing_policy_to_local_requested_model_input( state, parts, @@ -115,15 +121,22 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts( input.required_capabilities.as_ref(), LocalCandidatePersistencePolicyKind::SameFormatProviderDecision, ); + let model_directive_resolution = input + .model_directive_policy + .resolve_reasoning(spec_metadata.api_format, Some(&input.requested_model)); + let routing_model = model_directive_resolution + .base_model() + .unwrap_or(&input.requested_model); let (candidates, preselection_skipped) = planner_state .list_selectable_candidates_with_skip_reasons( spec_metadata.api_format, - &input.requested_model, + routing_model, spec_metadata.require_streaming, input.required_capabilities.as_ref(), Some(&input.auth_snapshot), input.client_session_affinity.as_ref(), current_unix_secs(), + false, ) .await?; let outcome = materialize_local_execution_candidates_with_serving( @@ -212,15 +225,22 @@ pub(crate) async fn build_local_same_format_provider_candidate_attempt_source<'a input.required_capabilities.as_ref(), LocalCandidatePersistencePolicyKind::SameFormatProviderDecision, ); + let model_directive_resolution = input + .model_directive_policy + .resolve_reasoning(spec_metadata.api_format, Some(&input.requested_model)); + let routing_model = model_directive_resolution + .base_model() + .unwrap_or(&input.requested_model); let (candidates, preselection_skipped) = planner_state .list_selectable_candidates_with_skip_reasons( spec_metadata.api_format, - &input.requested_model, + routing_model, spec_metadata.require_streaming, input.required_capabilities.as_ref(), Some(&input.auth_snapshot), input.client_session_affinity.as_ref(), current_unix_secs(), + false, ) .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 dc8de9ee8..5fb16dc56 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 @@ -17,7 +17,7 @@ use crate::ai_serving::planner::report_context::{ use crate::ai_serving::planner::spec_metadata::local_same_format_provider_spec_metadata; use crate::ai_serving::planner::CandidateFailureDiagnostic; use crate::ai_serving::planner::{ - build_ai_execution_decision_response, resolve_transport_request_gzip_policy, + build_ai_execution_decision_response, resolve_transport_request_encoding_policy, AiExecutionDecisionResponseParts, }; use crate::ai_serving::transport::{ @@ -184,7 +184,7 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_ compatibility_edits: _, request_redacted: _, } = resolved; - let request_gzip = resolve_transport_request_gzip_policy(&transport); + let request_encoding = resolve_transport_request_encoding_policy(&transport); let mut decision = build_ai_execution_decision_response(AiExecutionDecisionResponseParts { decision_is_stream: spec_metadata.require_streaming, @@ -194,6 +194,7 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_ request_id: trace_id.to_string(), candidate_id: candidate_id.to_string(), provider_name: transport.provider.name.clone(), + provider_type: transport.provider.provider_type.clone(), provider_id: candidate.provider_id.clone(), endpoint_id: candidate.endpoint_id.clone(), key_id: candidate.key_id.clone(), @@ -211,8 +212,8 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_ provider_request_body: Some(provider_request_body), provider_request_body_base64: None, content_type: Some("application/json".to_string()), - content_encoding: None, - request_gzip, + content_encoding: request_encoding.content_encoding, + request_gzip: request_encoding.request_gzip, proxy, transport_profile, timeouts: resolve_transport_execution_timeouts(&transport), @@ -221,7 +222,11 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_ report_context: Some(report_context), auth_context: input.auth_context.clone(), }); - apply_provider_request_routing_policy_to_decision(input, &mut decision)?; + apply_provider_request_routing_policy_to_decision( + input, + &mut decision, + Some(transport.as_ref()), + )?; Ok(Some(decision)) } 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 afb1aa557..a888f49fe 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 @@ -40,7 +40,9 @@ use super::{ LocalSameFormatProviderCandidateAttempt, LocalSameFormatProviderDecisionInput, LocalSameFormatProviderSpec, }; -use crate::ai_serving::planner::standard::same_format_provider_request_body_failure_extra_data; +use crate::ai_serving::planner::standard::{ + codex_model_capabilities_for_transport, same_format_provider_request_body_failure_extra_data, +}; pub(crate) fn resolve_same_format_provider_transport_unsupported_reason_for_trace( transport: &GatewayProviderTransportSnapshot, @@ -51,6 +53,7 @@ pub(crate) fn resolve_same_format_provider_transport_unsupported_reason_for_trac "openai:chat" => "openai:chat", "openai:responses" => "openai:responses", "openai:responses:compact" => "openai:responses:compact", + "openai:search" => "openai:search", "openai:embedding" => "openai:embedding", "openai:rerank" => "openai:rerank", "claude:messages" => "claude:messages", @@ -137,13 +140,26 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts( else { return Ok(None); }; - let enable_model_directives = - crate::system_features::reasoning_model_directive_enabled_for_api_format_and_model( - state, - spec.api_format, - Some(&input.requested_model), - ) - .await; + let model_directive_resolution = input + .model_directive_policy + .resolve_reasoning(spec.api_format, Some(&input.requested_model)); + let model_directive_mapping = + match model_directive_resolution.mapping_patch_for_mapped_model(&prepared.mapped_model) { + Ok(mapping) => mapping, + Err(skip_reason) => { + 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 effective_headers = input.effective_headers(&parts.headers); let redaction = resolve_provider_chat_pii_redaction( state, @@ -169,7 +185,7 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts( prepared.force_body_stream_field, prepared.kiro_auth.as_ref(), prepared.is_claude_code, - enable_model_directives, + false, ) else { mark_skipped_local_same_format_provider_candidate_with_extra_data( @@ -196,18 +212,11 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts( }; let mut base_provider_request_body = base_provider_request.body; let mut compatibility_edits = base_provider_request.compatibility_edits; - if let Some(mapping) = - crate::system_features::reasoning_model_directive_mapping_for_api_format_and_model( - state, - spec.api_format, - Some(&input.requested_model), - ) - .await - { + if let Some(mapping) = model_directive_mapping.as_ref() { let before_mapping = base_provider_request_body.clone(); crate::ai_serving::apply_model_directive_mapping_patch( &mut base_provider_request_body, - &mapping, + mapping, ); if before_mapping != base_provider_request_body { compatibility_edits.push(SameFormatProviderCompatibilityEdit { @@ -230,6 +239,48 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts( } } + let source_model = body_json + .get("model") + .and_then(Value::as_str) + .unwrap_or(input.requested_model.as_str()); + let codex_model_capabilities = codex_model_capabilities_for_transport( + &transport, + prepared.provider_api_format.as_str(), + prepared.mapped_model.as_str(), + source_model, + ); + if crate::ai_serving::finalize_openai_provider_request_with_codex_model_capabilities( + &mut base_provider_request_body, + crate::ai_serving::OpenAiProviderRequestFinalization { + source_api_format: spec.api_format, + provider_api_format: prepared.provider_api_format.as_str(), + provider_type: transport.provider.provider_type.as_str(), + provider_model: prepared.mapped_model.as_str(), + source_model, + body_rules: transport.endpoint.body_rules.as_ref(), + upstream_is_stream: prepared.upstream_is_stream, + require_body_stream_field: request_requires_body_stream_field( + body_json, + prepared.force_body_stream_field, + ), + }, + codex_model_capabilities.as_ref(), + ) + .is_err() + { + mark_skipped_local_same_format_provider_candidate( + state, + input, + trace_id, + candidate, + attempt.candidate_index, + &attempt.candidate_id, + "provider_request_body_missing", + ) + .await; + return Ok(None); + } + let antigravity_auth = if prepared.is_antigravity { let mut antigravity_support = classify_local_antigravity_request_support( &transport, @@ -467,6 +518,27 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts( .await; return Ok(None); }; + crate::ai_serving::apply_codex_openai_special_headers( + &mut provider_request_headers, + &provider_request_body, + effective_headers, + transport.provider.provider_type.as_str(), + prepared.provider_api_format.as_str(), + Some(trace_id), + transport.key.decrypted_auth_config.as_deref(), + ); + let provider_model = provider_request_body + .get("model") + .and_then(Value::as_str) + .unwrap_or(prepared.mapped_model.as_str()); + crate::ai_serving::apply_codex_openai_responses_lite_header_with_capabilities( + &mut provider_request_headers, + transport.provider.provider_type.as_str(), + prepared.provider_api_format.as_str(), + provider_model, + source_model, + codex_model_capabilities.as_ref(), + ); request_identity_response_encoding_when_redacted( &mut provider_request_headers, redaction.redacted, diff --git a/apps/aether-gateway/src/ai_serving/planner/request_gzip.rs b/apps/aether-gateway/src/ai_serving/planner/request_gzip.rs index 9fa314322..9de8cc4d0 100644 --- a/apps/aether-gateway/src/ai_serving/planner/request_gzip.rs +++ b/apps/aether-gateway/src/ai_serving/planner/request_gzip.rs @@ -1,23 +1,50 @@ use aether_ai_serving::AiRequestGzipPolicy; use serde_json::Value; -use crate::ai_serving::is_openai_responses_family_format; +use crate::ai_serving::{normalize_api_format_alias, parse_codex_auth_identity}; use super::state::GatewayProviderTransportSnapshot; -const DEFAULT_CODEX_REQUEST_GZIP_MIN_BYTES: usize = 64 * 1024; - -pub(crate) fn resolve_transport_request_gzip_policy( - transport: &GatewayProviderTransportSnapshot, -) -> Option { - transport_request_gzip_policy_from_config(transport.endpoint.config.as_ref()) - .or_else(|| transport_request_gzip_policy_from_config(transport.provider.config.as_ref())) - .or_else(|| default_transport_request_gzip_policy(transport)) +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub(crate) struct TransportRequestEncodingPolicy { + pub content_encoding: Option, + pub request_gzip: Option, } -fn default_transport_request_gzip_policy( +pub(crate) fn resolve_transport_request_encoding_policy( transport: &GatewayProviderTransportSnapshot, -) -> Option { +) -> TransportRequestEncodingPolicy { + if transport + .provider + .provider_type + .trim() + .eq_ignore_ascii_case("codex") + && normalize_api_format_alias(transport.endpoint.api_format.as_str()) + == "openai:responses:compact" + { + return TransportRequestEncodingPolicy::default(); + } + + let request_gzip = transport_request_gzip_policy_from_config( + transport.endpoint.config.as_ref(), + ) + .or_else(|| transport_request_gzip_policy_from_config(transport.provider.config.as_ref())); + if request_gzip.is_some() { + return TransportRequestEncodingPolicy { + content_encoding: None, + request_gzip, + }; + } + + TransportRequestEncodingPolicy { + content_encoding: default_transport_request_content_encoding(transport), + request_gzip: None, + } +} + +fn default_transport_request_content_encoding( + transport: &GatewayProviderTransportSnapshot, +) -> Option { if !transport .provider .provider_type @@ -26,19 +53,24 @@ fn default_transport_request_gzip_policy( { return None; } - if !is_codex_request_gzip_endpoint_api_format(transport.endpoint.api_format.as_str()) { + if !is_codex_request_compression_api_format(transport.endpoint.api_format.as_str()) { + return None; + } + let auth_type = + crate::ai_serving::transport::auth::resolve_local_auth_type_for_transport_format(transport); + let uses_codex_backend = auth_type == "oauth" + || (auth_type == "bearer" + && parse_codex_auth_identity(transport.key.decrypted_auth_config.as_deref()) + .uses_codex_backend); + if !uses_codex_backend { return None; } - Some(AiRequestGzipPolicy { - enabled: Some(true), - min_bytes: Some(DEFAULT_CODEX_REQUEST_GZIP_MIN_BYTES), - }) + Some("zstd".to_string()) } -fn is_codex_request_gzip_endpoint_api_format(api_format: &str) -> bool { - is_openai_responses_family_format(api_format) - || api_format.trim().eq_ignore_ascii_case("openai:image") +fn is_codex_request_compression_api_format(api_format: &str) -> bool { + normalize_api_format_alias(api_format) == "openai:responses" } fn transport_request_gzip_policy_from_config( @@ -216,6 +248,16 @@ mod tests { } } + fn resolved_gzip_policy( + transport: &GatewayProviderTransportSnapshot, + ) -> Option { + resolve_transport_request_encoding_policy(transport).request_gzip + } + + fn resolved_content_encoding(transport: &GatewayProviderTransportSnapshot) -> Option { + resolve_transport_request_encoding_policy(transport).content_encoding + } + #[test] fn endpoint_request_gzip_policy_overrides_provider_policy() { let transport = sample_transport( @@ -226,7 +268,7 @@ mod tests { ); assert_eq!( - resolve_transport_request_gzip_policy(&transport), + resolved_gzip_policy(&transport), Some(AiRequestGzipPolicy { enabled: Some(true), min_bytes: Some(1024), @@ -244,7 +286,7 @@ mod tests { ); assert_eq!( - resolve_transport_request_gzip_policy(&transport), + resolved_gzip_policy(&transport), Some(AiRequestGzipPolicy { enabled: Some(false), min_bytes: None, @@ -265,7 +307,7 @@ mod tests { ); assert_eq!( - resolve_transport_request_gzip_policy(&transport), + resolved_gzip_policy(&transport), Some(AiRequestGzipPolicy { enabled: Some(true), min_bytes: Some(4096), @@ -283,7 +325,7 @@ mod tests { ); assert_eq!( - resolve_transport_request_gzip_policy(&transport), + resolved_gzip_policy(&transport), Some(AiRequestGzipPolicy { enabled: Some(true), min_bytes: Some(1), @@ -292,35 +334,73 @@ mod tests { } #[test] - fn codex_responses_endpoint_gets_default_request_gzip_policy() { - let transport = sample_transport("codex", "openai:responses", None, None); + fn codex_responses_endpoint_uses_zstd_without_a_size_threshold() { + let mut transport = sample_transport("codex", "openai:responses", None, None); + transport.key.auth_type = "oauth".to_string(); assert_eq!( - resolve_transport_request_gzip_policy(&transport), - Some(AiRequestGzipPolicy { - enabled: Some(true), - min_bytes: Some(DEFAULT_CODEX_REQUEST_GZIP_MIN_BYTES), - }) + resolved_content_encoding(&transport).as_deref(), + Some("zstd") ); + assert_eq!(resolved_gzip_policy(&transport), None); } #[test] - fn codex_image_endpoint_gets_default_request_gzip_policy() { - let transport = sample_transport("codex", "openai:image", None, None); + fn codex_responses_api_key_auth_does_not_enable_default_compression() { + let transport = sample_transport("codex", "openai:responses", None, None); + + assert_eq!(resolved_content_encoding(&transport), None); + assert_eq!(resolved_gzip_policy(&transport), None); + } + + #[test] + fn codex_responses_bearer_auth_uses_identity_metadata_for_backend_compression() { + let mut transport = sample_transport("codex", "openai:responses", None, None); + transport.key.auth_type = "bearer".to_string(); + transport.key.decrypted_auth_config = + Some(r#"{"provider_type":"codex","account_id":"account-1"}"#.to_string()); assert_eq!( - resolve_transport_request_gzip_policy(&transport), - Some(AiRequestGzipPolicy { - enabled: Some(true), - min_bytes: Some(DEFAULT_CODEX_REQUEST_GZIP_MIN_BYTES), - }) + resolved_content_encoding(&transport).as_deref(), + Some("zstd") ); + assert_eq!(resolved_gzip_policy(&transport), None); + } + + #[test] + fn codex_image_endpoint_does_not_get_responses_request_gzip_policy() { + let transport = sample_transport("codex", "openai:image", None, None); + + assert_eq!(resolved_content_encoding(&transport), None); + assert_eq!(resolved_gzip_policy(&transport), None); + } + + #[test] + fn codex_compact_endpoint_does_not_get_default_request_gzip_policy() { + let transport = sample_transport("codex", "openai:responses:compact", None, None); + + assert_eq!(resolved_content_encoding(&transport), None); + assert_eq!(resolved_gzip_policy(&transport), None); + } + + #[test] + fn codex_compact_endpoint_rejects_an_explicit_request_gzip_policy() { + let transport = sample_transport( + "codex", + "openai:responses:compact", + None, + Some(json!({"request_gzip": {"enabled": true, "min_bytes": 2048}})), + ); + + assert_eq!(resolved_gzip_policy(&transport), None); + assert_eq!(resolved_content_encoding(&transport), None); } #[test] fn non_codex_endpoint_does_not_get_default_request_gzip_policy() { let transport = sample_transport("openai", "openai:responses", None, None); - assert_eq!(resolve_transport_request_gzip_policy(&transport), None); + assert_eq!(resolved_content_encoding(&transport), None); + assert_eq!(resolved_gzip_policy(&transport), None); } } diff --git a/apps/aether-gateway/src/ai_serving/planner/route.rs b/apps/aether-gateway/src/ai_serving/planner/route.rs index a08ccd192..362a6b623 100644 --- a/apps/aether-gateway/src/ai_serving/planner/route.rs +++ b/apps/aether-gateway/src/ai_serving/planner/route.rs @@ -77,6 +77,7 @@ mod tests { auth_endpoint_signature: None, execution_runtime_candidate: true, local_auth_rejection: None, + model_directive_policy: Default::default(), } } diff --git a/apps/aether-gateway/src/ai_serving/planner/specialized/files/decision.rs b/apps/aether-gateway/src/ai_serving/planner/specialized/files/decision.rs index 67416febc..3500358f5 100644 --- a/apps/aether-gateway/src/ai_serving/planner/specialized/files/decision.rs +++ b/apps/aether-gateway/src/ai_serving/planner/specialized/files/decision.rs @@ -7,7 +7,7 @@ use crate::ai_serving::planner::report_context::{ }; use crate::ai_serving::planner::spec_metadata::local_gemini_files_spec_metadata; use crate::ai_serving::planner::{ - build_ai_execution_decision_response, resolve_transport_request_gzip_policy, + build_ai_execution_decision_response, resolve_transport_request_encoding_policy, AiExecutionDecisionResponseParts, }; use crate::ai_serving::transport::{ @@ -124,7 +124,7 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat upstream_url, file_name: _, } = resolved; - let request_gzip = resolve_transport_request_gzip_policy(&transport); + let request_encoding = resolve_transport_request_encoding_policy(&transport); let mut decision = build_ai_execution_decision_response(AiExecutionDecisionResponseParts { decision_is_stream: spec_metadata.require_streaming, @@ -134,6 +134,7 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat request_id: trace_id.to_string(), candidate_id: candidate_id.clone(), provider_name: transport.provider.name.clone(), + provider_type: transport.provider.provider_type.clone(), provider_id: candidate.provider_id.clone(), endpoint_id: candidate.endpoint_id.clone(), key_id: candidate.key_id.clone(), @@ -156,8 +157,8 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat .map(str::trim) .filter(|value| !value.is_empty()) .map(ToOwned::to_owned), - content_encoding: None, - request_gzip, + content_encoding: request_encoding.content_encoding, + request_gzip: request_encoding.request_gzip, proxy, transport_profile, timeouts: resolve_transport_execution_timeouts(&transport), @@ -166,6 +167,10 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat report_context: Some(report_context), auth_context: input.auth_context.clone(), }); - apply_provider_request_routing_policy_to_decision(input, &mut decision)?; + apply_provider_request_routing_policy_to_decision( + input, + &mut decision, + Some(transport.as_ref()), + )?; Ok(Some(decision)) } diff --git a/apps/aether-gateway/src/ai_serving/planner/specialized/files/support.rs b/apps/aether-gateway/src/ai_serving/planner/specialized/files/support.rs index 33b94a401..c2d736d61 100644 --- a/apps/aether-gateway/src/ai_serving/planner/specialized/files/support.rs +++ b/apps/aether-gateway/src/ai_serving/planner/specialized/files/support.rs @@ -53,7 +53,9 @@ pub(super) async fn resolve_local_gemini_files_decision_input( state, auth_context, None, + decision.auth_endpoint_signature.as_deref(), Some(&explicit_required_capabilities), + &decision.model_directive_policy, ) .await { diff --git a/apps/aether-gateway/src/ai_serving/planner/specialized/image/decision.rs b/apps/aether-gateway/src/ai_serving/planner/specialized/image/decision.rs index e65bcb816..5cbecc94e 100644 --- a/apps/aether-gateway/src/ai_serving/planner/specialized/image/decision.rs +++ b/apps/aether-gateway/src/ai_serving/planner/specialized/image/decision.rs @@ -5,7 +5,7 @@ use crate::ai_serving::planner::report_context::{ }; use crate::ai_serving::planner::spec_metadata::local_openai_image_spec_metadata; use crate::ai_serving::planner::{ - build_ai_execution_decision_response, resolve_transport_request_gzip_policy, + build_ai_execution_decision_response, resolve_transport_request_encoding_policy, AiExecutionDecisionResponseParts, }; use crate::ai_serving::transport::{ @@ -84,11 +84,7 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat serde_json::Value::Bool(true), ); } - let upstream_is_stream = resolved - .provider_request_body - .get("stream") - .and_then(serde_json::Value::as_bool) - .unwrap_or(spec_metadata.require_streaming); + let upstream_is_stream = resolved.upstream_is_stream; let effective_headers = input.effective_headers(&parts.headers); let report_context = append_execution_contract_fields_to_value( build_local_execution_report_context(LocalExecutionReportContextParts { @@ -135,7 +131,7 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat spec_metadata.api_format, provider_api_format.as_str(), ); - let request_gzip = resolve_transport_request_gzip_policy(&transport); + let request_encoding = resolve_transport_request_encoding_policy(&transport); let mut decision = build_ai_execution_decision_response(AiExecutionDecisionResponseParts { decision_is_stream: spec_metadata.require_streaming, @@ -145,6 +141,7 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat request_id: trace_id.to_string(), candidate_id: candidate_id.clone(), provider_name: transport.provider.name.clone(), + provider_type: transport.provider.provider_type.clone(), provider_id: candidate.provider_id.clone(), endpoint_id: candidate.endpoint_id.clone(), key_id: candidate.key_id.clone(), @@ -162,8 +159,8 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat provider_request_body: Some(resolved.provider_request_body), provider_request_body_base64: None, content_type: Some("application/json".to_string()), - content_encoding: None, - request_gzip, + content_encoding: request_encoding.content_encoding, + request_gzip: request_encoding.request_gzip, proxy, transport_profile, timeouts: resolve_transport_execution_timeouts(&transport), @@ -172,6 +169,10 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat report_context: Some(report_context), auth_context: input.auth_context.clone(), }); - apply_provider_request_routing_policy_to_decision(input, &mut decision)?; + apply_provider_request_routing_policy_to_decision( + input, + &mut decision, + Some(transport.as_ref()), + )?; Ok(Some(decision)) } diff --git a/apps/aether-gateway/src/ai_serving/planner/specialized/image/request.rs b/apps/aether-gateway/src/ai_serving/planner/specialized/image/request.rs index ea191261a..37cc057cb 100644 --- a/apps/aether-gateway/src/ai_serving/planner/specialized/image/request.rs +++ b/apps/aether-gateway/src/ai_serving/planner/specialized/image/request.rs @@ -16,8 +16,8 @@ use crate::ai_serving::transport::{ ProviderOpenAiImageHeadersInput, StandardProviderRequestHeadersInput, GROK_CHAT_PATH, }; use crate::ai_serving::{ - apply_codex_openai_responses_special_body_edits, apply_codex_openai_responses_special_headers, - build_chatgpt_web_image_request_body, + apply_codex_openai_special_headers, build_chatgpt_web_image_request_body, + build_codex_openai_image_api_provider_request_body, build_gemini_image_request_body_from_openai_image_request, build_openai_image_api_provider_request_body, build_openai_image_provider_request_body, default_model_for_openai_image_operation, normalize_openai_image_request, @@ -48,6 +48,7 @@ pub(super) struct LocalOpenAiImageCandidatePayloadParts { pub(super) upstream_url: String, pub(super) input_summary: Value, pub(super) transport_profile: Option, + pub(super) upstream_is_stream: bool, } pub(super) async fn resolve_local_openai_image_candidate_payload_parts( @@ -130,7 +131,10 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts( parts, body_json, body_base64, - openai_image_normalize_options_for_provider(&transport.provider.provider_type), + openai_image_normalize_options_for_provider( + &transport.provider.provider_type, + Some(prepared_candidate.mapped_model.as_str()), + ), ); let Some(normalized_request) = normalized_request else { mark_skipped_local_openai_image_candidate_with_failure_diagnostic( @@ -174,29 +178,56 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts( } else { build_openai_image_upstream_url(transport, Some(parts.uri.path()), parts.uri.query()) }; - let mut provider_request_body = if is_chatgpt_web { - match build_chatgpt_web_image_request_body(parts, body_json, body_base64) { - Ok(body) => body, - Err(err) => err.to_error_json(), - } - } else if is_codex || is_grok { - build_openai_image_provider_request_body(&normalized_request) + let upstream_is_stream = + crate::ai_serving::planner::common::resolve_upstream_is_stream_for_provider( + transport.endpoint.config.as_ref(), + transport.provider.provider_type.as_str(), + spec_metadata.api_format, + spec_metadata.require_streaming && candidate.supports_streaming, + false, + ); + let provider_request_body = if is_chatgpt_web { + Some( + match build_chatgpt_web_image_request_body(parts, body_json, body_base64) { + Ok(body) => body, + Err(err) => err.to_error_json(), + }, + ) + } else if is_codex { + build_codex_openai_image_api_provider_request_body( + &normalized_request, + Some(prepared_candidate.mapped_model.as_str()), + upstream_is_stream, + ) + } else if is_grok { + Some(build_openai_image_provider_request_body( + &normalized_request, + )) } else { build_openai_image_api_provider_request_body( &normalized_request, Some(prepared_candidate.mapped_model.as_str()), + upstream_is_stream, ) }; - if !is_chatgpt_web { - apply_codex_openai_responses_special_body_edits( - &mut provider_request_body, - transport.provider.provider_type.as_str(), - spec_metadata.api_format, - transport.endpoint.body_rules.as_ref(), - Some(candidate.key_id.as_str()), - ); - } - + let Some(provider_request_body) = provider_request_body else { + mark_skipped_local_openai_image_candidate_with_failure_diagnostic( + state, + input, + trace_id, + candidate, + attempt.candidate_index, + &attempt.candidate_id, + "provider_request_body_missing", + CandidateFailureDiagnostic::provider_request_body_missing( + spec_metadata.api_format, + spec_metadata.api_format, + "codex_openai_images_request_contract", + ), + ) + .await; + return None; + }; let Some(mut provider_request_headers) = (if is_grok { build_grok_browser_headers(GrokHeaderInput { transport, @@ -214,10 +245,12 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts( headers: effective_headers, auth_header: &auth_header, auth_value: &auth_value, - accept: if is_codex || is_chatgpt_web { - "text/event-stream" + accept: if is_codex { + None + } else if upstream_is_stream { + Some("text/event-stream") } else { - "application/json" + Some("application/json") }, header_rules: transport.endpoint.header_rules.as_ref(), provider_request_body: &provider_request_body, @@ -245,7 +278,7 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts( provider_request_headers.insert("x-aether-chatgpt-web-image".to_string(), "1".to_string()); } else if is_grok { } else { - apply_codex_openai_responses_special_headers( + apply_codex_openai_special_headers( &mut provider_request_headers, &provider_request_body, effective_headers, @@ -287,6 +320,7 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts( upstream_url, input_summary, transport_profile, + upstream_is_stream, }) } @@ -403,7 +437,14 @@ async fn resolve_local_openai_image_to_gemini_candidate_payload_parts( return None; } }; - let upstream_is_stream = spec_metadata.require_streaming; + let upstream_is_stream = + crate::ai_serving::planner::common::resolve_upstream_is_stream_for_provider( + transport.endpoint.config.as_ref(), + transport.provider.provider_type.as_str(), + provider_api_format, + spec_metadata.require_streaming && candidate.supports_streaming, + false, + ); let Some(upstream_url) = crate::ai_serving::planner::standard::build_standard_upstream_url( parts, transport, @@ -474,6 +515,7 @@ async fn resolve_local_openai_image_to_gemini_candidate_payload_parts( upstream_url, input_summary: converted.summary_json, transport_profile: None, + upstream_is_stream, }) } diff --git a/apps/aether-gateway/src/ai_serving/planner/specialized/image/support.rs b/apps/aether-gateway/src/ai_serving/planner/specialized/image/support.rs index 58b7b6bc2..3a59fcf56 100644 --- a/apps/aether-gateway/src/ai_serving/planner/specialized/image/support.rs +++ b/apps/aether-gateway/src/ai_serving/planner/specialized/image/support.rs @@ -58,7 +58,9 @@ pub(super) async fn resolve_local_openai_image_decision_input( state, auth_context, Some(requested_model.as_str()), + decision.auth_endpoint_signature.as_deref(), None, + &decision.model_directive_policy, ) .await { @@ -124,6 +126,7 @@ pub(super) async fn list_local_openai_image_candidate_attempts( matches_client_format.then_some(&input.auth_snapshot), input.client_session_affinity.as_ref(), current_unix_secs(), + false, ) .await { @@ -144,8 +147,8 @@ pub(super) async fn list_local_openai_image_candidate_attempts( auth_snapshot_allows_cross_format_candidate( &input.auth_snapshot, &input.requested_model, + None, candidate, - false, ) }); } @@ -197,6 +200,7 @@ pub(super) async fn build_local_openai_image_candidate_attempt_source<'a>( matches_client_format.then_some(&input.auth_snapshot), input.client_session_affinity.as_ref(), current_unix_secs(), + false, ) .await { @@ -206,16 +210,16 @@ pub(super) async fn build_local_openai_image_candidate_attempt_source<'a>( auth_snapshot_allows_cross_format_candidate( &input.auth_snapshot, &input.requested_model, + None, candidate, - false, ) }); format_skipped.retain(|candidate| { auth_snapshot_allows_cross_format_candidate( &input.auth_snapshot, &input.requested_model, + None, &candidate.candidate, - false, ) }); } diff --git a/apps/aether-gateway/src/ai_serving/planner/specialized/video/decision.rs b/apps/aether-gateway/src/ai_serving/planner/specialized/video/decision.rs index 8d38e6ad8..6307329b8 100644 --- a/apps/aether-gateway/src/ai_serving/planner/specialized/video/decision.rs +++ b/apps/aether-gateway/src/ai_serving/planner/specialized/video/decision.rs @@ -5,7 +5,7 @@ use crate::ai_serving::planner::report_context::{ }; use crate::ai_serving::planner::spec_metadata::local_video_create_spec_metadata; use crate::ai_serving::planner::{ - build_ai_execution_decision_response, resolve_transport_request_gzip_policy, + build_ai_execution_decision_response, resolve_transport_request_encoding_policy, AiExecutionDecisionResponseParts, }; use crate::ai_serving::transport::{ @@ -104,7 +104,7 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat provider_request_body, upstream_url, } = resolved; - let request_gzip = resolve_transport_request_gzip_policy(&transport); + let request_encoding = resolve_transport_request_encoding_policy(&transport); let mut decision = build_ai_execution_decision_response(AiExecutionDecisionResponseParts { decision_is_stream: false, @@ -114,6 +114,7 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat request_id: trace_id.to_string(), candidate_id: candidate_id.clone(), provider_name: transport.provider.name.clone(), + provider_type: transport.provider.provider_type.clone(), provider_id: candidate.provider_id.clone(), endpoint_id: candidate.endpoint_id.clone(), key_id: candidate.key_id.clone(), @@ -137,8 +138,8 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat .map(str::trim) .filter(|value| !value.is_empty()) .map(ToOwned::to_owned), - content_encoding: None, - request_gzip, + content_encoding: request_encoding.content_encoding, + request_gzip: request_encoding.request_gzip, proxy, transport_profile, timeouts: resolve_transport_execution_timeouts(&transport), @@ -147,6 +148,10 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat report_context: Some(report_context), auth_context: input.auth_context.clone(), }); - apply_provider_request_routing_policy_to_decision(input, &mut decision)?; + apply_provider_request_routing_policy_to_decision( + input, + &mut decision, + Some(transport.as_ref()), + )?; Ok(Some(decision)) } diff --git a/apps/aether-gateway/src/ai_serving/planner/specialized/video/support.rs b/apps/aether-gateway/src/ai_serving/planner/specialized/video/support.rs index 2b382b94f..ffaf3eedf 100644 --- a/apps/aether-gateway/src/ai_serving/planner/specialized/video/support.rs +++ b/apps/aether-gateway/src/ai_serving/planner/specialized/video/support.rs @@ -62,7 +62,9 @@ pub(super) async fn resolve_local_video_create_decision_input( state, auth_context, Some(requested_model.as_str()), + decision.auth_endpoint_signature.as_deref(), None, + &decision.model_directive_policy, ) .await { @@ -130,6 +132,7 @@ pub(super) async fn list_local_video_create_candidate_attempts( Some(&input.auth_snapshot), input.client_session_affinity.as_ref(), current_unix_secs(), + false, ) .await { @@ -186,6 +189,7 @@ pub(super) async fn build_local_video_create_candidate_attempt_source<'a>( Some(&input.auth_snapshot), input.client_session_affinity.as_ref(), current_unix_secs(), + false, ) .await { diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/codex.rs b/apps/aether-gateway/src/ai_serving/planner/standard/codex.rs index aeee12699..9bdc65adb 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/codex.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/codex.rs @@ -3,5 +3,42 @@ mod tests; pub(crate) use crate::ai_serving::{ - apply_codex_openai_responses_special_body_edits, apply_codex_openai_responses_special_headers, + apply_codex_openai_responses_special_body_edits, apply_codex_openai_special_headers, }; + +pub(crate) fn codex_model_capabilities_for_transport( + transport: &crate::ai_serving::GatewayProviderTransportSnapshot, + provider_api_format: &str, + provider_model: &str, + source_model: &str, +) -> Option { + codex_model_capabilities( + &transport.provider.provider_type, + provider_api_format, + provider_model, + source_model, + transport.key.upstream_metadata.as_ref(), + ) +} + +fn codex_model_capabilities( + provider_type: &str, + provider_api_format: &str, + provider_model: &str, + source_model: &str, + upstream_metadata: Option<&serde_json::Value>, +) -> Option { + let uses_codex_model_catalog = + crate::ai_serving::is_openai_responses_family_format(provider_api_format) + || crate::ai_serving::api_format_alias_matches(provider_api_format, "openai:search"); + if !provider_type.trim().eq_ignore_ascii_case("codex") || !uses_codex_model_catalog { + return None; + } + Some( + crate::ai_serving::resolve_codex_responses_model_capabilities( + provider_model, + source_model, + upstream_metadata, + ), + ) +} diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/codex/tests.rs b/apps/aether-gateway/src/ai_serving/planner/standard/codex/tests.rs index b002ea28e..10bef0e6b 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/codex/tests.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/codex/tests.rs @@ -1,16 +1,56 @@ use std::collections::BTreeMap; use super::{ - apply_codex_openai_responses_special_body_edits, apply_codex_openai_responses_special_headers, + apply_codex_openai_responses_special_body_edits, apply_codex_openai_special_headers, + codex_model_capabilities, }; use crate::ai_serving::planner::standard::build_local_openai_responses_request_body; use http::{HeaderMap, HeaderValue}; use serde_json::json; +#[test] +fn search_uses_live_codex_model_catalog_capabilities() { + let metadata = crate::ai_serving::build_codex_model_catalog_metadata(&[json!({ + "slug": "gpt-search-custom", + "default_reasoning_level": "low", + "supported_reasoning_levels": [ + {"effort": "low"}, + {"effort": "max"} + ], + "supports_parallel_tool_calls": true + })]); + + let capabilities = codex_model_capabilities( + "codex", + "openai:search", + "gpt-search-custom", + "gpt-search-custom", + Some(&metadata), + ) + .expect("Search should resolve capabilities from the Codex model catalog"); + + assert_eq!( + capabilities.default_reasoning_effort.as_deref(), + Some("low") + ); + assert_eq!( + capabilities.supported_reasoning_efforts, + vec!["low".to_string(), "max".to_string()] + ); + assert!(codex_model_capabilities( + "codex", + "openai:chat", + "gpt-search-custom", + "gpt-search-custom", + Some(&metadata), + ) + .is_none()); +} + #[test] fn applies_codex_defaults_when_body_rules_do_not_handle_fields() { let mut body = json!({ - "model": "gpt-5", + "model": "gpt-5.4", "max_output_tokens": 128, "temperature": 0.3, "top_p": 0.9, @@ -31,10 +71,11 @@ fn applies_codex_defaults_when_body_rules_do_not_handle_fields() { assert!(body.get("top_p").is_none()); assert!(body.get("metadata").is_none()); assert_eq!(body["store"], false); - assert_eq!(body["instructions"], ""); + assert!(body.get("instructions").is_none()); assert_eq!(body["include"], json!(["reasoning.encrypted_content"])); assert_eq!(body["parallel_tool_calls"], true); - assert!(body.get("reasoning").is_none()); + assert_eq!(body["reasoning"]["effort"], "medium"); + assert!(body["reasoning"].get("summary").is_none()); } #[test] @@ -80,7 +121,7 @@ fn strips_store_for_compact_even_when_body_rules_handle_it() { {"action":"set","path":"top_p","value":0.5} ]); let mut body = json!({ - "model": "gpt-5", + "model": "gpt-5.4", "max_output_tokens": 128, "metadata": {"client": "desktop", "mode": "custom"}, "store": true, @@ -99,12 +140,29 @@ fn strips_store_for_compact_even_when_body_rules_handle_it() { assert!(body.get("max_output_tokens").is_none()); assert!(body.get("store").is_none()); assert_eq!(body["instructions"], "Keep custom"); - assert_eq!(body["metadata"]["mode"], "custom"); - assert_eq!(body["top_p"], 0.5); + assert!(body.get("metadata").is_none()); + assert!(body.get("top_p").is_none()); + assert_eq!(body["parallel_tool_calls"], true); + assert!(body.as_object().is_some_and(|object| { + object.keys().all(|field| { + matches!( + field.as_str(), + "model" + | "input" + | "instructions" + | "tools" + | "parallel_tool_calls" + | "reasoning" + | "service_tier" + | "prompt_cache_key" + | "text" + ) + }) + })); } #[test] -fn injects_stable_prompt_cache_key_for_codex_requests() { +fn does_not_synthesize_prompt_cache_key_from_api_key_identity() { let mut body = json!({ "model": "gpt-5", "input": "hello", @@ -118,10 +176,7 @@ fn injects_stable_prompt_cache_key_for_codex_requests() { Some("key-123"), ); - assert_eq!( - body["prompt_cache_key"], - "53363264-dbb0-5f9d-b9c7-3e92c45c5bdf" - ); + assert!(body.get("prompt_cache_key").is_none()); } #[test] @@ -144,57 +199,96 @@ fn keeps_existing_prompt_cache_key_for_codex_requests() { } #[test] -fn injects_chatgpt_account_id_and_session_headers_for_codex_requests() { +fn injects_identity_headers_without_deriving_session_headers_from_body() { let mut headers = BTreeMap::new(); let body = json!({ "model": "gpt-5", "prompt_cache_key": "172c39e6-c0a0-5a70-8b63-e0f8e0d185a3", }); - apply_codex_openai_responses_special_headers( + apply_codex_openai_special_headers( &mut headers, &body, &HeaderMap::new(), "codex", "openai:responses", Some("trace-codex-123"), - Some(r#"{"account_id":"acc-123"}"#), + Some(r#"{"account_id":"acc-123","is_fedramp":true}"#), + ); + + assert_eq!( + headers.get("chatgpt-account-id"), + Some(&"acc-123".to_string()) + ); + assert_eq!(headers.get("x-client-request-id"), None); + assert_eq!( + headers.get("user-agent"), + Some(&"codex_cli_rs/0.144.1".to_string()) + ); + assert_eq!(headers.get("originator"), Some(&"codex_cli_rs".to_string())); + assert!(!headers.contains_key("version")); + assert_eq!(headers.get("x-openai-fedramp"), Some(&"true".to_string())); + assert_eq!(headers.get("session-id"), None); + assert_eq!(headers.get("thread-id"), None); +} + +#[test] +fn injects_only_codex_client_headers_for_images_requests() { + let mut headers = BTreeMap::new(); + apply_codex_openai_special_headers( + &mut headers, + &json!({ + "model": "gpt-image-2", + "prompt": "draw a city" + }), + &HeaderMap::new(), + "codex", + "openai:image", + Some("trace-codex-image-123"), + Some(r#"{"account_id":"acc-123","is_fedramp":true}"#), ); assert_eq!( headers.get("chatgpt-account-id"), Some(&"acc-123".to_string()) ); - assert_eq!( - headers.get("x-client-request-id"), - Some(&"trace-codex-123".to_string()) - ); assert_eq!( headers.get("user-agent"), - Some( - &"codex-tui/0.122.0 (Mac OS 15.2.0; arm64) vscode/2.6.11 (codex-tui; 0.122.0)" - .to_string() - ) - ); - assert_eq!(headers.get("originator"), Some(&"codex-tui".to_string())); - assert_eq!( - headers.get("session_id"), - Some(&"ab5ecce4f0d110fe".to_string()) - ); - assert_eq!( - headers.get("conversation_id"), - Some(&"ab5ecce4f0d110fe".to_string()) + Some(&"codex_cli_rs/0.144.1".to_string()) ); + assert_eq!(headers.get("originator"), Some(&"codex_cli_rs".to_string())); + assert!(!headers.contains_key("version")); + assert_eq!(headers.get("x-openai-fedramp"), Some(&"true".to_string())); + for name in ["x-client-request-id", "session-id", "thread-id"] { + assert!( + !headers.contains_key(name), + "unexpected Images header: {name}" + ); + } } #[test] -fn respects_existing_codex_request_and_session_headers() { +fn preserves_client_context_headers_and_enforces_codex_provider_identity() { let mut headers = BTreeMap::new(); headers.insert( "x-client-request-id".to_string(), "kept-by-rule-request".to_string(), ); - headers.insert("session_id".to_string(), "kept-by-rule".to_string()); + headers.insert("session-id".to_string(), "kept-by-rule-session".to_string()); + headers.insert("thread-id".to_string(), "kept-by-rule-thread".to_string()); + headers.insert( + "chatgpt-account-id".to_string(), + "configured-spoof".to_string(), + ); + headers.insert( + "x-openai-fedramp".to_string(), + "configured-false".to_string(), + ); + headers.insert( + "User-Agent".to_string(), + "AsyncOpenAI/Python 2.44.0".to_string(), + ); + headers.insert("ORIGINATOR".to_string(), "sdk-client".to_string()); let body = json!({ "model": "gpt-5", "prompt_cache_key": "172c39e6-c0a0-5a70-8b63-e0f8e0d185a3", @@ -205,12 +299,12 @@ fn respects_existing_codex_request_and_session_headers() { HeaderValue::from_static("user-specified-request"), ); original_headers.insert( - "session_id", + "session-id", HeaderValue::from_static("user-specified-session"), ); original_headers.insert( - "conversation_id", - HeaderValue::from_static("user-specified-conversation"), + "thread-id", + HeaderValue::from_static("user-specified-thread"), ); original_headers.insert( "user-agent", @@ -220,64 +314,92 @@ fn respects_existing_codex_request_and_session_headers() { "originator", HeaderValue::from_static("user-specified-originator"), ); + original_headers.insert("version", HeaderValue::from_static("user-version")); + original_headers.insert("x-openai-fedramp", HeaderValue::from_static("user-fedramp")); + original_headers.insert( + "chatgpt-account-id", + HeaderValue::from_static("user-account"), + ); - apply_codex_openai_responses_special_headers( + apply_codex_openai_special_headers( &mut headers, &body, &original_headers, "codex", "openai:responses", Some("trace-codex-123"), - Some(r#"{"account_id":"acc-123"}"#), + Some(r#"{"account_id":"acc-123","is_fedramp":true}"#), ); assert_eq!( headers.get("x-client-request-id"), Some(&"kept-by-rule-request".to_string()) ); - assert!(!headers.contains_key("user-agent")); - assert!(!headers.contains_key("originator")); - assert_eq!(headers.get("session_id"), Some(&"kept-by-rule".to_string())); - assert!(!headers.contains_key("conversation_id")); + assert_eq!( + headers.get("user-agent"), + Some(&"codex_cli_rs/0.144.1".to_string()) + ); + assert_eq!(headers.get("originator"), Some(&"codex_cli_rs".to_string())); + assert_eq!( + headers + .keys() + .filter(|name| name.eq_ignore_ascii_case("user-agent")) + .count(), + 1 + ); + assert_eq!( + headers + .keys() + .filter(|name| name.eq_ignore_ascii_case("originator")) + .count(), + 1 + ); + assert!(!headers.contains_key("version")); + assert_eq!( + headers.get("chatgpt-account-id"), + Some(&"acc-123".to_string()) + ); + assert_eq!(headers.get("x-openai-fedramp"), Some(&"true".to_string())); + assert_eq!( + headers.get("session-id"), + Some(&"kept-by-rule-session".to_string()) + ); + assert_eq!( + headers.get("thread-id"), + Some(&"kept-by-rule-thread".to_string()) + ); } #[test] -fn skips_conversation_id_for_compact_codex_requests() { +fn compact_does_not_derive_session_headers_from_body() { let mut headers = BTreeMap::new(); let body = json!({ "model": "gpt-5", "prompt_cache_key": "172c39e6-c0a0-5a70-8b63-e0f8e0d185a3", }); - apply_codex_openai_responses_special_headers( + apply_codex_openai_special_headers( &mut headers, &body, &HeaderMap::new(), "codex", "openai:responses:compact", Some("trace-codex-compact-123"), - Some(r#"{"account_id":"acc-123"}"#), + Some(r#"{"account_id":"acc-123","is_fedramp":true}"#), ); assert_eq!( headers.get("chatgpt-account-id"), Some(&"acc-123".to_string()) ); - assert_eq!( - headers.get("x-client-request-id"), - Some(&"trace-codex-compact-123".to_string()) - ); + assert_eq!(headers.get("x-client-request-id"), None); assert_eq!( headers.get("user-agent"), - Some( - &"codex-tui/0.122.0 (Mac OS 15.2.0; arm64) vscode/2.6.11 (codex-tui; 0.122.0)" - .to_string() - ) + Some(&"codex_cli_rs/0.144.1".to_string()) ); - assert_eq!(headers.get("originator"), Some(&"codex-tui".to_string())); - assert_eq!( - headers.get("session_id"), - Some(&"ab5ecce4f0d110fe".to_string()) - ); - assert!(!headers.contains_key("conversation_id")); + assert_eq!(headers.get("originator"), Some(&"codex_cli_rs".to_string())); + assert!(!headers.contains_key("version")); + assert_eq!(headers.get("x-openai-fedramp"), Some(&"true".to_string())); + assert_eq!(headers.get("session-id"), None); + assert_eq!(headers.get("thread-id"), None); } diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/family/candidates.rs b/apps/aether-gateway/src/ai_serving/planner/standard/family/candidates.rs index 9ff137958..8cea21bd6 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/family/candidates.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/family/candidates.rs @@ -60,7 +60,9 @@ pub(super) async fn resolve_local_standard_decision_input( state, auth_context, Some(requested_model.as_str()), + decision.auth_endpoint_signature.as_deref(), None, + &decision.model_directive_policy, ) .await { @@ -119,6 +121,7 @@ pub(super) async fn materialize_local_standard_candidate_attempts( ); let preselection = preselect_local_execution_candidates_with_serving( planner_state, + &input.model_directive_policy, spec_metadata.api_format, &input.requested_model, false, @@ -243,6 +246,7 @@ pub(super) async fn build_local_standard_candidate_attempt_source<'a>( let (source, candidate_count) = build_lazy_requested_model_execution_candidate_attempt_source_with_serving( planner_state, + &input.model_directive_policy, trace_id, spec_metadata.api_format, &input.requested_model, @@ -338,6 +342,7 @@ async fn maybe_append_gemini_image_openai_image_preselection( let image_preselection = preselect_local_execution_candidates_for_api_formats_with_serving( planner_state, + &input.model_directive_policy, spec_metadata.api_format, &input.requested_model, spec_metadata.require_streaming, 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 033f53130..151cde139 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 @@ -15,7 +15,7 @@ use crate::ai_serving::planner::report_context::{ use crate::ai_serving::planner::spec_metadata::local_standard_spec_metadata; use crate::ai_serving::planner::CandidateFailureDiagnostic; use crate::ai_serving::planner::{ - build_ai_execution_decision_response, resolve_transport_request_gzip_policy, + build_ai_execution_decision_response, resolve_transport_request_encoding_policy, AiExecutionDecisionResponseParts, }; use crate::ai_serving::transport::{ @@ -176,7 +176,7 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate( transport_profile: _, request_redacted: _, } = resolved; - let request_gzip = resolve_transport_request_gzip_policy(&transport); + let request_encoding = resolve_transport_request_encoding_policy(&transport); let mut decision = build_ai_execution_decision_response(AiExecutionDecisionResponseParts { decision_is_stream: spec_metadata.require_streaming, @@ -186,6 +186,7 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate( request_id: trace_id.to_string(), candidate_id: candidate_id.to_string(), provider_name: candidate.provider_name.clone(), + provider_type: transport.provider.provider_type.clone(), provider_id: candidate.provider_id.clone(), endpoint_id: candidate.endpoint_id.clone(), key_id: candidate.key_id.clone(), @@ -203,8 +204,8 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate( provider_request_body: Some(provider_request_body), provider_request_body_base64: None, content_type: Some("application/json".to_string()), - content_encoding: None, - request_gzip, + content_encoding: request_encoding.content_encoding, + request_gzip: request_encoding.request_gzip, proxy, transport_profile, timeouts, @@ -213,7 +214,11 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate( report_context: Some(report_context), auth_context: input.auth_context.clone(), }); - apply_provider_request_routing_policy_to_decision(input, &mut decision)?; + apply_provider_request_routing_policy_to_decision( + input, + &mut decision, + Some(transport.as_ref()), + )?; Ok(Some(decision)) } @@ -372,6 +377,7 @@ mod tests { routing_policy: None, routing_trace_seed: None, routing_context: None, + model_directive_policy: Default::default(), } } @@ -475,6 +481,7 @@ mod tests { } else { "gpt-4o-upstream".to_string() }, + supports_streaming: true, mapping_matched_model: None, } } diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/family/request.rs b/apps/aether-gateway/src/ai_serving/planner/standard/family/request.rs index c6a631360..e2c29c99e 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/family/request.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/family/request.rs @@ -21,9 +21,9 @@ use crate::ai_serving::planner::redaction::{ }; use crate::ai_serving::planner::spec_metadata::local_standard_spec_metadata; use crate::ai_serving::planner::standard::{ - apply_codex_openai_responses_special_headers, apply_deepseek_tool_call_thinking_compat, - is_deepseek_provider, request_body_build_failure_extra_data, - request_conversion_failure_extra_data, + apply_codex_openai_special_headers, apply_deepseek_tool_call_thinking_compat, + codex_model_capabilities_for_transport, is_deepseek_provider, + request_body_build_failure_extra_data, request_conversion_failure_extra_data, }; use crate::ai_serving::transport::kiro::{ build_kiro_provider_headers, build_kiro_provider_request_body, @@ -44,7 +44,9 @@ use crate::ai_serving::transport::{ }; use crate::ai_serving::{ build_openai_image_request_body_from_gemini_image_request, gemini_request_is_image_generation, + project_codex_openai_image_api_request_body, project_openai_image_api_request_body, CandidateFailureDiagnostic, GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth, + OpenAiImageOperation, }; use crate::{AppState, GatewayError}; @@ -313,7 +315,13 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts( { return Ok( resolve_local_gemini_image_to_openai_image_candidate_payload_parts( - state, parts, trace_id, body_json, input, attempt, + state, + parts, + trace_id, + body_json, + input, + attempt, + spec_metadata.require_streaming, ) .await, ); @@ -555,13 +563,27 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts( ); let force_body_stream_field = endpoint_config_forces_body_stream_field(transport.endpoint.config.as_ref()); - let enable_model_directives = - crate::system_features::reasoning_model_directive_enabled_for_api_format_and_model( - state, - provider_api_format, - Some(&input.requested_model), - ) - .await; + let model_directive_resolution = input + .model_directive_policy + .resolve_reasoning(provider_api_format, Some(&input.requested_model)); + let model_directive_mapping = match model_directive_resolution + .mapping_patch_for_mapped_model(&prepared_candidate.mapped_model) + { + Ok(mapping) => mapping, + Err(skip_reason) => { + mark_skipped_local_standard_candidate( + state, + input, + trace_id, + candidate, + attempt.candidate_index, + &attempt.candidate_id, + skip_reason, + ) + .await; + return Ok(None); + } + }; let redaction = resolve_provider_chat_pii_redaction( state, parts, @@ -588,7 +610,7 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts( }, Some(input.auth_context.api_key_id.as_str()), Some(effective_headers), - enable_model_directives, + false, ) { Some(body) => body, None => { @@ -655,18 +677,8 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts( provider_api_format, Some(body_json), ); - if let Some(mapping) = - crate::system_features::reasoning_model_directive_mapping_for_api_format_and_model( - state, - provider_api_format, - Some(&input.requested_model), - ) - .await - { - crate::ai_serving::apply_model_directive_mapping_patch( - &mut provider_request_body, - &mapping, - ); + if let Some(mapping) = model_directive_mapping.as_ref() { + crate::ai_serving::apply_model_directive_mapping_patch(&mut provider_request_body, mapping); // Directive mapping is a deep-merge patch and may overwrite/add `stream`; // re-enforce stream-field policy afterward. enforce_provider_body_stream_policy( @@ -712,6 +724,64 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts( ); } + let normalized_provider_api_format = + crate::ai_serving::normalize_api_format_alias(provider_api_format); + if matches!( + normalized_provider_api_format.as_str(), + "openai:chat" | "openai:responses" | "openai:responses:compact" + ) { + let source_model = body_json + .get("model") + .and_then(Value::as_str) + .unwrap_or(input.requested_model.as_str()); + let codex_model_capabilities = codex_model_capabilities_for_transport( + transport, + provider_api_format, + prepared_candidate.mapped_model.as_str(), + source_model, + ); + if crate::ai_serving::finalize_openai_provider_request_with_codex_model_capabilities( + &mut provider_request_body, + crate::ai_serving::OpenAiProviderRequestFinalization { + source_api_format: spec_metadata.api_format, + provider_api_format, + provider_type: transport.provider.provider_type.as_str(), + provider_model: prepared_candidate.mapped_model.as_str(), + source_model, + body_rules: transport.endpoint.body_rules.as_ref(), + upstream_is_stream, + require_body_stream_field: request_requires_body_stream_field( + body_json, + force_body_stream_field, + ), + }, + codex_model_capabilities.as_ref(), + ) + .is_err() + { + mark_skipped_local_standard_candidate_with_extra_data( + state, + input, + trace_id, + candidate, + attempt.candidate_index, + &attempt.candidate_id, + "provider_request_body_build_failed", + request_conversion_failure_extra_data( + body_json, + spec_metadata.api_format, + provider_api_format, + Some(prepared_candidate.mapped_model.as_str()), + Some(parts.uri.path()), + upstream_is_stream, + "standard_family_request_finalization", + ), + ) + .await; + return Ok(None); + } + } + if let Some(kiro_auth) = kiro_auth.as_ref() { return Ok(build_kiro_cross_format_payload_parts( state, @@ -752,8 +822,6 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts( .await); } - let normalized_provider_api_format = - crate::ai_serving::normalize_api_format_alias(provider_api_format); if normalized_provider_api_format == "gemini:generate_content" && is_gemini_cli_provider_transport(transport) { @@ -838,7 +906,7 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts( return Ok(None); }; let mut provider_request_headers = resolved_headers.headers; - apply_codex_openai_responses_special_headers( + apply_codex_openai_special_headers( &mut provider_request_headers, &provider_request_body, effective_headers, @@ -988,7 +1056,7 @@ async fn build_gemini_cli_cross_format_payload_parts( }; let mut provider_request_headers = resolved.headers.headers; - apply_codex_openai_responses_special_headers( + apply_codex_openai_special_headers( &mut provider_request_headers, &resolved.body, effective_headers, @@ -1146,6 +1214,7 @@ async fn resolve_local_gemini_image_to_openai_image_candidate_payload_parts( body_json: &serde_json::Value, input: &LocalStandardDecisionInput, attempt: &LocalStandardCandidateAttempt, + client_requires_streaming: bool, ) -> Option { let client_api_format = "gemini:generate_content"; let provider_api_format = "openai:image"; @@ -1221,9 +1290,44 @@ async fn resolve_local_gemini_image_to_openai_image_candidate_payload_parts( return None; }; - let upstream_is_stream = true; - let upstream_url = - build_openai_image_upstream_url(transport, Some("/v1/images/generations"), None); + let upstream_is_stream = resolve_upstream_is_stream_for_provider( + transport.endpoint.config.as_ref(), + transport.provider.provider_type.as_str(), + provider_api_format, + client_requires_streaming && candidate.supports_streaming, + false, + ); + let is_codex = transport + .provider + .provider_type + .trim() + .eq_ignore_ascii_case("codex"); + let mut provider_request_body = converted.body_json; + if upstream_is_stream { + provider_request_body + .as_object_mut()? + .insert("stream".to_string(), Value::Bool(true)); + } + provider_request_body = project_openai_image_api_request_body( + &provider_request_body, + &prepared_candidate.mapped_model, + converted.operation, + crate::image_capabilities::openai_image_provider_max_generation_count_for_model( + transport.provider.provider_type.as_str(), + Some(prepared_candidate.mapped_model.as_str()), + ), + )?; + if is_codex { + provider_request_body = project_codex_openai_image_api_request_body( + &provider_request_body, + converted.operation, + )?; + } + let request_path = match converted.operation { + OpenAiImageOperation::Generate => "/v1/images/generations", + OpenAiImageOperation::Edit => "/v1/images/edits", + }; + let upstream_url = build_openai_image_upstream_url(transport, Some(request_path), None); let effective_headers = input.effective_headers(&parts.headers); let Some(mut provider_request_headers) = build_openai_image_headers(ProviderOpenAiImageHeadersInput { @@ -1231,9 +1335,15 @@ async fn resolve_local_gemini_image_to_openai_image_candidate_payload_parts( headers: effective_headers, auth_header: &prepared_candidate.auth_header, auth_value: &prepared_candidate.auth_value, - accept: "text/event-stream", + accept: if is_codex { + None + } else if upstream_is_stream { + Some("text/event-stream") + } else { + Some("application/json") + }, header_rules: transport.endpoint.header_rules.as_ref(), - provider_request_body: &converted.body_json, + provider_request_body: &provider_request_body, original_request_body: body_json, }) else { @@ -1254,9 +1364,9 @@ async fn resolve_local_gemini_image_to_openai_image_candidate_payload_parts( .await; return None; }; - apply_codex_openai_responses_special_headers( + apply_codex_openai_special_headers( &mut provider_request_headers, - &converted.body_json, + &provider_request_body, effective_headers, transport.provider.provider_type.as_str(), provider_api_format, @@ -1269,7 +1379,7 @@ async fn resolve_local_gemini_image_to_openai_image_candidate_payload_parts( auth_value: prepared_candidate.auth_value, mapped_model: converted.mapped_model, provider_api_format: provider_api_format.to_string(), - provider_request_body: converted.body_json, + provider_request_body, provider_request_headers, upstream_url, upstream_is_stream, 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 45796eee2..03d9699a6 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/mod.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/mod.rs @@ -15,7 +15,8 @@ mod normalize; mod openai; pub(crate) use self::codex::{ - apply_codex_openai_responses_special_body_edits, apply_codex_openai_responses_special_headers, + apply_codex_openai_responses_special_body_edits, apply_codex_openai_special_headers, + codex_model_capabilities_for_transport, }; pub(crate) use self::deepseek::{apply_deepseek_tool_call_thinking_compat, is_deepseek_provider}; pub(crate) use self::family::{ @@ -25,9 +26,11 @@ pub(crate) use self::family::{ pub(crate) use self::normalize::{ build_cross_format_openai_chat_request_body, build_cross_format_openai_chat_upstream_url, build_cross_format_openai_responses_request_body, + build_cross_format_openai_responses_request_body_with_codex_model_capabilities, build_cross_format_openai_responses_upstream_url, build_local_openai_chat_request_body, build_local_openai_chat_upstream_url, build_local_openai_responses_request_body, - build_local_openai_responses_upstream_url, + build_local_openai_responses_request_body_with_codex_model_capabilities, + build_local_openai_responses_upstream_url, validate_final_openai_provider_request, }; pub(crate) use self::openai::{ build_local_openai_chat_stream_attempt_source_for_kind, @@ -297,7 +300,7 @@ mod tests { let converted = build_standard_request_body( &request, "claude:messages", - "gpt-5", + "gpt-5.4", "codex", "openai:responses", "/v1/messages", @@ -309,7 +312,7 @@ mod tests { assert!(converted.get("metadata").is_none()); assert_eq!(converted["store"], false); - assert_eq!(converted["instructions"], ""); + assert!(converted.get("instructions").is_none()); assert_eq!(converted["include"], json!(["reasoning.encrypted_content"])); assert_eq!(converted["parallel_tool_calls"], true); assert_eq!(converted["reasoning"]["effort"], "medium"); diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/normalize.rs b/apps/aether-gateway/src/ai_serving/planner/standard/normalize.rs index 906caec11..898497435 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/normalize.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/normalize.rs @@ -12,9 +12,38 @@ pub(crate) use self::chat::{ }; pub(crate) use self::responses::{ build_cross_format_openai_responses_request_body, + build_cross_format_openai_responses_request_body_with_codex_model_capabilities, build_cross_format_openai_responses_upstream_url, build_local_openai_responses_request_body, + build_local_openai_responses_request_body_with_codex_model_capabilities, build_local_openai_responses_upstream_url, }; pub(super) use crate::ai_serving::planner::common::{ enforce_provider_body_stream_policy, request_requires_body_stream_field, }; + +pub(crate) fn validate_final_openai_provider_request( + provider_api_format: &str, + mapped_model: &str, + source_request_body: &serde_json::Value, + provider_request_body: &serde_json::Value, +) -> Option<()> { + let provider_model = provider_request_body + .get("model") + .and_then(serde_json::Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .unwrap_or(mapped_model); + let source_model = source_request_body + .get("model") + .and_then(serde_json::Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .unwrap_or(mapped_model); + crate::ai_serving::validate_openai_provider_request_contract( + provider_api_format, + provider_model, + source_model, + provider_request_body, + ) + .ok() +} diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/normalize/chat.rs b/apps/aether-gateway/src/ai_serving/planner/standard/normalize/chat.rs index 35a8a9549..c113e88f6 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/normalize/chat.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/normalize/chat.rs @@ -9,7 +9,10 @@ use crate::ai_serving::{ GatewayProviderTransportSnapshot, }; -use super::{enforce_provider_body_stream_policy, request_requires_body_stream_field}; +use super::{ + enforce_provider_body_stream_policy, request_requires_body_stream_field, + validate_final_openai_provider_request, +}; pub(crate) fn build_local_openai_chat_request_body( body_json: &Value, @@ -39,6 +42,12 @@ pub(crate) fn build_local_openai_chat_request_body( upstream_is_stream, request_requires_body_stream_field(body_json, force_body_stream_field), ); + validate_final_openai_provider_request( + "openai:chat", + mapped_model, + body_json, + &provider_request_body, + )?; Some(provider_request_body) } @@ -92,6 +101,12 @@ pub(crate) fn build_cross_format_openai_chat_request_body( upstream_is_stream, request_requires_body_stream_field(body_json, force_body_stream_field), ); + validate_final_openai_provider_request( + provider_api_format, + mapped_model, + body_json, + &provider_request_body, + )?; Some(provider_request_body) } diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/normalize/responses.rs b/apps/aether-gateway/src/ai_serving/planner/standard/normalize/responses.rs index 2042eac30..dabdc5895 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/normalize/responses.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/normalize/responses.rs @@ -2,14 +2,16 @@ use serde_json::Value; use crate::ai_serving::transport::apply_standard_provider_request_body_rules_with_request_headers; use crate::ai_serving::{ - apply_codex_openai_responses_special_body_edits, apply_openai_responses_compact_special_body_edits, build_cross_format_openai_responses_request_body_with_model_directives as surface_build_cross_format_openai_responses_request_body, build_local_openai_responses_request_body_with_model_directives as surface_build_local_openai_responses_request_body, GatewayProviderTransportSnapshot, }; -use super::{enforce_provider_body_stream_policy, request_requires_body_stream_field}; +use super::{ + enforce_provider_body_stream_policy, request_requires_body_stream_field, + validate_final_openai_provider_request, +}; pub(crate) fn build_local_openai_responses_request_body( body_json: &Value, @@ -19,9 +21,35 @@ pub(crate) fn build_local_openai_responses_request_body( provider_type: &str, provider_api_format: &str, body_rules: Option<&Value>, - user_api_key_id: Option<&str>, + _user_api_key_id: Option<&str>, request_headers: &http::HeaderMap, enable_model_directives: bool, +) -> Option { + build_local_openai_responses_request_body_with_codex_model_capabilities( + body_json, + mapped_model, + require_streaming, + force_body_stream_field, + provider_type, + provider_api_format, + body_rules, + request_headers, + None, + enable_model_directives, + ) +} + +pub(crate) fn build_local_openai_responses_request_body_with_codex_model_capabilities( + body_json: &Value, + mapped_model: &str, + require_streaming: bool, + force_body_stream_field: bool, + provider_type: &str, + provider_api_format: &str, + body_rules: Option<&Value>, + request_headers: &http::HeaderMap, + model_capabilities: Option<&crate::ai_serving::CodexResponsesModelCapabilities>, + enable_model_directives: bool, ) -> Option { let provider_request_body = surface_build_local_openai_responses_request_body( body_json, @@ -36,12 +64,18 @@ pub(crate) fn build_local_openai_responses_request_body( body_json, request_headers, )?; - apply_codex_openai_responses_special_body_edits( + let source_model = body_json + .get("model") + .and_then(Value::as_str) + .unwrap_or(mapped_model); + crate::ai_serving::apply_codex_openai_responses_special_body_edits_with_source_model_and_capabilities( &mut provider_request_body, provider_type, provider_api_format, + mapped_model, + source_model, + model_capabilities, body_rules, - user_api_key_id, ); apply_openai_responses_compact_special_body_edits( &mut provider_request_body, @@ -53,6 +87,12 @@ pub(crate) fn build_local_openai_responses_request_body( require_streaming, request_requires_body_stream_field(body_json, force_body_stream_field), ); + validate_final_openai_provider_request( + provider_api_format, + mapped_model, + body_json, + &provider_request_body, + )?; Some(provider_request_body) } @@ -65,9 +105,37 @@ pub(crate) fn build_cross_format_openai_responses_request_body( force_body_stream_field: bool, provider_type: &str, body_rules: Option<&Value>, - user_api_key_id: Option<&str>, + _user_api_key_id: Option<&str>, request_headers: &http::HeaderMap, enable_model_directives: bool, +) -> Option { + build_cross_format_openai_responses_request_body_with_codex_model_capabilities( + body_json, + mapped_model, + client_api_format, + provider_api_format, + upstream_is_stream, + force_body_stream_field, + provider_type, + body_rules, + request_headers, + None, + enable_model_directives, + ) +} + +pub(crate) fn build_cross_format_openai_responses_request_body_with_codex_model_capabilities( + body_json: &Value, + mapped_model: &str, + client_api_format: &str, + provider_api_format: &str, + upstream_is_stream: bool, + force_body_stream_field: bool, + provider_type: &str, + body_rules: Option<&Value>, + request_headers: &http::HeaderMap, + model_capabilities: Option<&crate::ai_serving::CodexResponsesModelCapabilities>, + enable_model_directives: bool, ) -> Option { let provider_request_body = surface_build_cross_format_openai_responses_request_body( body_json, @@ -84,12 +152,18 @@ pub(crate) fn build_cross_format_openai_responses_request_body( body_json, request_headers, )?; - apply_codex_openai_responses_special_body_edits( + let source_model = body_json + .get("model") + .and_then(Value::as_str) + .unwrap_or(mapped_model); + crate::ai_serving::apply_codex_openai_responses_special_body_edits_with_source_model_and_capabilities( &mut provider_request_body, provider_type, provider_api_format, + mapped_model, + source_model, + model_capabilities, body_rules, - user_api_key_id, ); apply_openai_responses_compact_special_body_edits( &mut provider_request_body, @@ -101,6 +175,12 @@ pub(crate) fn build_cross_format_openai_responses_request_body( upstream_is_stream, request_requires_body_stream_field(body_json, force_body_stream_field), ); + validate_final_openai_provider_request( + provider_api_format, + mapped_model, + body_json, + &provider_request_body, + )?; Some(provider_request_body) } diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/normalize/tests.rs b/apps/aether-gateway/src/ai_serving/planner/standard/normalize/tests.rs index 6694b3657..0297f282b 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/normalize/tests.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/normalize/tests.rs @@ -6,8 +6,8 @@ use http::Request; use serde_json::{json, Value}; use super::{ - build_cross_format_openai_responses_request_body, build_local_openai_responses_request_body, - build_local_openai_responses_upstream_url, + build_cross_format_openai_responses_request_body, build_local_openai_chat_request_body, + build_local_openai_responses_request_body, build_local_openai_responses_upstream_url, }; fn object_keys(value: &Value) -> Vec<&str> { @@ -146,12 +146,10 @@ fn local_openai_responses_wrapper_preserves_body_order_after_edits() { "reasoning", "tool_choice", "parallel_tool_calls", - "instructions", - "prompt_cache_key", ] ); assert_eq!(provider_request_body["parallel_tool_calls"], json!(true)); - assert_eq!(provider_request_body["instructions"], json!("")); + assert!(provider_request_body.get("instructions").is_none()); } #[test] @@ -181,18 +179,49 @@ fn local_openai_responses_compact_wrapper_strips_store_for_same_format_requests( } #[test] -fn local_openai_responses_compact_wrapper_strips_include_for_codex_requests() { +fn local_codex_compact_wrapper_applies_the_complete_request_projection() { let body_json = json!({ - "model": "gpt-5.4", - "input": [], + "model": "gpt-5.6-sol", + "input": [{ + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "hello"}] + }], + "instructions": "Work carefully", + "client_metadata": {"origin": "codex"}, "include": ["reasoning.encrypted_content"], "store": true, - "stream": true + "stream": true, + "stream_options": {"reasoning_summary_delivery": "sequential_cutoff"}, + "tool_choice": "auto", + "parallel_tool_calls": true, + "reasoning": {"effort": "max", "summary": "auto", "context": "all_turns"}, + "text": {"verbosity": "medium"}, + "tools": [{ + "type": "function", + "name": "lookup", + "parameters": {"type": "object", "properties": {}} + }], + "service_tier": "priority", + "prompt_cache_key": "thread-compact" }); - let provider_request_body = build_local_openai_responses_request_body( + let regular = build_local_openai_responses_request_body( &body_json, - "gpt-5.4", + "gpt-5.6-sol", + true, + false, + "codex", + "openai:responses", + None, + Some("key-123"), + &http::HeaderMap::new(), + false, + ) + .expect("local Codex Responses body should build"); + let compact = build_local_openai_responses_request_body( + &body_json, + "gpt-5.6-sol", false, false, "codex", @@ -202,22 +231,44 @@ fn local_openai_responses_compact_wrapper_strips_include_for_codex_requests() { &http::HeaderMap::new(), false, ) - .expect("local codex compact body should build"); + .expect("local Codex Compact body should build"); - assert!(provider_request_body.get("include").is_none()); - assert!(provider_request_body.get("store").is_none()); - assert!(provider_request_body.get("stream").is_none()); - assert_eq!(provider_request_body["instructions"], ""); - assert_eq!( - provider_request_body["prompt_cache_key"], - "3d2e2842-74cb-55dd-803a-b8940b3500c2" - ); + for field in [ + "client_metadata", + "include", + "store", + "stream", + "stream_options", + "tool_choice", + ] { + assert!( + regular.get(field).is_some(), + "Responses should contain {field}" + ); + assert!(compact.get(field).is_none(), "Compact should omit {field}"); + } + for field in [ + "model", + "input", + "instructions", + "parallel_tool_calls", + "reasoning", + "text", + "tools", + "service_tier", + "prompt_cache_key", + ] { + assert_eq!( + compact[field], regular[field], + "Compact should preserve {field}" + ); + } } #[test] fn local_openai_responses_wrapper_applies_model_directive_before_body_rules() { let body_json = json!({ - "model": "gpt-5.4-max", + "model": "gpt-5.6-sol-max", "input": "hello", "reasoning": {"effort": "low", "summary": "auto"} }); @@ -227,7 +278,7 @@ fn local_openai_responses_wrapper_applies_model_directive_before_body_rules() { let provider_request_body = build_local_openai_responses_request_body( &body_json, - "gpt-5.4", + "gpt-5.6-sol", false, false, "openai", @@ -244,6 +295,132 @@ fn local_openai_responses_wrapper_applies_model_directive_before_body_rules() { assert_eq!(provider_request_body["metadata"]["override_seen"], true); } +#[test] +fn final_openai_provider_contract_uses_the_mapped_model_for_reasoning() { + let alias = json!({ + "model": "deployment-alias", + "input": "hello", + "reasoning": {"effort": "max"} + }); + assert!(build_local_openai_responses_request_body( + &alias, + "gpt-5.6-sol", + false, + false, + "openai", + "openai:responses", + None, + None, + &http::HeaderMap::new(), + false, + ) + .is_some()); + assert!(build_local_openai_responses_request_body( + &alias, + "gpt-5.4", + false, + false, + "openai", + "openai:responses", + None, + None, + &http::HeaderMap::new(), + false, + ) + .is_none()); + + let minimal = json!({ + "model": "deployment-alias", + "messages": [{"role": "user", "content": "hello"}], + "reasoning_effort": "minimal" + }); + assert!(build_local_openai_chat_request_body( + &minimal, + "gpt-5.6-terra", + false, + false, + None, + &http::HeaderMap::new(), + false, + ) + .is_none()); + + let opaque_mapping = json!({ + "model": "gpt-5.6-sol-max", + "input": "hello", + "reasoning": {"effort": "max", "mode": "pro"}, + "prompt_cache_options": {"mode": "explicit", "ttl": "30m"} + }); + assert!(build_local_openai_responses_request_body( + &opaque_mapping, + "azure-production", + false, + false, + "openai", + "openai:responses", + None, + None, + &http::HeaderMap::new(), + false, + ) + .is_some()); + assert!(build_local_openai_responses_request_body( + &opaque_mapping, + "gpt-5.4", + false, + false, + "openai", + "openai:responses", + None, + None, + &http::HeaderMap::new(), + false, + ) + .is_none()); +} + +#[test] +fn final_openai_provider_contract_validates_body_rule_output() { + let body = json!({ + "model": "gpt-5.6-sol", + "input": "hello", + "reasoning": {"effort": "max"} + }); + let model_override = json!([ + {"action":"set","path":"model","value":"gpt-5.4"} + ]); + assert!(build_local_openai_responses_request_body( + &body, + "gpt-5.6-sol", + false, + false, + "openai", + "openai:responses", + Some(&model_override), + None, + &http::HeaderMap::new(), + false, + ) + .is_none()); + + let cache_override = json!([ + {"action":"set","path":"prompt_cache_options.ttl","value":"1h"} + ]); + assert!(build_local_openai_responses_request_body( + &json!({"model":"gpt-5.6-sol","input":"hello"}), + "gpt-5.6-sol", + false, + false, + "openai", + "openai:responses", + Some(&cache_override), + None, + &http::HeaderMap::new(), + false, + ) + .is_none()); +} + #[test] fn local_openai_responses_upstream_url_preserves_codex_base_path() { let request = Request::builder() @@ -371,7 +548,7 @@ fn applies_codex_defaults_unless_body_rules_handle_the_field() { } #[test] -fn injects_codex_prompt_cache_key_for_openai_responses_cross_format_requests() { +fn omits_codex_prompt_cache_key_for_openai_responses_cross_format_requests() { let body_json = json!({ "model": "claude-sonnet-4-5", "messages": [{ @@ -395,14 +572,11 @@ fn injects_codex_prompt_cache_key_for_openai_responses_cross_format_requests() { ) .expect("claude cli to codex request should build"); - assert_eq!( - provider_request_body["prompt_cache_key"], - "4ee6ea6e-3ac6-5a18-8cb8-1f8b956419e5" - ); + assert!(provider_request_body.get("prompt_cache_key").is_none()); } #[test] -fn injects_codex_prompt_cache_key_for_openai_chat_cross_format_requests() { +fn omits_codex_prompt_cache_key_for_openai_chat_cross_format_requests() { let body_json = json!({ "model": "gpt-5", "messages": [{ @@ -425,8 +599,5 @@ fn injects_codex_prompt_cache_key_for_openai_chat_cross_format_requests() { ) .expect("openai chat to codex request should build"); - assert_eq!( - provider_request_body["prompt_cache_key"], - "4ee6ea6e-3ac6-5a18-8cb8-1f8b956419e5" - ); + assert!(provider_request_body.get("prompt_cache_key").is_none()); } diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/payload.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/payload.rs index f035425f5..113829caf 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/payload.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/payload.rs @@ -6,7 +6,7 @@ use crate::ai_serving::planner::report_context::{ insert_provider_stream_event_api_format, LocalExecutionReportContextParts, }; use crate::ai_serving::planner::{ - build_ai_execution_decision_response, resolve_transport_request_gzip_policy, + build_ai_execution_decision_response, resolve_transport_request_encoding_policy, AiExecutionDecisionResponseParts, }; use crate::ai_serving::transport::{ @@ -44,6 +44,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate candidate_id, .. } = attempt; + let upstream_is_stream = upstream_is_stream && eligible.candidate.supports_streaming; let payload_started_at = std::time::Instant::now(); let Some(resolved) = resolve_local_openai_chat_candidate_payload_parts( state, @@ -72,6 +73,14 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate payload_started_at.elapsed().as_millis() as u64, ); let candidate = &eligible.candidate; + let upstream_is_stream = + crate::ai_serving::planner::common::resolve_upstream_is_stream_for_provider( + resolved.transport.endpoint.config.as_ref(), + resolved.transport.provider.provider_type.as_str(), + resolved.provider_api_format.as_str(), + upstream_is_stream, + false, + ); let prompt_cache_key = resolved .provider_request_body @@ -208,7 +217,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate "stream_candidate_report_context", report_context_started_at.elapsed().as_millis() as u64, ); - let request_gzip = resolve_transport_request_gzip_policy(&transport); + let request_encoding = resolve_transport_request_encoding_policy(&transport); let decision_started_at = std::time::Instant::now(); let mut decision = build_ai_execution_decision_response(AiExecutionDecisionResponseParts { @@ -219,6 +228,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate request_id: trace_id.to_string(), candidate_id: candidate_id.clone(), provider_name: transport.provider.name.clone(), + provider_type: transport.provider.provider_type.clone(), provider_id: candidate.provider_id.clone(), endpoint_id: candidate.endpoint_id.clone(), key_id: candidate.key_id.clone(), @@ -236,8 +246,8 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate provider_request_body: Some(provider_request_body), provider_request_body_base64: None, content_type: Some("application/json".to_string()), - content_encoding: None, - request_gzip, + content_encoding: request_encoding.content_encoding, + request_gzip: request_encoding.request_gzip, proxy, transport_profile, timeouts, @@ -246,7 +256,11 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate report_context: Some(report_context), auth_context: input.auth_context.clone(), }); - apply_provider_request_routing_policy_to_decision(input, &mut decision)?; + apply_provider_request_routing_policy_to_decision( + input, + &mut decision, + Some(transport.as_ref()), + )?; observe_gateway_stage_ms( "stream_candidate_decision_build", decision_started_at.elapsed().as_millis() as u64, 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 ffb393206..b8e142673 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 @@ -25,11 +25,11 @@ use crate::ai_serving::planner::redaction::{ request_identity_response_encoding_when_redacted, resolve_provider_chat_pii_redaction, }; use crate::ai_serving::planner::standard::{ - apply_codex_openai_responses_special_body_edits, apply_codex_openai_responses_special_headers, + apply_codex_openai_responses_special_body_edits, apply_codex_openai_special_headers, apply_deepseek_tool_call_thinking_compat, build_cross_format_openai_chat_request_body, build_cross_format_openai_chat_upstream_url, build_local_openai_chat_request_body, - build_local_openai_chat_upstream_url, request_body_build_failure_extra_data, - request_conversion_failure_extra_data, + build_local_openai_chat_upstream_url, codex_model_capabilities_for_transport, + request_body_build_failure_extra_data, request_conversion_failure_extra_data, }; use crate::ai_serving::transport::antigravity::is_antigravity_provider_transport; use crate::ai_serving::transport::auth::resolve_local_openai_bearer_auth; @@ -56,7 +56,10 @@ use crate::ai_serving::transport::{ use crate::ai_serving::{ ai_local_execution_contract_for_formats, request_conversion_direct_auth, request_conversion_kind, CandidateFailureDiagnostic, GatewayProviderTransportSnapshot, - LocalResolvedOAuthRequestAuth, + LocalResolvedOAuthRequestAuth, OpenAiImageOperation, +}; +use crate::ai_serving::{ + project_codex_openai_image_api_request_body, project_openai_image_api_request_body, }; use crate::ai_serving::{ConversionMode, ExecutionStrategy}; use crate::stage_metrics::observe_gateway_stage_ms; @@ -88,37 +91,7 @@ pub(crate) struct LocalOpenAiChatCandidatePayloadParts { } #[derive(Default)] -pub(crate) struct LocalOpenAiChatRequestPreparation { - model_directives_enabled: BTreeMap<(String, String), bool>, -} - -impl LocalOpenAiChatRequestPreparation { - async fn model_directives_enabled( - &mut self, - state: &AppState, - provider_api_format: &str, - requested_model: &str, - ) -> bool { - let key = ( - provider_api_format.trim().to_ascii_lowercase(), - requested_model.trim().to_string(), - ); - if let Some(enabled) = self.model_directives_enabled.get(&key) { - crate::stage_metrics::record_openai_chat_model_directive_cache_hit(); - return *enabled; - } - crate::stage_metrics::record_openai_chat_model_directive_cache_miss(); - let enabled = - crate::system_features::reasoning_model_directive_enabled_for_api_format_and_model( - state, - provider_api_format, - Some(requested_model), - ) - .await; - self.model_directives_enabled.insert(key, enabled); - enabled - } -} +pub(crate) struct LocalOpenAiChatRequestPreparation; fn is_grok_text_provider_api_format(provider_api_format: &str) -> bool { matches!( @@ -127,6 +100,65 @@ fn is_grok_text_provider_api_format(provider_api_format: &str) -> bool { ) } +fn finalize_openai_chat_provider_request_body( + provider_request_body: &mut Value, + custom_directive_mapping: Option<&Value>, + provider_api_format: &str, + upstream_is_stream: bool, + force_body_stream_field: bool, + original_body: &Value, + transport: &GatewayProviderTransportSnapshot, + mapped_model: &str, +) -> bool { + if let Some(mapping) = custom_directive_mapping { + crate::ai_serving::apply_model_directive_mapping_patch(provider_request_body, mapping); + } + + // Mapping and endpoint body rules can both write `stream`. The resolved transport + // policy is authoritative and therefore runs after every body mutation. + enforce_provider_body_stream_policy( + provider_request_body, + provider_api_format, + upstream_is_stream, + request_requires_body_stream_field(original_body, force_body_stream_field), + ); + apply_deepseek_tool_call_thinking_compat( + provider_request_body, + transport.provider.provider_type.as_str(), + transport.endpoint.base_url.as_str(), + provider_api_format, + Some(original_body), + ); + let source_model = original_body + .get("model") + .and_then(Value::as_str) + .unwrap_or(mapped_model); + let codex_model_capabilities = codex_model_capabilities_for_transport( + transport, + provider_api_format, + mapped_model, + source_model, + ); + crate::ai_serving::finalize_openai_provider_request_with_codex_model_capabilities( + provider_request_body, + crate::ai_serving::OpenAiProviderRequestFinalization { + source_api_format: "openai:chat", + provider_api_format, + provider_type: transport.provider.provider_type.as_str(), + provider_model: mapped_model, + source_model, + body_rules: transport.endpoint.body_rules.as_ref(), + upstream_is_stream, + require_body_stream_field: request_requires_body_stream_field( + original_body, + force_body_stream_field, + ), + }, + codex_model_capabilities.as_ref(), + ) + .is_ok() +} + #[allow(clippy::too_many_arguments)] pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts( state: &AppState, @@ -134,7 +166,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts( trace_id: &str, body_json: &serde_json::Value, input: &LocalOpenAiChatDecisionInput, - mut preparation: Option<&mut LocalOpenAiChatRequestPreparation>, + _preparation: Option<&mut LocalOpenAiChatRequestPreparation>, eligible: &EligibleLocalExecutionCandidate, candidate_index: u32, candidate_id: &str, @@ -151,18 +183,9 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts( let force_body_stream_field = endpoint_config_forces_body_stream_field(transport.endpoint.config.as_ref()); let model_directives_started_at = std::time::Instant::now(); - let enable_model_directives = if let Some(preparation) = preparation { - preparation - .model_directives_enabled(state, provider_api_format, &input.requested_model) - .await - } else { - crate::system_features::reasoning_model_directive_enabled_for_api_format_and_model( - state, - provider_api_format, - Some(&input.requested_model), - ) - .await - }; + let model_directive_resolution = input + .model_directive_policy + .resolve_reasoning(provider_api_format, Some(&input.requested_model)); observe_gateway_stage_ms( "openai_chat_payload_model_directives", model_directives_started_at.elapsed().as_millis() as u64, @@ -220,15 +243,33 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts( return Ok(None); } }; + let model_directive_mapping = match model_directive_resolution + .mapping_patch_for_mapped_model(&prepared_candidate.mapped_model) + { + Ok(mapping) => mapping, + Err(skip_reason) => { + mark_skipped_local_openai_chat_candidate( + state, + input, + trace_id, + candidate, + candidate_index, + candidate_id, + skip_reason, + ) + .await; + return Ok(None); + } + }; - let Some(provider_request_body) = build_local_openai_chat_request_body( + let Some(mut provider_request_body) = build_local_openai_chat_request_body( body_json, &prepared_candidate.mapped_model, upstream_is_stream, force_body_stream_field, transport.endpoint.body_rules.as_ref(), effective_headers, - enable_model_directives, + false, ) else { mark_skipped_local_openai_chat_candidate_with_extra_data( state, @@ -247,6 +288,33 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts( .await; return Ok(None); }; + if !finalize_openai_chat_provider_request_body( + &mut provider_request_body, + model_directive_mapping.as_ref(), + provider_api_format, + upstream_is_stream, + force_body_stream_field, + body_json, + transport, + &prepared_candidate.mapped_model, + ) { + mark_skipped_local_openai_chat_candidate_with_extra_data( + state, + input, + trace_id, + candidate, + candidate_index, + candidate_id, + "provider_request_body_build_failed", + request_body_build_failure_extra_data( + body_json, + "openai:chat", + provider_api_format, + ), + ) + .await; + return Ok(None); + } let upstream_url = build_grok_upstream_url(transport, GROK_CHAT_PATH); let Some(mut provider_request_headers) = build_grok_browser_headers(GrokHeaderInput { @@ -379,6 +447,24 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts( return Ok(None); } }; + let model_directive_mapping = match model_directive_resolution + .mapping_patch_for_mapped_model(&prepared_candidate.mapped_model) + { + Ok(mapping) => mapping, + Err(skip_reason) => { + mark_skipped_local_openai_chat_candidate( + state, + input, + trace_id, + candidate, + candidate_index, + candidate_id, + skip_reason, + ) + .await; + return Ok(None); + } + }; observe_gateway_stage_ms( "openai_chat_payload_auth_prepare", auth_prepare_started_at.elapsed().as_millis() as u64, @@ -392,7 +478,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts( force_body_stream_field, transport.endpoint.body_rules.as_ref(), effective_headers, - enable_model_directives, + false, ) else { mark_skipped_local_openai_chat_candidate_with_extra_data( state, @@ -415,13 +501,33 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts( "openai_chat_payload_body_build", body_build_started_at.elapsed().as_millis() as u64, ); - apply_deepseek_tool_call_thinking_compat( + if !finalize_openai_chat_provider_request_body( &mut provider_request_body, - transport.provider.provider_type.as_str(), - transport.endpoint.base_url.as_str(), + model_directive_mapping.as_ref(), "openai:chat", - Some(body_json), - ); + upstream_is_stream, + force_body_stream_field, + body_json, + transport, + &prepared_candidate.mapped_model, + ) { + mark_skipped_local_openai_chat_candidate_with_extra_data( + state, + input, + trace_id, + candidate, + candidate_index, + candidate_id, + "provider_request_body_build_failed", + request_body_build_failure_extra_data( + body_json, + "openai:chat", + provider_api_format, + ), + ) + .await; + return Ok(None); + } let Some(upstream_url) = build_local_openai_chat_upstream_url(parts, transport) else { mark_skipped_local_openai_chat_candidate_with_failure_diagnostic( @@ -475,7 +581,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts( return Ok(None); }; let mut provider_request_headers = resolved_headers.headers; - apply_codex_openai_responses_special_headers( + apply_codex_openai_special_headers( &mut provider_request_headers, &provider_request_body, effective_headers, @@ -651,6 +757,24 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts( } } }; + let model_directive_mapping = match model_directive_resolution + .mapping_patch_for_mapped_model(&prepared_candidate.mapped_model) + { + Ok(mapping) => mapping, + Err(skip_reason) => { + mark_skipped_local_openai_chat_candidate( + state, + input, + trace_id, + candidate, + candidate_index, + candidate_id, + skip_reason, + ) + .await; + return Ok(None); + } + }; let Some(mut provider_request_body) = build_cross_format_openai_chat_request_body( body_json, @@ -666,7 +790,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts( }, Some(input.auth_context.api_key_id.as_str()), effective_headers, - enable_model_directives, + false, ) else { mark_skipped_local_openai_chat_candidate_with_extra_data( state, @@ -689,34 +813,37 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts( .await; return Ok(None); }; - if let Some(mapping) = - crate::system_features::reasoning_model_directive_mapping_for_api_format_and_model( - state, - provider_api_format.as_str(), - Some(&input.requested_model), - ) - .await - { - crate::ai_serving::apply_model_directive_mapping_patch( - &mut provider_request_body, - &mapping, - ); - // Directive mapping is a deep-merge patch and may overwrite/add `stream`; - // re-enforce stream-field policy afterward. - enforce_provider_body_stream_policy( - &mut provider_request_body, - provider_api_format.as_str(), - upstream_is_stream, - request_requires_body_stream_field(body_json, force_body_stream_field), - ); - } - apply_deepseek_tool_call_thinking_compat( + if !finalize_openai_chat_provider_request_body( &mut provider_request_body, - transport.provider.provider_type.as_str(), - transport.endpoint.base_url.as_str(), + model_directive_mapping.as_ref(), provider_api_format.as_str(), - Some(body_json), - ); + upstream_is_stream, + force_body_stream_field, + body_json, + transport, + &prepared_candidate.mapped_model, + ) { + mark_skipped_local_openai_chat_candidate_with_extra_data( + state, + input, + trace_id, + candidate, + candidate_index, + candidate_id, + "provider_request_body_build_failed", + request_conversion_failure_extra_data( + body_json, + "openai:chat", + provider_api_format.as_str(), + Some(prepared_candidate.mapped_model.as_str()), + Some(parts.uri.path()), + upstream_is_stream, + "openai_chat_request_conversion", + ), + ) + .await; + return Ok(None); + } if let Some(kiro_auth) = kiro_auth.as_ref() { return Ok(build_kiro_openai_chat_cross_format_payload_parts( @@ -845,7 +972,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts( return Ok(None); }; let mut provider_request_headers = resolved_headers.headers; - apply_codex_openai_responses_special_headers( + apply_codex_openai_special_headers( &mut provider_request_headers, &provider_request_body, effective_headers, @@ -997,7 +1124,7 @@ async fn build_antigravity_openai_chat_cross_format_payload_parts( } }; let mut provider_request_headers = resolved.headers.headers; - apply_codex_openai_responses_special_headers( + apply_codex_openai_special_headers( &mut provider_request_headers, &resolved.body, effective_headers, @@ -1148,7 +1275,7 @@ async fn build_gemini_cli_openai_chat_cross_format_payload_parts( } }; let mut provider_request_headers = resolved.headers.headers; - apply_codex_openai_responses_special_headers( + apply_codex_openai_special_headers( &mut provider_request_headers, &resolved.body, effective_headers, @@ -1255,6 +1382,19 @@ async fn resolve_openai_chat_to_openai_image_payload_parts( .provider_type .trim() .eq_ignore_ascii_case("chatgpt_web"); + let is_codex = transport + .provider + .provider_type + .trim() + .eq_ignore_ascii_case("codex"); + let upstream_is_stream = + crate::ai_serving::planner::common::resolve_upstream_is_stream_for_provider( + transport.endpoint.config.as_ref(), + transport.provider.provider_type.as_str(), + provider_api_format, + upstream_is_stream, + false, + ); let Some((mut provider_request_body, image_request_summary)) = (if is_chatgpt_web { build_chatgpt_web_image_provider_body_from_openai_chat_body( body_json, @@ -1263,7 +1403,7 @@ async fn resolve_openai_chat_to_openai_image_payload_parts( } else { build_openai_image_provider_body_from_openai_chat_body( body_json, - &input.requested_model, + &prepared_candidate.mapped_model, upstream_is_stream, ) }) else { @@ -1280,24 +1420,70 @@ async fn resolve_openai_chat_to_openai_image_payload_parts( .await; return Ok(None); }; + let Some(operation) = openai_image_operation_from_summary(&image_request_summary) else { + return Ok(None); + }; if !is_chatgpt_web { - apply_codex_openai_responses_special_body_edits( - &mut provider_request_body, - transport.provider.provider_type.as_str(), - provider_api_format, - transport.endpoint.body_rules.as_ref(), - Some(candidate.key_id.as_str()), - ); + let Some(projected) = project_openai_image_api_request_body( + &provider_request_body, + &prepared_candidate.mapped_model, + operation, + crate::image_capabilities::openai_image_provider_max_generation_count_for_model( + transport.provider.provider_type.as_str(), + Some(prepared_candidate.mapped_model.as_str()), + ), + ) else { + mark_skipped_local_openai_chat_candidate_with_extra_data( + state, + input, + trace_id, + candidate, + candidate_index, + candidate_id, + "provider_request_body_build_failed", + request_body_build_failure_extra_data( + body_json, + "openai:chat", + provider_api_format, + ), + ) + .await; + return Ok(None); + }; + provider_request_body = projected; + } + if is_codex { + let Some(projected) = + project_codex_openai_image_api_request_body(&provider_request_body, operation) + else { + mark_skipped_local_openai_chat_candidate_with_extra_data( + state, + input, + trace_id, + candidate, + candidate_index, + candidate_id, + "provider_request_body_build_failed", + request_body_build_failure_extra_data( + body_json, + "openai:chat", + provider_api_format, + ), + ) + .await; + return Ok(None); + }; + provider_request_body = projected; } let upstream_url = if is_chatgpt_web { chatgpt_web_image_internal_url(&transport.endpoint.base_url) } else { - build_openai_image_upstream_url( - transport, - Some("/v1/images/generations"), - parts.uri.query(), - ) + let request_path = match operation { + OpenAiImageOperation::Generate => "/v1/images/generations", + OpenAiImageOperation::Edit => "/v1/images/edits", + }; + build_openai_image_upstream_url(transport, Some(request_path), parts.uri.query()) }; let Some(mut provider_request_headers) = build_openai_image_headers(ProviderOpenAiImageHeadersInput { @@ -1305,7 +1491,13 @@ async fn resolve_openai_chat_to_openai_image_payload_parts( headers: &parts.headers, auth_header: &prepared_candidate.auth_header, auth_value: &prepared_candidate.auth_value, - accept: "text/event-stream", + accept: if is_codex { + None + } else if upstream_is_stream { + Some("text/event-stream") + } else { + Some("application/json") + }, header_rules: transport.endpoint.header_rules.as_ref(), provider_request_body: &provider_request_body, original_request_body: body_json, @@ -1331,7 +1523,7 @@ async fn resolve_openai_chat_to_openai_image_payload_parts( if is_chatgpt_web { provider_request_headers.insert("x-aether-chatgpt-web-image".to_string(), "1".to_string()); } else { - apply_codex_openai_responses_special_headers( + apply_codex_openai_special_headers( &mut provider_request_headers, &provider_request_body, &parts.headers, @@ -1386,46 +1578,24 @@ fn build_openai_image_provider_body_from_openai_chat_body( copy_openai_chat_image_option(body_json, &mut image_options, "input_fidelity"); copy_openai_chat_image_option(body_json, &mut image_options, "partial_images"); - let input = if images.is_empty() { - serde_json::json!([{ - "role": "user", - "content": prompt, - }]) - } else { - let mut content = vec![serde_json::json!({ - "type": "input_text", - "text": prompt, - })]; - content.extend(images); - serde_json::json!([{ - "role": "user", - "content": content, - }]) - }; - let mut body = serde_json::Map::new(); - if let Some(model) = body_json - .get("model") - .and_then(Value::as_str) - .map(str::trim) - .filter(|value| !value.is_empty()) - .or_else(|| { - let requested_model = requested_model.trim(); - (!requested_model.is_empty()).then_some(requested_model) - }) - { - body.insert("model".to_string(), Value::String(model.to_string())); + let requested_model = requested_model.trim(); + if requested_model.is_empty() { + return None; } - body.insert("input".to_string(), input); - let mut image_tool = image_options.clone(); - image_tool.insert( - "type".to_string(), - Value::String("image_generation".to_string()), - ); body.insert( - "tools".to_string(), - Value::Array(vec![Value::Object(image_tool)]), + "model".to_string(), + Value::String(requested_model.to_string()), ); + body.insert("prompt".to_string(), Value::String(prompt)); + body.extend(image_options.clone()); + if operation == "edit" { + let image_urls = openai_image_inputs_as_api_urls(&images); + if image_urls.len() != images.len() { + return None; + } + body.insert("images".to_string(), Value::Array(image_urls)); + } if upstream_is_stream { body.insert("stream".to_string(), Value::Bool(true)); } @@ -1451,6 +1621,14 @@ fn build_openai_image_provider_body_from_openai_chat_body( Some((Value::Object(body), Value::Object(summary))) } +fn openai_image_operation_from_summary(summary: &Value) -> Option { + match summary.get("operation")?.as_str()? { + "generate" => Some(OpenAiImageOperation::Generate), + "edit" => Some(OpenAiImageOperation::Edit), + _ => None, + } +} + fn build_chatgpt_web_image_provider_body_from_openai_chat_body( body_json: &Value, requested_model: &str, @@ -1603,6 +1781,20 @@ fn openai_image_inputs_as_urls(images: &[Value]) -> Vec { .collect() } +fn openai_image_inputs_as_api_urls(images: &[Value]) -> Vec { + images + .iter() + .filter_map(|image| { + image + .get("image_url") + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(|value| json!({ "image_url": value })) + }) + .collect() +} + fn chatgpt_web_ratio_for_size(size: &str) -> String { let Some((width, height)) = size.split_once('x') else { return "1:1".to_string(); @@ -2011,6 +2203,7 @@ mod tests { routing_policy: None, routing_trace_seed: None, routing_context: None, + model_directive_policy: Default::default(), } } @@ -2096,6 +2289,7 @@ mod tests { global_model_id: "global-model-1".to_string(), global_model_name: "gemini-2.5-pro".to_string(), selected_provider_model_name: "gemini-2.5-pro".to_string(), + supports_streaming: true, mapping_matched_model: None, }, transport: Arc::new(sample_gemini_cli_transport()), @@ -2138,6 +2332,238 @@ mod tests { eligible } + fn sample_openai_chat_eligible(provider_type: &str) -> EligibleLocalExecutionCandidate { + let mut transport = sample_gemini_cli_transport(); + transport.provider.name = provider_type.to_string(); + transport.provider.provider_type = provider_type.to_string(); + transport.endpoint.api_format = "openai:chat".to_string(); + transport.endpoint.api_family = Some("openai".to_string()); + transport.endpoint.endpoint_kind = Some("chat_completions".to_string()); + transport.endpoint.base_url = if provider_type == "grok" { + "https://grok.com".to_string() + } else { + "https://api.openai.test".to_string() + }; + transport.endpoint.custom_path = None; + transport.key.api_formats = Some(vec!["openai:chat".to_string()]); + transport.key.upstream_metadata = None; + if provider_type == "grok" { + transport.key.auth_type = "oauth".to_string(); + transport.key.decrypted_api_key.clear(); + transport.key.decrypted_auth_config = + Some(json!({ "sso_token": "test-session" }).to_string()); + } else { + transport.key.auth_type = "bearer".to_string(); + transport.key.decrypted_api_key = "test-api-key".to_string(); + transport.key.decrypted_auth_config = None; + } + + let mut eligible = sample_gemini_cli_eligible(); + eligible.candidate.provider_name = provider_type.to_string(); + eligible.candidate.provider_type = provider_type.to_string(); + eligible.candidate.endpoint_api_format = "openai:chat".to_string(); + eligible.candidate.global_model_name = "gpt-5.6-sol".to_string(); + eligible.candidate.selected_provider_model_name = "gpt-5.6-sol".to_string(); + eligible.transport = Arc::new(transport); + eligible.provider_api_format = "openai:chat".to_string(); + eligible + } + + fn sample_custom_directive_input() -> LocalOpenAiChatDecisionInput { + let mut input = sample_input(); + input.requested_model = "gpt-5.6-sol-high".to_string(); + input.model_directive_policy = + crate::system_features::ModelDirectivePolicySnapshot::from_config_values( + Some(&json!(true)), + Some(&json!({ + "reasoning_effort": { + "api_formats": { + "openai:chat": { + "suffixes": ["high"], + "mappings": { + "high": { + "reasoning_effort": "low", + "stream": false + } + } + } + } + } + })), + ); + input + } + + fn sample_alias_max_directive_input() -> LocalOpenAiChatDecisionInput { + let mut input = sample_input(); + input.requested_model = "deployment-alias-max".to_string(); + input.model_directive_policy = + crate::system_features::ModelDirectivePolicySnapshot::from_config_values( + Some(&json!(true)), + None, + ); + input + } + + #[tokio::test] + async fn alias_reasoning_directive_is_constrained_by_the_mapped_openai_model() { + let state = AppState::new().expect("state should build"); + let request = http::Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header(http::header::CONTENT_TYPE, "application/json") + .body(()) + .expect("request should build"); + let (parts, _) = request.into_parts(); + let body_json = json!({ + "model": "deployment-alias-max", + "messages": [{"role": "user", "content": "hello"}] + }); + + let mut supported = sample_openai_chat_eligible("custom"); + supported.candidate.selected_provider_model_name = "gpt-5.6-sol".to_string(); + let payload = resolve_local_openai_chat_candidate_payload_parts( + &state, + &parts, + "trace-alias-max-gpt-5.6-sol", + &body_json, + &sample_alias_max_directive_input(), + None, + &supported, + 0, + "candidate-0", + "openai_chat_sync", + "openai_chat_sync_success", + false, + ) + .await + .expect("candidate resolution should not fail") + .expect("GPT-5.6 candidate should build a payload"); + assert_eq!(payload.provider_request_body["reasoning_effort"], "max"); + + let mut unsupported = sample_openai_chat_eligible("custom"); + unsupported.candidate.selected_provider_model_name = "gpt-5.4".to_string(); + let payload = resolve_local_openai_chat_candidate_payload_parts( + &state, + &parts, + "trace-alias-max-gpt-5.4", + &body_json, + &sample_alias_max_directive_input(), + None, + &unsupported, + 0, + "candidate-0", + "openai_chat_sync", + "openai_chat_sync_success", + false, + ) + .await + .expect("candidate resolution should not fail"); + assert!(payload.is_none(), "GPT-5.4 must reject the max directive"); + } + + #[tokio::test] + async fn custom_policy_suffix_patch_is_applied_after_candidate_mapping() { + let state = AppState::new().expect("state should build"); + let request = http::Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header(http::header::CONTENT_TYPE, "application/json") + .body(()) + .expect("request should build"); + let (parts, _) = request.into_parts(); + let body_json = json!({ + "model": "deployment-alias-VendorFuture", + "messages": [{"role": "user", "content": "hello"}] + }); + let mut input = sample_input(); + input.requested_model = "deployment-alias-VendorFuture".to_string(); + input.model_directive_policy = + crate::system_features::ModelDirectivePolicySnapshot::from_config_values( + Some(&json!(true)), + Some(&json!({ + "reasoning_effort": { + "api_formats": { + "openai:chat": { + "suffixes": ["VendorFuture"], + "mappings": { + "VendorFuture": { + "reasoning_effort": "high" + } + } + } + } + } + })), + ); + let payload = resolve_local_openai_chat_candidate_payload_parts( + &state, + &parts, + "trace-custom-policy-suffix", + &body_json, + &input, + None, + &sample_openai_chat_eligible("custom"), + 0, + "candidate-0", + "openai_chat_sync", + "openai_chat_sync_success", + false, + ) + .await + .expect("candidate resolution should not fail") + .expect("custom directive candidate should build a payload"); + + assert_eq!(payload.provider_request_body["model"], "gpt-5.6-sol"); + assert_eq!(payload.provider_request_body["reasoning_effort"], "high"); + } + + #[tokio::test] + async fn same_format_and_grok_chat_apply_the_same_custom_directive_finalization() { + let state = AppState::new().expect("state should build"); + let request = http::Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header(http::header::CONTENT_TYPE, "application/json") + .body(()) + .expect("request should build"); + let (parts, _) = request.into_parts(); + let body_json = json!({ + "model": "gpt-5.6-sol-high", + "messages": [{"role": "user", "content": "hello"}], + "stream": true + }); + + for provider_type in ["custom", "grok"] { + let payload = resolve_local_openai_chat_candidate_payload_parts( + &state, + &parts, + &format!("trace-directive-{provider_type}"), + &body_json, + &sample_custom_directive_input(), + None, + &sample_openai_chat_eligible(provider_type), + 0, + "candidate-0", + OPENAI_CHAT_STREAM_PLAN_KIND, + "openai_chat_stream_success", + true, + ) + .await + .expect("candidate resolution should not fail") + .expect("same-format candidate should build a payload"); + + assert_eq!( + payload.provider_request_body["reasoning_effort"], "low", + "custom mapping must be authoritative for {provider_type}" + ); + assert_eq!( + payload.provider_request_body["stream"], true, + "stream policy must be re-applied after mapping for {provider_type}" + ); + } + } + #[tokio::test] async fn openai_chat_to_gemini_cli_wraps_cross_format_body_in_v1internal_envelope() { let state = AppState::new().expect("state should build"); @@ -2333,7 +2759,7 @@ mod tests { } #[test] - fn openai_chat_image_bridge_body_injects_image_generation_tool() { + fn openai_chat_image_bridge_builds_images_api_body() { let body_json = json!({ "model": "gpt-image-2", "messages": [ @@ -2347,13 +2773,22 @@ mod tests { build_openai_image_provider_body_from_openai_chat_body(&body_json, "gpt-image-2", true) .expect("chat image body should convert"); - assert_eq!(provider_body["tools"][0]["type"], "image_generation"); - assert_eq!(provider_body["tools"][0]["size"], "1024x1024"); - assert_eq!(provider_body["tools"][0]["output_format"], "png"); assert_eq!(provider_body["model"], "gpt-image-2"); + assert_eq!(provider_body["prompt"], "Draw a glass city"); + assert_eq!(provider_body["size"], "1024x1024"); + assert_eq!(provider_body["output_format"], "png"); assert_eq!(provider_body["stream"], true); - assert_eq!(provider_body["input"][0]["content"], "Draw a glass city"); + assert!(provider_body.get("tools").is_none()); + assert!(provider_body.get("input").is_none()); assert_eq!(summary["operation"], "generate"); assert_eq!(summary["output_format"], "png"); + + let (sync_provider_body, _) = build_openai_image_provider_body_from_openai_chat_body( + &body_json, + "gpt-image-2", + false, + ) + .expect("chat image body should convert for a sync upstream"); + assert!(sync_provider_body.get("stream").is_none()); } } diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/support.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/support.rs index d3cc2c337..f57237ec2 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/support.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/support.rs @@ -293,6 +293,7 @@ pub(crate) async fn build_lazy_local_openai_chat_candidate_attempt_source<'a>( ); build_lazy_requested_model_execution_candidate_attempt_source_with_serving( planner_state, + &input.model_directive_policy, trace_id, "openai:chat", &input.requested_model, diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/candidates.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/candidates.rs index f5a0be53c..e0db3b2b1 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/candidates.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/candidates.rs @@ -21,6 +21,7 @@ pub(crate) async fn list_local_openai_chat_candidates( > { let outcome = preselect_local_execution_candidates_with_serving( PlannerAppState::new(state), + &input.model_directive_policy, "openai:chat", &input.requested_model, require_streaming, diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/resolve.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/resolve.rs index 42a08bfaf..24b123087 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/resolve.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/resolve.rs @@ -65,7 +65,9 @@ pub(crate) async fn resolve_local_openai_chat_decision_input( state, auth_context.clone(), Some(requested_model.as_str()), + decision.auth_endpoint_signature.as_deref(), None, + &decision.model_directive_policy, ) .await { 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 f4849f1b1..1a7d78675 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 @@ -111,7 +111,7 @@ pub(crate) async fn build_local_openai_chat_stream_attempt_source<'a>( input, candidates, prefetched_attempts: VecDeque::new(), - request_preparation: LocalOpenAiChatRequestPreparation::default(), + request_preparation: LocalOpenAiChatRequestPreparation, }, candidate_count, ))) diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/plan_builders/stream.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/plan_builders/stream.rs index b60a245bb..f41c5c1f5 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/plan_builders/stream.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/plan_builders/stream.rs @@ -129,7 +129,7 @@ pub(crate) fn build_openai_chat_stream_plan_from_decision( headers: std::mem::take(&mut provider_request_headers), content_type, body: RequestBody::from_json(provider_request_body_value), - stream: true, + stream: effective_upstream_is_stream, }, ); @@ -229,7 +229,7 @@ pub(crate) fn build_openai_responses_stream_plan_from_decision( headers: std::mem::take(&mut provider_request_headers), content_type, body: RequestBody::from_json(provider_request_body_value), - stream: true, + stream: effective_upstream_is_stream, }, ); @@ -291,6 +291,7 @@ mod tests { request_id: Some("req_123".to_string()), candidate_id: Some("cand_123".to_string()), provider_name: Some("Codex".to_string()), + provider_type: Some("codex".to_string()), provider_id: Some("prov_123".to_string()), endpoint_id: Some("ep_123".to_string()), key_id: Some("key_123".to_string()), @@ -385,6 +386,46 @@ mod tests { ); } + #[test] + fn build_compact_stream_plan_preserves_non_stream_upstream_mode() { + let parts = http::Request::builder() + .uri("http://localhost/v1/responses/compact") + .body(()) + .expect("request should build") + .into_parts() + .0; + let mut payload = sample_responses_payload(); + payload.decision_kind = Some("openai_responses_compact_stream".to_string()); + payload.upstream_url = Some("https://example.com/v1/responses/compact".to_string()); + payload.provider_api_format = Some("openai:responses:compact".to_string()); + payload.client_api_format = Some("openai:responses:compact".to_string()); + payload.upstream_is_stream = false; + payload.provider_request_body = Some(json!({ + "model": "gpt-5.6-sol", + "input": [], + "instructions": "You are Codex.", + "tools": [], + "parallel_tool_calls": true, + "reasoning": {"effort": "high"}, + "prompt_cache_key": "cache-key", + "text": {"verbosity": "low"} + })); + + let built = + build_openai_responses_stream_plan_from_decision(&parts, &json!({}), payload, true) + .expect("plan build should succeed") + .expect("plan should be produced"); + + assert!(!built.plan.stream); + assert!(built + .plan + .body + .json_body + .as_ref() + .is_some_and(|body| body.get("stream").is_none())); + assert!(built.plan.headers.get("accept").is_none()); + } + #[test] fn build_openai_chat_stream_plan_fallback_preserves_complete_same_format_headers() { let parts = http::Request::builder() @@ -404,6 +445,7 @@ mod tests { request_id: Some("req_stream_456".to_string()), candidate_id: Some("cand_stream_456".to_string()), provider_name: Some("OpenAI".to_string()), + provider_type: Some("openai".to_string()), provider_id: Some("prov_stream_456".to_string()), endpoint_id: Some("ep_stream_456".to_string()), key_id: Some("key_stream_456".to_string()), @@ -462,7 +504,7 @@ mod tests { } #[test] - fn build_openai_chat_stream_plan_keeps_downstream_stream_for_force_non_stream_upstream() { + fn build_openai_chat_stream_plan_preserves_force_non_stream_upstream_mode() { fn force_non_stream_payload(provider_request_body: Option) -> AiExecutionDecision { AiExecutionDecision { action: "stream".to_string(), @@ -472,6 +514,7 @@ mod tests { request_id: Some("req_force_non_stream".to_string()), candidate_id: Some("cand_force_non_stream".to_string()), provider_name: Some("OpenAI".to_string()), + provider_type: Some("openai".to_string()), provider_id: Some("prov_force_non_stream".to_string()), endpoint_id: Some("ep_force_non_stream".to_string()), key_id: Some("key_force_non_stream".to_string()), @@ -523,7 +566,7 @@ mod tests { .expect("plan build should succeed") .expect("plan should be produced"); - assert!(built.plan.stream); + assert!(!built.plan.stream); assert_eq!( built .plan @@ -548,7 +591,7 @@ mod tests { .expect("fallback plan build should succeed") .expect("fallback plan should be produced"); - assert!(built.plan.stream); + assert!(!built.plan.stream); assert_eq!( built .plan @@ -579,6 +622,7 @@ mod tests { request_id: Some("req_stream_789".to_string()), candidate_id: Some("cand_stream_789".to_string()), provider_name: Some("Claude".to_string()), + provider_type: Some("anthropic".to_string()), provider_id: Some("prov_stream_789".to_string()), endpoint_id: Some("ep_stream_789".to_string()), key_id: Some("key_stream_789".to_string()), diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/plan_builders/sync.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/plan_builders/sync.rs index ae44d683f..f1047ad18 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/plan_builders/sync.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/plan_builders/sync.rs @@ -257,6 +257,7 @@ mod tests { request_id: Some("req_123".to_string()), candidate_id: Some("cand_123".to_string()), provider_name: Some("Codex".to_string()), + provider_type: Some("codex".to_string()), provider_id: Some("prov_123".to_string()), endpoint_id: Some("ep_123".to_string()), key_id: Some("key_123".to_string()), @@ -369,6 +370,7 @@ mod tests { request_id: Some("req_456".to_string()), candidate_id: Some("cand_456".to_string()), provider_name: Some("OpenAI".to_string()), + provider_type: Some("openai".to_string()), provider_id: Some("prov_456".to_string()), endpoint_id: Some("ep_456".to_string()), key_id: Some("key_456".to_string()), @@ -440,6 +442,7 @@ mod tests { request_id: Some("req_789".to_string()), candidate_id: Some("cand_789".to_string()), provider_name: Some("Claude".to_string()), + provider_type: Some("anthropic".to_string()), provider_id: Some("prov_789".to_string()), endpoint_id: Some("ep_789".to_string()), key_id: Some("key_789".to_string()), diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/payload.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/payload.rs index 75785d27b..5b318b047 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/payload.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/payload.rs @@ -9,7 +9,7 @@ use crate::ai_serving::planner::report_context::{ }; use crate::ai_serving::planner::spec_metadata::local_openai_responses_spec_metadata; use crate::ai_serving::planner::{ - build_ai_execution_decision_response, resolve_transport_request_gzip_policy, + build_ai_execution_decision_response, resolve_transport_request_encoding_policy, AiExecutionDecisionResponseParts, }; use crate::ai_serving::transport::{ @@ -204,7 +204,7 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand image_request_summary: _, request_redacted: _, } = resolved; - let request_gzip = resolve_transport_request_gzip_policy(&transport); + let request_encoding = resolve_transport_request_encoding_policy(&transport); let mut decision = build_ai_execution_decision_response(AiExecutionDecisionResponseParts { decision_is_stream: spec_metadata.require_streaming, @@ -214,6 +214,7 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand request_id: trace_id.to_string(), candidate_id: candidate_id.clone(), provider_name: transport.provider.name.clone(), + provider_type: transport.provider.provider_type.clone(), provider_id: candidate.provider_id.clone(), endpoint_id: candidate.endpoint_id.clone(), key_id: candidate.key_id.clone(), @@ -231,8 +232,8 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand provider_request_body: Some(provider_request_body), provider_request_body_base64: None, content_type: Some("application/json".to_string()), - content_encoding: None, - request_gzip, + content_encoding: request_encoding.content_encoding, + request_gzip: request_encoding.request_gzip, proxy, transport_profile, timeouts, @@ -241,6 +242,10 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand report_context: Some(report_context), auth_context: input.auth_context.clone(), }); - apply_provider_request_routing_policy_to_decision(input, &mut decision)?; + apply_provider_request_routing_policy_to_decision( + input, + &mut decision, + Some(transport.as_ref()), + )?; Ok(Some(decision)) } diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/request.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/request.rs index e6edbcf61..120c872a7 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/request.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/request.rs @@ -27,11 +27,12 @@ use crate::ai_serving::planner::redaction::{ }; use crate::ai_serving::planner::spec_metadata::local_openai_responses_spec_metadata; use crate::ai_serving::planner::standard::{ - apply_codex_openai_responses_special_body_edits, apply_codex_openai_responses_special_headers, - apply_deepseek_tool_call_thinking_compat, build_cross_format_openai_responses_request_body, - build_cross_format_openai_responses_upstream_url, build_local_openai_responses_request_body, - build_local_openai_responses_upstream_url, request_body_build_failure_extra_data, - request_conversion_failure_extra_data, + apply_codex_openai_special_headers, apply_deepseek_tool_call_thinking_compat, + build_cross_format_openai_responses_request_body_with_codex_model_capabilities, + build_cross_format_openai_responses_upstream_url, + build_local_openai_responses_request_body_with_codex_model_capabilities, + build_local_openai_responses_upstream_url, codex_model_capabilities_for_transport, + request_body_build_failure_extra_data, request_conversion_failure_extra_data, }; use crate::ai_serving::transport::antigravity::is_antigravity_provider_transport; use crate::ai_serving::transport::auth::{ @@ -58,7 +59,10 @@ use crate::ai_serving::transport::{ use crate::ai_serving::{ ai_local_execution_contract_for_formats, request_conversion_direct_auth, request_conversion_kind, CandidateFailureDiagnostic, GatewayProviderTransportSnapshot, - LocalResolvedOAuthRequestAuth, PlannerAppState, + LocalResolvedOAuthRequestAuth, OpenAiImageOperation, PlannerAppState, +}; +use crate::ai_serving::{ + project_codex_openai_image_api_request_body, project_openai_image_api_request_body, }; use crate::ai_serving::{ConversionMode, ExecutionStrategy}; use crate::{AppState, GatewayError}; @@ -282,13 +286,26 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts( let auth_header = prepared_candidate.auth_header; let auth_value = prepared_candidate.auth_value; let mapped_model = prepared_candidate.mapped_model; - let enable_model_directives = - crate::system_features::reasoning_model_directive_enabled_for_api_format_and_model( - state, - provider_api_format, - Some(&input.requested_model), - ) - .await; + let model_directive_resolution = input + .model_directive_policy + .resolve_reasoning(provider_api_format, Some(&input.requested_model)); + let model_directive_mapping = + match model_directive_resolution.mapping_patch_for_mapped_model(&mapped_model) { + Ok(mapping) => mapping, + Err(skip_reason) => { + mark_skipped_local_openai_responses_candidate( + state, + input, + trace_id, + candidate, + candidate_index, + candidate_id, + skip_reason, + ) + .await; + return Ok(None); + } + }; let redaction = resolve_provider_chat_pii_redaction( state, parts, @@ -311,9 +328,19 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts( let force_body_stream_field = endpoint_config_forces_body_stream_field(transport.endpoint.config.as_ref()); let effective_headers = input.effective_headers(&parts.headers); + let source_model = body_json + .get("model") + .and_then(Value::as_str) + .unwrap_or(input.requested_model.as_str()); + let codex_model_capabilities = codex_model_capabilities_for_transport( + &transport, + provider_api_format, + mapped_model.as_str(), + source_model, + ); let Some(mut base_provider_request_body) = (if is_grok && is_grok_text_provider_api_format(provider_api_format) { - build_local_openai_responses_request_body( + build_local_openai_responses_request_body_with_codex_model_capabilities( body_json, &mapped_model, upstream_is_stream, @@ -321,12 +348,12 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts( transport.provider.provider_type.as_str(), spec_metadata.api_format, transport.endpoint.body_rules.as_ref(), - Some(input.auth_context.api_key_id.as_str()), effective_headers, - enable_model_directives, + codex_model_capabilities.as_ref(), + false, ) } else if needs_bidirectional_conversion { - build_cross_format_openai_responses_request_body( + build_cross_format_openai_responses_request_body_with_codex_model_capabilities( body_json, &mapped_model, spec_metadata.api_format, @@ -339,12 +366,12 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts( } else { transport.endpoint.body_rules.as_ref() }, - Some(input.auth_context.api_key_id.as_str()), effective_headers, - enable_model_directives, + codex_model_capabilities.as_ref(), + false, ) } else { - build_local_openai_responses_request_body( + build_local_openai_responses_request_body_with_codex_model_capabilities( body_json, &mapped_model, upstream_is_stream, @@ -356,9 +383,9 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts( } else { transport.endpoint.body_rules.as_ref() }, - Some(input.auth_context.api_key_id.as_str()), effective_headers, - enable_model_directives, + codex_model_capabilities.as_ref(), + false, ) }) else { @@ -383,17 +410,10 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts( .await; return Ok(None); }; - if let Some(mapping) = - crate::system_features::reasoning_model_directive_mapping_for_api_format_and_model( - state, - provider_api_format, - Some(&input.requested_model), - ) - .await - { + if let Some(mapping) = model_directive_mapping.as_ref() { crate::ai_serving::apply_model_directive_mapping_patch( &mut base_provider_request_body, - &mapping, + mapping, ); // Directive mapping is a deep-merge patch and may overwrite/add `stream`; // re-enforce stream-field policy afterward. @@ -411,6 +431,46 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts( provider_api_format, Some(body_json), ); + if crate::ai_serving::finalize_openai_provider_request_with_codex_model_capabilities( + &mut base_provider_request_body, + crate::ai_serving::OpenAiProviderRequestFinalization { + source_api_format: spec_metadata.api_format, + provider_api_format, + provider_type: transport.provider.provider_type.as_str(), + provider_model: mapped_model.as_str(), + source_model, + body_rules: transport.endpoint.body_rules.as_ref(), + upstream_is_stream, + require_body_stream_field: request_requires_body_stream_field( + body_json, + force_body_stream_field, + ), + }, + codex_model_capabilities.as_ref(), + ) + .is_err() + { + mark_skipped_local_openai_responses_candidate_with_extra_data( + state, + input, + trace_id, + candidate, + candidate_index, + candidate_id, + "provider_request_body_build_failed", + request_conversion_failure_extra_data( + body_json, + spec_metadata.api_format, + provider_api_format, + Some(mapped_model.as_str()), + Some(parts.uri.path()), + upstream_is_stream, + "openai_responses_request_conversion", + ), + ) + .await; + return Ok(None); + } let provider_request_body = base_provider_request_body; if let Some(kiro_auth) = kiro_auth.as_ref() { @@ -612,7 +672,11 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts( }; let mut provider_request_headers = resolved_headers.headers; if !is_grok { - apply_codex_openai_responses_special_headers( + apply_local_auth_config_header_overrides( + &mut provider_request_headers, + transport.key.decrypted_auth_config.as_deref(), + ); + apply_codex_openai_special_headers( &mut provider_request_headers, &provider_request_body, effective_headers, @@ -621,9 +685,13 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts( Some(trace_id), transport.key.decrypted_auth_config.as_deref(), ); - apply_local_auth_config_header_overrides( + crate::ai_serving::apply_codex_openai_responses_lite_header_with_capabilities( &mut provider_request_headers, - transport.key.decrypted_auth_config.as_deref(), + transport.provider.provider_type.as_str(), + provider_api_format, + mapped_model.as_str(), + source_model, + codex_model_capabilities.as_ref(), ); } request_identity_response_encoding_when_redacted( @@ -788,7 +856,11 @@ async fn build_antigravity_openai_responses_payload_parts( } }; let mut provider_request_headers = resolved.headers.headers; - apply_codex_openai_responses_special_headers( + apply_local_auth_config_header_overrides( + &mut provider_request_headers, + resolved.transport.key.decrypted_auth_config.as_deref(), + ); + apply_codex_openai_special_headers( &mut provider_request_headers, &resolved.body, effective_headers, @@ -797,10 +869,6 @@ async fn build_antigravity_openai_responses_payload_parts( Some(trace_id), resolved.transport.key.decrypted_auth_config.as_deref(), ); - apply_local_auth_config_header_overrides( - &mut provider_request_headers, - resolved.transport.key.decrypted_auth_config.as_deref(), - ); provider_request_headers.insert("accept".to_string(), "text/event-stream".to_string()); request_identity_response_encoding_when_redacted( &mut provider_request_headers, @@ -941,7 +1009,11 @@ async fn build_gemini_cli_openai_responses_payload_parts( } }; let mut provider_request_headers = resolved.headers.headers; - apply_codex_openai_responses_special_headers( + apply_local_auth_config_header_overrides( + &mut provider_request_headers, + resolved.transport.key.decrypted_auth_config.as_deref(), + ); + apply_codex_openai_special_headers( &mut provider_request_headers, &resolved.body, effective_headers, @@ -950,10 +1022,6 @@ async fn build_gemini_cli_openai_responses_payload_parts( Some(trace_id), resolved.transport.key.decrypted_auth_config.as_deref(), ); - apply_local_auth_config_header_overrides( - &mut provider_request_headers, - resolved.transport.key.decrypted_auth_config.as_deref(), - ); request_identity_response_encoding_when_redacted( &mut provider_request_headers, request_redacted, @@ -1182,11 +1250,16 @@ async fn resolve_openai_responses_to_openai_image_payload_parts( .provider_type .trim() .eq_ignore_ascii_case("chatgpt_web"); + let is_codex = transport + .provider + .provider_type + .trim() + .eq_ignore_ascii_case("codex"); let upstream_is_stream = resolve_upstream_is_stream_for_provider( transport.endpoint.config.as_ref(), transport.provider.provider_type.as_str(), provider_api_format, - spec_metadata.require_streaming, + spec_metadata.require_streaming && candidate.supports_streaming, false, ); let Some((mut provider_request_body, image_request_summary)) = (if is_chatgpt_web { @@ -1197,7 +1270,7 @@ async fn resolve_openai_responses_to_openai_image_payload_parts( } else { build_openai_image_provider_body_from_openai_responses_body( body_json, - &input.requested_model, + &prepared_candidate.mapped_model, upstream_is_stream, ) }) else { @@ -1218,25 +1291,31 @@ async fn resolve_openai_responses_to_openai_image_payload_parts( .await; return None; }; - + let operation = openai_image_operation_from_summary(&image_request_summary)?; if !is_chatgpt_web { - apply_codex_openai_responses_special_body_edits( - &mut provider_request_body, - transport.provider.provider_type.as_str(), - provider_api_format, - transport.endpoint.body_rules.as_ref(), - Some(candidate.key_id.as_str()), - ); + provider_request_body = project_openai_image_api_request_body( + &provider_request_body, + &prepared_candidate.mapped_model, + operation, + crate::image_capabilities::openai_image_provider_max_generation_count_for_model( + transport.provider.provider_type.as_str(), + Some(prepared_candidate.mapped_model.as_str()), + ), + )?; + } + if is_codex { + provider_request_body = + project_codex_openai_image_api_request_body(&provider_request_body, operation)?; } let upstream_url = if is_chatgpt_web { chatgpt_web_image_internal_url(&transport.endpoint.base_url) } else { - build_openai_image_upstream_url( - transport, - Some("/v1/images/generations"), - parts.uri.query(), - ) + let request_path = match operation { + OpenAiImageOperation::Generate => "/v1/images/generations", + OpenAiImageOperation::Edit => "/v1/images/edits", + }; + build_openai_image_upstream_url(transport, Some(request_path), parts.uri.query()) }; let Some(mut provider_request_headers) = build_openai_image_headers(ProviderOpenAiImageHeadersInput { @@ -1244,7 +1323,13 @@ async fn resolve_openai_responses_to_openai_image_payload_parts( headers: &parts.headers, auth_header: &prepared_candidate.auth_header, auth_value: &prepared_candidate.auth_value, - accept: "text/event-stream", + accept: if is_codex { + None + } else if upstream_is_stream { + Some("text/event-stream") + } else { + Some("application/json") + }, header_rules: transport.endpoint.header_rules.as_ref(), provider_request_body: &provider_request_body, original_request_body: body_json, @@ -1270,7 +1355,11 @@ async fn resolve_openai_responses_to_openai_image_payload_parts( if is_chatgpt_web { provider_request_headers.insert("x-aether-chatgpt-web-image".to_string(), "1".to_string()); } else { - apply_codex_openai_responses_special_headers( + apply_local_auth_config_header_overrides( + &mut provider_request_headers, + transport.key.decrypted_auth_config.as_deref(), + ); + apply_codex_openai_special_headers( &mut provider_request_headers, &provider_request_body, &parts.headers, @@ -1279,10 +1368,6 @@ async fn resolve_openai_responses_to_openai_image_payload_parts( Some(trace_id), transport.key.decrypted_auth_config.as_deref(), ); - apply_local_auth_config_header_overrides( - &mut provider_request_headers, - transport.key.decrypted_auth_config.as_deref(), - ); } let (execution_strategy, conversion_mode) = @@ -1314,59 +1399,68 @@ fn build_openai_image_provider_body_from_openai_responses_body( upstream_is_stream: bool, ) -> Option<(Value, Value)> { let object = body_json.as_object()?; - let input = object.get("input")?.clone(); let tool = openai_responses_image_generation_tool(object); + let (prompt, images) = collect_openai_responses_image_prompt_and_images(object.get("input"))?; + let operation = if images.is_empty() { + OpenAiImageOperation::Generate + } else { + OpenAiImageOperation::Edit + }; + if let Some(action) = tool + .as_ref() + .and_then(|tool| tool.get("action")) + .and_then(Value::as_str) + { + let expected = operation.as_str(); + if !action.trim().eq_ignore_ascii_case(expected) { + return None; + } + } let mut body = serde_json::Map::new(); - body.insert("input".to_string(), input); - if let Some(model) = object - .get("model") - .and_then(Value::as_str) - .map(str::trim) - .filter(|value| !value.is_empty()) - .or_else(|| { - let requested_model = requested_model.trim(); - (!requested_model.is_empty()).then_some(requested_model) - }) - { - body.insert("model".to_string(), Value::String(model.to_string())); + let requested_model = requested_model.trim(); + if requested_model.is_empty() { + return None; } + body.insert( + "model".to_string(), + Value::String(requested_model.to_string()), + ); + body.insert("prompt".to_string(), Value::String(prompt)); for key in [ + "background", + "quality", + "size", + "output_format", + "output_compression", + "moderation", + "input_fidelity", + "partial_images", + "n", "user", - "metadata", - "include", - "parallel_tool_calls", - "store", ] { - if let Some(value) = object.get(key) { + if let Some(value) = tool + .as_ref() + .and_then(|tool| tool.get(key)) + .or_else(|| object.get(key)) + { body.insert(key.to_string(), value.clone()); } } + if operation == OpenAiImageOperation::Edit { + let image_urls = openai_image_inputs_as_api_urls(&images); + if image_urls.len() != images.len() { + return None; + } + body.insert("images".to_string(), Value::Array(image_urls)); + } if upstream_is_stream { body.insert("stream".to_string(), Value::Bool(true)); - } else if let Some(value) = object.get("stream") { - body.insert("stream".to_string(), value.clone()); } - let image_tool = tool.clone().unwrap_or_else(|| { - let mut tool = serde_json::Map::new(); - tool.insert( - "type".to_string(), - Value::String("image_generation".to_string()), - ); - tool - }); - body.insert( - "tools".to_string(), - Value::Array(vec![Value::Object(image_tool)]), - ); - let mut summary = serde_json::Map::new(); summary.insert( "operation".to_string(), - tool.as_ref() - .and_then(|tool| tool.get("action")) - .cloned() - .unwrap_or_else(|| json!("generate")), + Value::String(operation.as_str().to_string()), ); for key in ["output_format", "partial_images", "size", "quality"] { let tool_value = tool.as_ref().and_then(|tool| tool.get(key)); @@ -1378,6 +1472,14 @@ fn build_openai_image_provider_body_from_openai_responses_body( Some((Value::Object(body), Value::Object(summary))) } +fn openai_image_operation_from_summary(summary: &Value) -> Option { + match summary.get("operation")?.as_str()? { + "generate" => Some(OpenAiImageOperation::Generate), + "edit" => Some(OpenAiImageOperation::Edit), + _ => None, + } +} + fn openai_responses_image_generation_tool( object: &serde_json::Map, ) -> Option> { @@ -1572,6 +1674,20 @@ fn openai_image_inputs_as_urls(images: &[Value]) -> Vec { .collect() } +fn openai_image_inputs_as_api_urls(images: &[Value]) -> Vec { + images + .iter() + .filter_map(|image| { + image + .get("image_url") + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(|value| json!({ "image_url": value })) + }) + .collect() +} + fn chatgpt_web_ratio_for_size(size: &str) -> String { let Some((width, height)) = size.split_once('x') else { return "1:1".to_string(); @@ -1769,7 +1885,7 @@ mod tests { use super::*; #[test] - fn openai_responses_image_bridge_body_preserves_image_generation_tool() { + fn openai_responses_image_bridge_builds_images_api_body() { let body_json = json!({ "model": "gpt-image-2", "input": "Draw a glass city", @@ -1792,14 +1908,74 @@ mod tests { ) .expect("responses image body should convert"); - assert_eq!(provider_body["tools"][0]["type"], "image_generation"); - assert_eq!(provider_body["tools"][0]["size"], "1024x1024"); - assert_eq!(provider_body["tools"][0]["output_format"], "png"); assert_eq!(provider_body["model"], "gpt-image-2"); - assert_eq!(provider_body["input"], "Draw a glass city"); + assert_eq!(provider_body["prompt"], "Draw a glass city"); + assert_eq!(provider_body["size"], "1024x1024"); + assert_eq!(provider_body["output_format"], "png"); assert_eq!(provider_body["stream"], true); + assert!(provider_body.get("tools").is_none()); + assert!(provider_body.get("input").is_none()); assert_eq!(summary["operation"], "generate"); assert_eq!(summary["output_format"], "png"); + + let (sync_provider_body, _) = build_openai_image_provider_body_from_openai_responses_body( + &body_json, + "gpt-image-2", + false, + ) + .expect("responses image body should convert for a sync upstream"); + assert!(sync_provider_body.get("stream").is_none()); + } + + #[test] + fn responses_image_bridge_uses_the_shared_mapped_model_projection() { + let body_json = json!({ + "model": "image-alias", + "input": "Draw a glass city", + "tools": [{ + "type": "image_generation", + "quality": "high", + "n": 2 + }], + "tool_choice": {"type": "image_generation"} + }); + let (body, _) = build_openai_image_provider_body_from_openai_responses_body( + &body_json, "dall-e-3", false, + ) + .expect("Responses image body should convert before provider projection"); + + assert!(project_openai_image_api_request_body( + &body, + "dall-e-3", + OpenAiImageOperation::Generate, + 1, + ) + .is_none()); + let single = json!({ + "model": "dall-e-3", + "prompt": "Draw a glass city", + "quality": "high", + "n": 1 + }); + let projected = project_openai_image_api_request_body( + &single, + "dall-e-3", + OpenAiImageOperation::Generate, + 1, + ) + .expect("DALL-E 3 single image request should project"); + assert_eq!(projected["quality"], "hd"); + + let codex_overflow = json!({ + "model": "gpt-image-2", + "prompt": "Draw a glass city", + "n": 11 + }); + assert!(project_codex_openai_image_api_request_body( + &codex_overflow, + OpenAiImageOperation::Generate + ) + .is_none()); } #[test] diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/support.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/support.rs index 537139ea3..c44d3782c 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/support.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/support.rs @@ -91,7 +91,9 @@ pub(crate) async fn resolve_local_openai_responses_decision_input( state, auth_context.clone(), Some(requested_model.as_str()), + decision.auth_endpoint_signature.as_deref(), None, + &decision.model_directive_policy, ) .await { @@ -171,6 +173,7 @@ pub(crate) async fn materialize_local_openai_responses_candidate_attempts( ); let preselection = preselect_local_execution_candidates_with_serving( planner_state, + &input.model_directive_policy, spec_metadata.api_format, &input.requested_model, spec_metadata.require_streaming, @@ -280,6 +283,7 @@ pub(crate) async fn build_local_openai_responses_candidate_attempt_source<'a>( Ok( build_lazy_requested_model_execution_candidate_attempt_source_with_serving( planner_state, + &input.model_directive_policy, trace_id, spec_metadata.api_format, &input.requested_model, @@ -365,6 +369,7 @@ pub(crate) async fn build_local_openai_responses_image_candidate_attempt_source< ); let preselection = preselect_local_execution_candidates_for_api_formats_with_serving( planner_state, + &input.model_directive_policy, spec_metadata.api_format, &input.requested_model, false, 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 6c56b72eb..99078f2e4 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 @@ -146,6 +146,7 @@ pub(crate) fn build_standard_stream_plan_from_decision( &provider_request_headers, &provider_request_body_value, )?; + let stream = payload.upstream_is_stream; let plan = build_ai_execution_plan_from_decision( &mut payload, AiExecutionPlanFromDecisionParts { @@ -155,7 +156,7 @@ pub(crate) fn build_standard_stream_plan_from_decision( headers: std::mem::take(&mut provider_request_headers), content_type, body: RequestBody::from_json(provider_request_body_value), - stream: true, + stream, }, ); diff --git a/apps/aether-gateway/src/ai_serving/planner/state/candidate_runtime.rs b/apps/aether-gateway/src/ai_serving/planner/state/candidate_runtime.rs index 7d635a8aa..a30961369 100644 --- a/apps/aether-gateway/src/ai_serving/planner/state/candidate_runtime.rs +++ b/apps/aether-gateway/src/ai_serving/planner/state/candidate_runtime.rs @@ -10,16 +10,15 @@ impl<'a> PlannerAppState<'a> { api_key_id: &str, requested_model: Option<&str>, explicit_required_capabilities: Option<&Value>, + model_directive_base_model: Option<&str>, ) -> Option { - let enable_model_directives = - crate::system_features::reasoning_model_directive_enabled(self.app()).await; crate::request_candidate_runtime::resolve_request_candidate_required_capabilities( self.app(), user_id, api_key_id, requested_model, explicit_required_capabilities, - enable_model_directives, + model_directive_base_model, ) .await } diff --git a/apps/aether-gateway/src/ai_serving/planner/state/scheduler.rs b/apps/aether-gateway/src/ai_serving/planner/state/scheduler.rs index 2ecf04662..22571c73b 100644 --- a/apps/aether-gateway/src/ai_serving/planner/state/scheduler.rs +++ b/apps/aether-gateway/src/ai_serving/planner/state/scheduler.rs @@ -20,14 +20,8 @@ impl<'a> PlannerAppState<'a> { auth_snapshot: Option<&GatewayAuthApiKeySnapshot>, client_session_affinity: Option<&ClientSessionAffinity>, now_unix_secs: u64, + enable_model_directives: bool, ) -> Result, GatewayError> { - let enable_model_directives = - crate::system_features::reasoning_model_directive_enabled_for_api_format_and_model( - self.app(), - api_format, - Some(global_model_name), - ) - .await; crate::scheduler::candidate::list_selectable_candidates( self.app().data.as_ref(), self.app(), @@ -52,6 +46,7 @@ impl<'a> PlannerAppState<'a> { auth_snapshot: Option<&GatewayAuthApiKeySnapshot>, client_session_affinity: Option<&ClientSessionAffinity>, now_unix_secs: u64, + enable_model_directives: bool, ) -> Result< ( Vec, @@ -63,14 +58,6 @@ impl<'a> PlannerAppState<'a> { let wait_interval = Duration::from_millis(API_KEY_CONCURRENCY_WAIT_POLL_INTERVAL_MS.max(1)); let wait_deadline = Instant::now() + wait_timeout; let mut attempt_now_unix_secs = now_unix_secs; - let enable_model_directives = - crate::system_features::reasoning_model_directive_enabled_for_api_format_and_model( - self.app(), - api_format, - Some(global_model_name), - ) - .await; - loop { let result = crate::scheduler::candidate::list_selectable_candidates_with_skip_reasons( self.app().data.as_ref(), diff --git a/apps/aether-gateway/src/ai_serving/pure/mod.rs b/apps/aether-gateway/src/ai_serving/pure/mod.rs index 3df3aa6e0..c1a22da71 100644 --- a/apps/aether-gateway/src/ai_serving/pure/mod.rs +++ b/apps/aether-gateway/src/ai_serving/pure/mod.rs @@ -3,10 +3,14 @@ pub(crate) use aether_ai_formats::api::{ aggregate_openai_chat_stream_sync_response, aggregate_openai_responses_stream_sync_response, aggregate_standard_chat_stream_sync_response, aggregate_standard_cli_stream_sync_response, api_format_alias_matches, api_format_storage_aliases, - apply_codex_openai_responses_chat_body_edits, apply_codex_openai_responses_special_body_edits, - apply_codex_openai_responses_special_headers, apply_model_directive_mapping_patch, + apply_codex_openai_compact_terminal_headers, apply_codex_openai_responses_chat_body_edits, + apply_codex_openai_responses_lite_header_with_capabilities, + apply_codex_openai_responses_special_body_edits, + apply_codex_openai_responses_special_body_edits_with_source_model_and_capabilities, + apply_codex_openai_special_headers, apply_model_directive_mapping_patch, apply_model_directive_overrides_from_model, apply_model_directive_overrides_from_request, apply_openai_responses_compact_special_body_edits, build_chatgpt_web_image_request_body, + build_codex_model_catalog_metadata, build_codex_openai_image_api_provider_request_body, build_core_error_body_for_client_format, build_cross_format_openai_chat_request_body, build_cross_format_openai_chat_request_body_with_model_directives, build_cross_format_openai_responses_request_body, @@ -41,17 +45,21 @@ pub(crate) use aether_ai_formats::api::{ convert_standard_chat_response, convert_standard_cli_response, copy_request_number_field, copy_request_number_field_as, core_error_background_report_kind, core_error_default_client_api_format, core_success_background_report_kind, + default_model_directive_mapping_patch, default_model_directive_suffixes, default_model_for_openai_image_operation, encode_done_sse, encode_json_sse, encode_kiro_sse_events, endpoint_config_forces_upstream_stream_policy, enforce_request_body_stream_field, estimate_kiro_tokens, extract_openai_text_content, + finalize_openai_provider_request, + finalize_openai_provider_request_with_codex_model_capabilities, find_kiro_real_thinking_end_tag, find_kiro_real_thinking_end_tag_at_buffer_end, - find_kiro_real_thinking_start_tag, force_upstream_streaming_for_provider, - gemini_request_is_image_generation, implicit_sync_finalize_report_kind, - is_core_error_finalize_kind, is_matching_stream_http_request, is_matching_stream_request, - is_openai_image_stream_request, is_openai_responses_family_format, is_openai_responses_format, - kiro_crc32, map_claude_stop_reason, map_openai_reasoning_effort_to_claude_output, - map_openai_reasoning_effort_to_gemini_budget, maybe_bridge_standard_sync_json_to_stream, - maybe_build_ai_surface_stream_rewriter, + find_kiro_real_thinking_start_tag, forbid_upstream_streaming_for_provider, + force_upstream_streaming_for_provider, gemini_request_is_image_generation, + implicit_sync_finalize_report_kind, is_core_error_finalize_kind, + is_matching_stream_http_request, is_matching_stream_request, is_openai_image_stream_request, + is_openai_responses_compact_format, is_openai_responses_family_format, + is_openai_responses_format, kiro_crc32, map_claude_stop_reason, + map_openai_reasoning_effort_to_claude_output, map_openai_reasoning_effort_to_gemini_budget, + maybe_bridge_standard_sync_json_to_stream, maybe_build_ai_surface_stream_rewriter, maybe_build_openai_chat_cross_format_sync_product_from_normalized_payload, maybe_build_openai_image_sync_finalize_product, maybe_build_openai_responses_cross_format_sync_product_from_normalized_payload, @@ -60,51 +68,60 @@ pub(crate) use aether_ai_formats::api::{ maybe_build_standard_cross_format_sync_product_from_normalized_payload, maybe_build_standard_same_format_sync_body_from_normalized_payload, maybe_build_standard_sync_finalize_product_from_normalized_payload, model_directive_base_model, - normalize_api_format_alias, normalize_claude_request_to_openai_chat_request, - normalize_gemini_request_to_openai_chat_request, normalize_openai_image_request, - normalize_openai_image_request_with_options, + model_directive_builtin_suffix_supported_for_source_model, + model_directive_suffix_has_builtin_mapping, normalize_api_format_alias, + normalize_claude_request_to_openai_chat_request, + normalize_gemini_request_to_openai_chat_request, normalize_openai_image_quality, + normalize_openai_image_request, normalize_openai_image_request_with_options, normalize_openai_responses_request_to_openai_chat_request, normalize_provider_private_report_context, normalize_provider_private_response_value, normalize_standard_request_to_openai_chat_request, openai_image_operation_from_path, - parse_direct_request_body, parse_openai_stop_sequences, parse_openai_tool_result_content, - prepare_local_success_response_parts, prepare_local_success_response_parts_owned, - provider_adaptation_allows_sync_finalize_envelope, provider_adaptation_anchor_api_format, - provider_adaptation_descriptor_for_envelope, provider_adaptation_descriptor_for_provider_type, + parse_codex_auth_identity, parse_direct_request_body, parse_model_directive, + parse_model_directive_with_suffixes, parse_openai_stop_sequences, + parse_openai_tool_result_content, prepare_local_success_response_parts, + prepare_local_success_response_parts_owned, project_codex_openai_image_api_request_body, + project_openai_image_api_request_body, provider_adaptation_allows_sync_finalize_envelope, + provider_adaptation_anchor_api_format, provider_adaptation_descriptor_for_envelope, + provider_adaptation_descriptor_for_provider_type, provider_adaptation_requires_eventstream_accept, provider_adaptation_should_unwrap_stream_envelope, provider_private_response_allows_sync_finalize, request_candidate_api_format_preference, request_candidate_api_formats, request_conversion_kind, request_conversion_requires_enable_flag, request_path_implies_stream_request, resolve_claude_stream_spec, resolve_claude_sync_spec, - 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_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_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, resolve_openai_responses_sync_spec, resolve_requested_gemini_image_model_for_request, resolve_requested_openai_image_model_for_request, + resolve_upstream_is_stream_for_provider as resolve_format_upstream_is_stream_for_provider, resolve_upstream_is_stream_from_endpoint_config, sanitize_request_path, sanitize_request_path_and_query, sanitize_request_query_string, stream_body_contains_error_event, supports_stream_execution_decision_kind, supports_sync_execution_decision_kind, sync_chat_response_conversion_kind, - sync_cli_response_conversion_kind, transform_provider_private_stream_line, value_as_u64, - AiControlPlanRequest, AiSurfaceFinalizeError, AiSurfaceStreamRewriter, CanonicalStreamFrame, + sync_cli_response_conversion_kind, transform_provider_private_stream_line, + validate_openai_provider_request_contract, value_as_u64, AiControlPlanRequest, + AiSurfaceFinalizeError, AiSurfaceStreamRewriter, CanonicalStreamFrame, ChatGptWebImageRequestError, ClaudeClientEmitter, ClaudeProviderState, - ExecutionRuntimeAuthContext, FinalizeStreamRewriteMode, FormatContext, GeminiClientEmitter, - GeminiImageRequestForOpenAi, GeminiProviderState, KiroToClaudeCliStreamState, - LocalCoreSyncErrorKind, LocalGeminiFilesSpec, LocalOpenAiImageSpec, LocalOpenAiResponsesSpec, - LocalSameFormatProviderFamily, LocalSameFormatProviderSpec, LocalStandardSourceFamily, - LocalStandardSourceMode, LocalStandardSpec, LocalSyncReportParts, LocalVideoCreateFamily, - LocalVideoCreateSpec, NormalizedOpenAiImageRequest, OpenAIChatClientEmitter, - OpenAIChatProviderState, OpenAIResponsesClientEmitter, OpenAIResponsesProviderState, - OpenAiImageNormalizeOptions, OpenAiImageOperation, OpenAiImageRequestForGemini, - OpenAiImageResponseFormat, OpenAiImageStreamState, OpenAiImageSyncFinalizeProduct, + CodexResponsesModelCapabilities, ExecutionRuntimeAuthContext, FinalizeStreamRewriteMode, + FormatContext, GeminiClientEmitter, GeminiImageRequestForOpenAi, GeminiProviderState, + KiroToClaudeCliStreamState, LocalCoreSyncErrorKind, LocalGeminiFilesSpec, LocalOpenAiImageSpec, + LocalOpenAiResponsesSpec, LocalSameFormatProviderFamily, LocalSameFormatProviderSpec, + LocalStandardSourceFamily, LocalStandardSourceMode, LocalStandardSpec, LocalSyncReportParts, + LocalVideoCreateFamily, LocalVideoCreateSpec, NormalizedOpenAiImageRequest, + OpenAIChatClientEmitter, OpenAIChatProviderState, OpenAIResponsesClientEmitter, + OpenAIResponsesProviderState, OpenAiImageNormalizeOptions, OpenAiImageOperation, + OpenAiImageRequestForGemini, OpenAiImageResponseFormat, OpenAiImageStreamState, + OpenAiImageSyncFinalizeProduct, OpenAiProviderRequestFinalization, ProviderAdaptationDescriptor, ProviderAdaptationSurface, ProviderPrivateStreamNormalizer, - RequestConversionKind, StandardCrossFormatSyncProduct, StandardSyncFinalizeNormalizedProduct, - StreamingStandardFormatMatrix, SyncChatResponseConversionKind, SyncCliResponseConversionKind, - SyncToStreamBridgeOutcome, ANTIGRAVITY_V1INTERNAL_ENVELOPE_NAME, CLAUDE_CHAT_STREAM_PLAN_KIND, + ReasoningEffort, RequestConversionKind, ServiceTier, StandardCrossFormatSyncProduct, + StandardSyncFinalizeNormalizedProduct, StreamingStandardFormatMatrix, + SyncChatResponseConversionKind, SyncCliResponseConversionKind, SyncToStreamBridgeOutcome, + ANTIGRAVITY_V1INTERNAL_ENVELOPE_NAME, 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, @@ -125,22 +142,28 @@ pub(crate) use aether_ai_formats::api::{ 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_VIDEO_CANCEL_SYNC_PLAN_KIND, GEMINI_VIDEO_CREATE_SYNC_FINALIZE_REPORT_KIND, GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND, - KIRO_ENVELOPE_NAME, KIRO_MAX_THINKING_BUFFER, OPENAI_CHAT_STREAM_PLAN_KIND, - OPENAI_CHAT_STREAM_SUCCESS_REPORT_KIND, OPENAI_CHAT_SYNC_ERROR_REPORT_KIND, - OPENAI_CHAT_SYNC_FINALIZE_REPORT_KIND, OPENAI_CHAT_SYNC_PLAN_KIND, - OPENAI_CHAT_SYNC_SUCCESS_REPORT_KIND, OPENAI_EMBEDDING_SYNC_PLAN_KIND, - OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_STREAM_SUCCESS_REPORT_KIND, - OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND, - OPENAI_IMAGE_SYNC_SUCCESS_REPORT_KIND, OPENAI_RERANK_SYNC_PLAN_KIND, - OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, OPENAI_RESPONSES_COMPACT_STREAM_SUCCESS_REPORT_KIND, + KIRO_ENVELOPE_NAME, KIRO_MAX_THINKING_BUFFER, MODEL_DIRECTIVE_API_FORMATS, + OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_STREAM_SUCCESS_REPORT_KIND, + OPENAI_CHAT_SYNC_ERROR_REPORT_KIND, OPENAI_CHAT_SYNC_FINALIZE_REPORT_KIND, + OPENAI_CHAT_SYNC_PLAN_KIND, OPENAI_CHAT_SYNC_SUCCESS_REPORT_KIND, + OPENAI_EMBEDDING_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND, + OPENAI_IMAGE_STREAM_SUCCESS_REPORT_KIND, OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND, + OPENAI_IMAGE_SYNC_PLAN_KIND, OPENAI_IMAGE_SYNC_SUCCESS_REPORT_KIND, + OPENAI_RERANK_SYNC_PLAN_KIND, OPENAI_RESPONSES_COMPACT_STREAM_PLAN_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_PLAN_KIND, OPENAI_RESPONSES_COMPACT_SYNC_SUCCESS_REPORT_KIND, OPENAI_RESPONSES_STREAM_PLAN_KIND, OPENAI_RESPONSES_STREAM_SUCCESS_REPORT_KIND, OPENAI_RESPONSES_SYNC_ERROR_REPORT_KIND, OPENAI_RESPONSES_SYNC_FINALIZE_REPORT_KIND, OPENAI_RESPONSES_SYNC_PLAN_KIND, - OPENAI_RESPONSES_SYNC_SUCCESS_REPORT_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, + OPENAI_RESPONSES_SYNC_SUCCESS_REPORT_KIND, OPENAI_SEARCH_SYNC_PLAN_KIND, + OPENAI_SEARCH_SYNC_SUCCESS_REPORT_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_VIDEO_CONTENT_PLAN_KIND, OPENAI_VIDEO_CREATE_SYNC_FINALIZE_REPORT_KIND, OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND, OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND, }; -pub(crate) use aether_ai_formats::{is_embedding_api_format, is_rerank_api_format}; +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, +}; diff --git a/apps/aether-gateway/src/api/ai/openai.rs b/apps/aether-gateway/src/api/ai/openai.rs index 999c356c6..85d9cbe2e 100644 --- a/apps/aether-gateway/src/api/ai/openai.rs +++ b/apps/aether-gateway/src/api/ai/openai.rs @@ -5,6 +5,7 @@ pub(crate) fn normalized_signature(api_format: &str) -> Option<&'static str> { "openai:rerank" => Some("openai:rerank"), "openai:responses" => Some("openai:responses"), "openai:responses:compact" => Some("openai:responses:compact"), + "openai:search" => Some("openai:search"), "openai:image" => Some("openai:image"), "openai:video" => Some("openai:video"), _ => None, @@ -18,6 +19,7 @@ pub(crate) fn local_path(api_format: &str) -> Option<&'static str> { "openai:rerank" => Some("/v1/rerank"), "openai:responses" => Some("/v1/responses"), "openai:responses:compact" => Some("/v1/responses/compact"), + "openai:search" => Some("/v1/alpha/search"), "openai:image" => Some("/v1/images/generations"), "openai:video" => Some("/v1/videos"), _ => None, diff --git a/apps/aether-gateway/src/api/ai/registry.rs b/apps/aether-gateway/src/api/ai/registry.rs index 5d92375c5..7ae1df4e6 100644 --- a/apps/aether-gateway/src/api/ai/registry.rs +++ b/apps/aether-gateway/src/api/ai/registry.rs @@ -15,6 +15,7 @@ const AI_POST_ROUTE_PATTERNS: &[&str] = &[ "/v1/messages/count_tokens", "/v1/responses", "/v1/responses/compact", + "/v1/alpha/search", "/v1/images/generations", "/v1/images/edits", "/v1/interactions", @@ -120,6 +121,7 @@ mod tests { "/api/v1/services/embeddings/multimodal-embedding/multimodal-embedding", ), ("openai:rerank", "openai", "rerank", "/v1/rerank"), + ("openai:search", "openai", "search", "/v1/alpha/search"), ("jina:rerank", "jina", "rerank", "/v1/rerank"), ] { assert_eq!( diff --git a/apps/aether-gateway/src/cache/candidate_page.rs b/apps/aether-gateway/src/cache/candidate_page.rs index 7ed3fe2c0..0c22546bf 100644 --- a/apps/aether-gateway/src/cache/candidate_page.rs +++ b/apps/aether-gateway/src/cache/candidate_page.rs @@ -92,6 +92,7 @@ pub(crate) struct CandidatePageCacheKey { preselection_mode: &'static str, use_api_format_alias_match: bool, client_session_affinity_hash: String, + model_directive_policy_hash: String, } #[derive(Debug, Clone, PartialEq, Eq, Hash)] @@ -140,6 +141,7 @@ impl CandidatePageCacheKey { preselection_mode: &'static str, use_api_format_alias_match: bool, client_session_affinity: Option<&ClientSessionAffinity>, + model_directive_policy_hash: &str, ) -> Self { Self { requested_model: normalize_text_key(requested_model), @@ -153,6 +155,7 @@ impl CandidatePageCacheKey { preselection_mode, use_api_format_alias_match, client_session_affinity_hash: client_session_affinity_key(client_session_affinity), + model_directive_policy_hash: normalize_text_key(model_directive_policy_hash), } } } @@ -171,6 +174,7 @@ impl CandidateResolvedPageCacheKey { preselection_mode: &'static str, use_api_format_alias_match: bool, client_session_affinity: Option<&ClientSessionAffinity>, + model_directive_policy_hash: &str, resolution_mode: AiCandidateResolutionMode, ) -> Self { Self { @@ -186,6 +190,7 @@ impl CandidateResolvedPageCacheKey { preselection_mode, use_api_format_alias_match, client_session_affinity, + model_directive_policy_hash, ), resolution_mode: resolution_mode_name(resolution_mode), } @@ -569,6 +574,7 @@ mod tests { "provider_endpoint_key_model", true, None, + "policy-a", ); let different_user = CandidatePageCacheKey::new( "gpt-4o", @@ -582,6 +588,7 @@ mod tests { "provider_endpoint_key_model", true, None, + "policy-a", ); let different_model = CandidatePageCacheKey::new( "gpt-4.1", @@ -595,6 +602,7 @@ mod tests { "provider_endpoint_key_model", true, None, + "policy-a", ); let different_format = CandidatePageCacheKey::new( "gpt-4o", @@ -608,6 +616,7 @@ mod tests { "provider_endpoint_key_model", true, None, + "policy-a", ); let different_capabilities = CandidatePageCacheKey::new( "gpt-4o", @@ -621,11 +630,90 @@ mod tests { "provider_endpoint_key_model", true, None, + "policy-a", + ); + let same_policy = CandidatePageCacheKey::new( + "gpt-4o", + "openai:chat", + true, + &auth_a, + Some(&json!({"vision": true})), + None, + Some("bearer"), + 7, + "provider_endpoint_key_model", + true, + None, + "policy-a", + ); + let different_policy = CandidatePageCacheKey::new( + "gpt-4o", + "openai:chat", + true, + &auth_a, + Some(&json!({"vision": true})), + None, + Some("bearer"), + 7, + "provider_endpoint_key_model", + true, + None, + "policy-b", ); + assert_eq!(base, same_policy); assert_ne!(base, different_user); assert_ne!(base, different_model); assert_ne!(base, different_format); assert_ne!(base, different_capabilities); + assert_ne!(base, different_policy); + + let resolved_base = CandidateResolvedPageCacheKey::new( + "gpt-4o", + "openai:chat", + true, + &auth_a, + Some(&json!({"vision": true})), + None, + Some("bearer"), + 7, + "provider_endpoint_key_model", + true, + None, + "policy-a", + AiCandidateResolutionMode::Standard, + ); + let resolved_same_policy = CandidateResolvedPageCacheKey::new( + "gpt-4o", + "openai:chat", + true, + &auth_a, + Some(&json!({"vision": true})), + None, + Some("bearer"), + 7, + "provider_endpoint_key_model", + true, + None, + "policy-a", + AiCandidateResolutionMode::Standard, + ); + let resolved_different_policy = CandidateResolvedPageCacheKey::new( + "gpt-4o", + "openai:chat", + true, + &auth_a, + Some(&json!({"vision": true})), + None, + Some("bearer"), + 7, + "provider_endpoint_key_model", + true, + None, + "policy-b", + AiCandidateResolutionMode::Standard, + ); + assert_eq!(resolved_base, resolved_same_policy); + assert_ne!(resolved_base, resolved_different_policy); } } diff --git a/apps/aether-gateway/src/client_session_affinity.rs b/apps/aether-gateway/src/client_session_affinity.rs index af802c51f..66e009184 100644 --- a/apps/aether-gateway/src/client_session_affinity.rs +++ b/apps/aether-gateway/src/client_session_affinity.rs @@ -101,6 +101,29 @@ pub(crate) fn client_session_affinity_from_request( client_session_scope_from_request(headers, body_json)?.scheduler_affinity() } +pub(crate) fn client_session_affinity_from_api_request( + api_format: &str, + headers: &http::HeaderMap, + body_json: Option<&Value>, +) -> Option { + client_session_scope_from_api_request(api_format, headers, body_json)?.scheduler_affinity() +} + +fn client_session_scope_from_api_request( + api_format: &str, + headers: &http::HeaderMap, + body_json: Option<&Value>, +) -> Option { + if api_format.trim().eq_ignore_ascii_case("openai:search") { + let request = ClientSessionRequest { headers, body_json }; + return explicit_aether_session_scope(&request, CodexSessionScopeAdapter.family()) + .or_else(|| codex_search_session_scope(&request)) + .or_else(|| client_session_scope_from_request(headers, body_json)); + } + + client_session_scope_from_request(headers, body_json) +} + pub(crate) fn client_session_scope_from_request( headers: &http::HeaderMap, body_json: Option<&Value>, @@ -113,6 +136,22 @@ pub(crate) fn client_session_scope_from_request( .or_else(|| extract_scope_from_other_specific_adapters(&request, client_family.as_str())) } +fn codex_search_session_scope(request: &ClientSessionRequest<'_>) -> Option { + let session_id = request + .body_json? + .get("id")? + .as_str() + .map(str::trim) + .filter(|value| !value.is_empty())?; + Some(ClientSessionScope::new( + CodexSessionScopeAdapter.family(), + session_id, + None, + header_value_str(request.headers, "chatgpt-account-id"), + ClientSessionSignalSource::Body, + )) +} + pub(crate) fn client_session_affinity_from_parts( parts: &http::request::Parts, body_json: Option<&Value>, @@ -307,7 +346,9 @@ impl ClientSessionScopeAdapter for GenericSessionScopeAdapter { } fn extract_scope(&self, request: &ClientSessionRequest<'_>) -> Option { - if let Some(root_session) = header_value_str(request.headers, "session_id") + if let Some(root_session) = header_value_str(request.headers, "session-id") + .or_else(|| header_value_str(request.headers, "thread-id")) + .or_else(|| header_value_str(request.headers, "session_id")) .or_else(|| header_value_str(request.headers, "conversation_id")) { return Some(ClientSessionScope::new( @@ -367,7 +408,9 @@ impl ClientSessionScopeAdapter for CodexSessionScopeAdapter { } fn extract_scope(&self, request: &ClientSessionRequest<'_>) -> Option { - header_value_str(request.headers, "session_id") + header_value_str(request.headers, "session-id") + .or_else(|| header_value_str(request.headers, "thread-id")) + .or_else(|| header_value_str(request.headers, "session_id")) .or_else(|| header_value_str(request.headers, "conversation_id")) .map(|root_session| { ClientSessionScope::new( @@ -785,6 +828,7 @@ fn has_header_with_prefix(headers: &http::HeaderMap, prefix: &str) -> bool { #[cfg(test)] mod tests { use super::{ + client_session_affinity_from_api_request, client_session_affinity_from_report_context_value, client_session_affinity_from_request, client_session_affinity_report_context_value, client_session_scope_from_request, ClientSessionSignalSource, AETHER_AGENT_ID_HEADER, AETHER_SESSION_ID_HEADER, @@ -841,9 +885,10 @@ mod tests { let mut headers = HeaderMap::new(); headers.insert( http::header::USER_AGENT, - HeaderValue::from_static("codex-tui/0.122.0"), + HeaderValue::from_static("codex_cli_rs/0.144.1"), ); - headers.insert("session_id", HeaderValue::from_static("codex-session")); + headers.insert("session-id", HeaderValue::from_static("codex-session")); + headers.insert("thread-id", HeaderValue::from_static("codex-thread")); let affinity = client_session_affinity_from_request(&headers, None).expect("affinity should build"); @@ -860,7 +905,7 @@ mod tests { let mut headers = HeaderMap::new(); headers.insert( http::header::USER_AGENT, - HeaderValue::from_static("codex-tui/0.122.0"), + HeaderValue::from_static("codex_cli_rs/0.144.1"), ); headers.insert( "x-client-request-id", @@ -1120,4 +1165,28 @@ mod tests { assert!(client_session_affinity_from_request(&headers, Some(&body)).is_none()); } + + #[test] + fn codex_search_uses_request_id_as_session_affinity() { + let mut headers = HeaderMap::new(); + headers.insert("chatgpt-account-id", HeaderValue::from_static("account-1")); + let body = json!({"id": "codex-session-1", "model": "gpt-5.6"}); + + let affinity = + client_session_affinity_from_api_request("openai:search", &headers, Some(&body)) + .expect("search affinity should build"); + + assert_eq!(affinity.client_family.as_deref(), Some("codex")); + assert_eq!( + affinity.session_key.as_deref(), + Some("account=account-1;session=codex-session-1") + ); + } + + #[test] + fn top_level_request_id_is_not_a_generic_session_signal() { + let body = json!({"id": "request-id", "model": "gpt-5.6"}); + + assert!(client_session_affinity_from_request(&HeaderMap::new(), Some(&body)).is_none()); + } } diff --git a/apps/aether-gateway/src/constants.rs b/apps/aether-gateway/src/constants.rs index 8052f2001..44e2f4788 100644 --- a/apps/aether-gateway/src/constants.rs +++ b/apps/aether-gateway/src/constants.rs @@ -125,6 +125,7 @@ pub(crate) const RUST_FRONTDOOR_OWNED_ROUTE_PATTERNS: &[&str] = &[ "/v1/messages/count_tokens", "/v1/responses", "/v1/responses/compact", + "/v1/alpha/search", "/v1/models/{model}:generateContent", "/v1/models/{model}:streamGenerateContent", "/v1/models/{model}:predictLongRunning", diff --git a/apps/aether-gateway/src/control/auth/gate.rs b/apps/aether-gateway/src/control/auth/gate.rs index f75a898a1..3a63bad65 100644 --- a/apps/aether-gateway/src/control/auth/gate.rs +++ b/apps/aether-gateway/src/control/auth/gate.rs @@ -70,12 +70,10 @@ pub(crate) async fn request_model_local_rejection( ) { if !contains_string(allowed_models, requested_model) && !model_directive_base_model_is_allowed_for_request( - state, decision, requested_model, allowed_models, ) - .await && !request_model_resolves_to_allowed_model( state, decision, @@ -90,31 +88,46 @@ pub(crate) async fn request_model_local_rejection( } } - balance_capacity_rejection( - state, - decision, - auth_context, - requested_model.as_deref(), - headers, - body, - ) - .await + Ok(None) } -async fn balance_capacity_rejection( +pub(crate) async fn execution_plan_balance_capacity_rejection( state: &AppState, decision: &GatewayControlDecision, - auth_context: &GatewayControlAuthContext, - requested_model: Option<&str>, - headers: &http::HeaderMap, - body: &Bytes, + plan: &aether_contracts::ExecutionPlan, + report_context: Option<&serde_json::Value>, ) -> Result, GatewayError> { - if auth_context.api_key_is_standalone { + let Some(auth_context) = decision.auth_context.as_ref() else { + return Ok(None); + }; + if auth_context.api_key_is_standalone || auth_context.local_rejection.is_some() { return Ok(None); } - if auth_context.local_rejection.is_some() { + let Some(available_usd) = available_balance_capacity_usd(state, auth_context).await? else { return Ok(None); + }; + match estimate_execution_plan_cost_upper_bound_usd(state, plan, report_context).await? { + Some(estimated_cost_usd) + if estimated_cost_usd <= available_usd + DAILY_QUOTA_EPSILON_USD => + { + Ok(None) + } + Some(_) | None if available_usd <= DAILY_QUOTA_EPSILON_USD => { + Ok(Some(GatewayLocalAuthRejection::BalanceDenied { + remaining: Some(0.0), + })) + } + Some(_) => Ok(Some(GatewayLocalAuthRejection::BalanceDenied { + remaining: Some(available_usd), + })), + None => Ok(None), } +} + +async fn available_balance_capacity_usd( + state: &AppState, + auth_context: &GatewayControlAuthContext, +) -> Result, GatewayError> { let quota = state .find_user_daily_quota_availability_for_auth(&auth_context.user_id) .await? @@ -130,36 +143,13 @@ async fn balance_capacity_rejection( let wallet_is_unlimited = wallet .as_ref() .is_some_and(|wallet| wallet.limit_mode.eq_ignore_ascii_case("unlimited")); - let available_usd = match quota.as_ref() { + Ok(match quota.as_ref() { Some(quota) if !quota.allow_wallet_overage => Some(quota.remaining_usd.max(0.0)), Some(_) if wallet_is_unlimited => None, Some(quota) => Some(quota.remaining_usd.max(0.0) + wallet_available_usd.unwrap_or(0.0)), None if wallet_is_unlimited => None, None => wallet_available_usd, - }; - let Some(available_usd) = available_usd else { - return Ok(None); - }; - if available_usd <= DAILY_QUOTA_EPSILON_USD { - return Ok(Some(GatewayLocalAuthRejection::BalanceDenied { - remaining: Some(0.0), - })); - } - let Some(requested_model) = requested_model else { - return Ok(None); - }; - let Some(estimated_cost_usd) = - estimate_request_cost_upper_bound_usd(state, decision, requested_model, headers, body) - .await? - else { - return Ok(None); - }; - if estimated_cost_usd > available_usd + DAILY_QUOTA_EPSILON_USD { - return Ok(Some(GatewayLocalAuthRejection::BalanceDenied { - remaining: Some(available_usd), - })); - } - Ok(None) + }) } fn wallet_finite_available_usd( @@ -173,262 +163,300 @@ fn wallet_finite_available_usd( Some(wallet.balance.max(0.0) + wallet.gift_balance.max(0.0)) } -async fn estimate_request_cost_upper_bound_usd( +async fn estimate_execution_plan_cost_upper_bound_usd( state: &AppState, - decision: &GatewayControlDecision, - requested_model: &str, - headers: &http::HeaderMap, - body: &Bytes, + plan: &aether_contracts::ExecutionPlan, + report_context: Option<&serde_json::Value>, ) -> Result, GatewayError> { - let Some(api_format) = decision - .auth_endpoint_signature - .as_deref() - .map(crate::ai_serving::normalize_api_format_alias) - .filter(|value| !value.trim().is_empty()) - else { + let api_format = crate::ai_serving::normalize_api_format_alias(&plan.provider_api_format); + let Some(task_type) = authorization_task_type(&api_format, report_context) else { return Ok(None); }; - let body = crate::headers::decoded_request_body_bytes(headers, body.as_ref()).ok(); - let Some(body) = body else { + let Some(body_json) = plan.body.json_body.as_ref() else { return Ok(None); }; - let body_json = serde_json::from_slice::(body.as_ref()).ok(); - let Some(input_tokens) = body_json - .as_ref() - .map(estimate_json_tokens) - .filter(|value| *value > 0) - else { + if !openai_request_input_is_self_contained(&api_format, body_json) { + return Ok(None); + } + let input_tokens = json_token_count_upper_bound(body_json); + let Ok(input_tokens) = i64::try_from(input_tokens) else { return Ok(None); }; - let max_output_tokens = body_json.as_ref().and_then(max_output_tokens_from_request); - let cache_key = auth_request_cost_upper_bound_cache_key( + let max_output_tokens = max_output_tokens_from_request(body_json) + .map(|value| value.saturating_mul(output_choice_count_upper_bound(&api_format, body_json))) + .and_then(|value| i64::try_from(value).ok()); + let requested_processing_tier = + aether_data_contracts::repository::usage::extract_provider_service_tier_from_body(Some( + body_json, + )); + let model_id = report_context_string_field(report_context, "model_id"); + let global_model_name = report_context_string_field(report_context, "global_model_name"); + let cache_ttl_minutes = + aether_data_contracts::repository::usage::resolve_provider_cache_ttl_minutes( + Some(&api_format), + plan.model_name.as_deref(), + global_model_name, + Some(body_json), + ); + if model_id.is_none() && global_model_name.is_none() { + return Ok(None); + } + let cache_key = execution_plan_cost_upper_bound_cache_key( + plan, + model_id, + global_model_name, &api_format, - requested_model, input_tokens, max_output_tokens, + requested_processing_tier.as_deref(), + cache_ttl_minutes, ); let ttl = state.frontdoor_runtime_guards.auth_capacity_cache_ttl; if ttl.is_zero() { - return estimate_request_cost_upper_bound_for_tokens( + return calculate_execution_plan_cost_upper_bound( state, + plan, + model_id, + global_model_name, &api_format, - requested_model, + task_type, input_tokens, max_output_tokens, + requested_processing_tier.as_deref(), + cache_ttl_minutes, ) .await; } state .auth_request_cost_upper_bound_cache .get_or_load(cache_key, ttl, || async { - estimate_request_cost_upper_bound_for_tokens( + calculate_execution_plan_cost_upper_bound( state, + plan, + model_id, + global_model_name, &api_format, - requested_model, + task_type, input_tokens, max_output_tokens, + requested_processing_tier.as_deref(), + cache_ttl_minutes, ) .await }) .await } -async fn estimate_request_cost_upper_bound_for_tokens( +#[allow(clippy::too_many_arguments)] +async fn calculate_execution_plan_cost_upper_bound( state: &AppState, + plan: &aether_contracts::ExecutionPlan, + model_id: Option<&str>, + global_model_name: Option<&str>, api_format: &str, - requested_model: &str, - input_tokens: u64, - max_output_tokens: Option, + task_type: &str, + input_tokens: i64, + max_output_tokens: Option, + requested_processing_tier: Option<&str>, + cache_ttl_minutes: Option, ) -> Result, GatewayError> { - let candidates = state - .list_minimal_candidate_selection_rows_for_api_format_and_requested_model( - api_format, - requested_model, - ) - .await?; - let mut max_estimate = None::; - for candidate in candidates { - let context = state + let context = match model_id { + Some(model_id) => state .data - .find_billing_model_context_by_model_id( - &candidate.provider_id, - Some(&candidate.key_id), - &candidate.model_id, + .find_billing_model_context_by_model_id(&plan.provider_id, Some(&plan.key_id), model_id) + .await + .map_err(|err| GatewayError::Internal(err.to_string()))?, + None => state + .data + .find_billing_model_context( + &plan.provider_id, + Some(&plan.key_id), + global_model_name.expect("global model name should exist"), ) .await - .map_err(|err| GatewayError::Internal(err.to_string()))?; - let Some(context) = context else { - continue; - }; - let Some(estimate) = estimate_cost_from_billing_context( - &context, - api_format, - input_tokens, - max_output_tokens, - ) else { - return Ok(None); - }; - max_estimate = Some(max_estimate.map_or(estimate, |current| current.max(estimate))); - } - Ok(max_estimate.filter(|value| value.is_finite() && *value >= 0.0)) + .map_err(|err| GatewayError::Internal(err.to_string()))?, + }; + let Some(context) = context else { + return Ok(None); + }; + let mut estimate = + aether_billing::BillingAuthorizationEstimateInput::new(task_type, input_tokens); + estimate.api_format = Some(api_format.to_string()); + estimate.requested_processing_tier = requested_processing_tier.map(ToOwned::to_owned); + estimate.cache_ttl_minutes = cache_ttl_minutes; + estimate.max_output_tokens = max_output_tokens; + aether_billing::BillingService::new() + .estimate_authorization_cost_upper_bound( + &aether_billing::BillingModelPricingSnapshot::from(context), + &estimate, + ) + .map_err(|err| GatewayError::Internal(err.to_string())) } -fn auth_request_cost_upper_bound_cache_key( +fn execution_plan_cost_upper_bound_cache_key( + plan: &aether_contracts::ExecutionPlan, + model_id: Option<&str>, + global_model_name: Option<&str>, api_format: &str, - requested_model: &str, - input_tokens: u64, - max_output_tokens: Option, + input_tokens: i64, + max_output_tokens: Option, + requested_processing_tier: Option<&str>, + cache_ttl_minutes: Option, ) -> String { format!( - "{}\x1f{}\x1f{}\x1f{}", + "{}\x1f{}\x1f{}\x1f{}\x1f{}\x1f{}\x1f{}\x1f{}\x1f{}", + plan.provider_id, + plan.key_id, + model_id.unwrap_or(""), + global_model_name.unwrap_or(""), api_format, - requested_model, input_tokens, max_output_tokens .map(|value| value.to_string()) - .unwrap_or_else(|| "none".to_string()) + .unwrap_or_else(|| "none".to_string()), + requested_processing_tier.unwrap_or("standard"), + cache_ttl_minutes + .map(|value| value.to_string()) + .unwrap_or_else(|| "none".to_string()), ) } -fn estimate_cost_from_billing_context( - context: &aether_data_contracts::repository::billing::StoredBillingModelContext, +fn authorization_task_type<'a>( api_format: &str, - input_tokens: u64, - max_output_tokens: Option, -) -> Option { - if context - .provider_billing_type - .as_deref() - .is_some_and(|value| value.eq_ignore_ascii_case("free_tier")) + report_context: Option<&'a serde_json::Value>, +) -> Option<&'a str> { + if report_context + .and_then(|context| context.get("image_request")) + .is_some() + || api_format == "openai:image" { - return Some(0.0); + return None; } - let price_per_request = context - .model_price_per_request - .or(context.default_price_per_request) - .filter(|value| value.is_finite() && *value >= 0.0) - .unwrap_or(0.0); - let tiered_pricing = effective_tiered_pricing(context); - let input_price_per_1m = tiered_price_per_1m(tiered_pricing, "input_price_per_1m") - .filter(|value| value.is_finite() && *value >= 0.0) - .unwrap_or(0.0); - let output_price_per_1m = tiered_price_per_1m(tiered_pricing, "output_price_per_1m") - .filter(|value| value.is_finite() && *value >= 0.0) - .unwrap_or(0.0); - let output_tokens = if output_price_per_1m > 0.0 { - max_output_tokens? - } else { - 0 - }; - let estimate = price_per_request - + (input_tokens as f64 * input_price_per_1m / 1_000_000.0) - + (output_tokens as f64 * output_price_per_1m / 1_000_000.0); - let rate_multiplier = rate_multiplier_for_api_format(context, api_format); - Some(estimate * rate_multiplier) + if api_format.ends_with(":embedding") { + return Some("embedding"); + } + if api_format.ends_with(":rerank") { + return Some("rerank"); + } + Some("chat") } -fn effective_tiered_pricing( - context: &aether_data_contracts::repository::billing::StoredBillingModelContext, -) -> Option<&serde_json::Value> { - context - .model_tiered_pricing - .as_ref() - .filter(|value| tiered_pricing_has_rates(value)) - .or(context.default_tiered_pricing.as_ref()) -} - -fn tiered_pricing_has_rates(value: &serde_json::Value) -> bool { - value - .get("tiers") - .and_then(serde_json::Value::as_array) - .is_some_and(|tiers| !tiers.is_empty()) - || ["input_price_per_1m", "output_price_per_1m"] - .iter() - .any(|field| { - value - .get(*field) - .and_then(serde_json::Value::as_f64) - .is_some() - }) -} - -fn rate_multiplier_for_api_format( - context: &aether_data_contracts::repository::billing::StoredBillingModelContext, - api_format: &str, -) -> f64 { - let normalized_api_format = api_format.trim().to_ascii_lowercase(); - let Some(mapping) = context - .provider_api_key_rate_multipliers - .as_ref() - .and_then(serde_json::Value::as_object) - else { - return 1.0; - }; - mapping - .get(&normalized_api_format) - .and_then(serde_json::Value::as_f64) - .filter(|value| value.is_finite() && *value >= 0.0) - .unwrap_or(1.0) -} - -fn tiered_price_per_1m(tiered_pricing: Option<&serde_json::Value>, field: &str) -> Option { - let value = tiered_pricing?; - value - .get(field) - .and_then(serde_json::Value::as_f64) - .or_else(|| { - value - .get("tiers") - .and_then(serde_json::Value::as_array) - .into_iter() - .flatten() - .filter_map(|tier| tier.get(field).and_then(serde_json::Value::as_f64)) - .filter(|price| price.is_finite() && *price >= 0.0) - .max_by(|left, right| left.partial_cmp(right).unwrap_or(std::cmp::Ordering::Equal)) - }) +fn report_context_string_field<'a>( + report_context: Option<&'a serde_json::Value>, + key: &str, +) -> Option<&'a str> { + report_context + .and_then(|context| context.get(key)) + .and_then(serde_json::Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) } fn max_output_tokens_from_request(value: &serde_json::Value) -> Option { ["max_tokens", "max_completion_tokens", "max_output_tokens"] .iter() - .find_map(|field| value.get(*field).and_then(serde_json::Value::as_u64)) + .filter_map(|field| value.get(*field).and_then(serde_json::Value::as_u64)) .filter(|value| *value > 0) + .max() } -fn estimate_json_tokens(value: &serde_json::Value) -> u64 { +fn output_choice_count_upper_bound(api_format: &str, value: &serde_json::Value) -> u64 { + if api_format != "openai:chat" { + return 1; + } + value + .get("n") + .and_then(serde_json::Value::as_u64) + .filter(|value| *value > 0) + .unwrap_or(1) +} + +fn openai_request_input_is_self_contained(api_format: &str, value: &serde_json::Value) -> bool { + if !api_format.starts_with("openai:") { + return false; + } + let Some(object) = value.as_object() else { + return true; + }; + if ["previous_response_id", "conversation"] + .iter() + .any(|key| object.get(*key).is_some_and(has_reference_value)) + { + return false; + } + if object + .get("prompt") + .and_then(serde_json::Value::as_object) + .and_then(|prompt| prompt.get("id")) + .is_some_and(has_reference_value) + { + return false; + } + !contains_indirect_request_input(value) +} + +fn contains_indirect_request_input(value: &serde_json::Value) -> bool { match value { - serde_json::Value::String(text) => estimate_text_tokens(text), - serde_json::Value::Array(items) => items - .iter() - .map(estimate_json_tokens) - .fold(0u64, u64::saturating_add), - serde_json::Value::Object(object) => object - .iter() - .map(|(key, value)| { - estimate_text_tokens(key).saturating_add(estimate_json_tokens(value)) - }) - .fold(0u64, u64::saturating_add), - _ => 1, + serde_json::Value::Array(items) => items.iter().any(contains_indirect_request_input), + serde_json::Value::Object(object) => { + let item_type = object + .get("type") + .and_then(serde_json::Value::as_str) + .map(|value| value.trim().to_ascii_lowercase()); + if item_type.as_deref().is_some_and(|item_type| { + matches!( + item_type, + "url" + | "item_reference" + | "input_file" + | "input_image" + | "input_audio" + | "image_url" + | "file_search" + | "web_search" + | "web_search_preview" + | "computer_use" + | "computer_use_preview" + | "code_interpreter" + | "mcp" + | "image_generation" + ) + }) { + return true; + } + if ["file_id", "file_uri", "fileUri", "vector_store_ids"] + .iter() + .any(|key| object.get(*key).is_some_and(has_reference_value)) + { + return true; + } + object.values().any(contains_indirect_request_input) + } + _ => false, } } -fn estimate_text_tokens(text: &str) -> u64 { - let chars = text.chars().count() as u64; - chars.div_ceil(4).max(1) +fn has_reference_value(value: &serde_json::Value) -> bool { + match value { + serde_json::Value::Null => false, + serde_json::Value::String(value) => !value.trim().is_empty(), + serde_json::Value::Array(values) => !values.is_empty(), + serde_json::Value::Object(values) => !values.is_empty(), + _ => true, + } } -async fn model_directive_base_model_is_allowed_for_request( - state: &AppState, +fn json_token_count_upper_bound(value: &serde_json::Value) -> u64 { + serde_json::to_vec(value) + .map(|bytes| u64::try_from(bytes.len()).unwrap_or(u64::MAX)) + .unwrap_or(u64::MAX) +} + +fn model_directive_base_model_is_allowed_for_request( decision: &GatewayControlDecision, requested_model: &str, allowed_models: &[String], ) -> bool { - let Some(base_model) = crate::ai_serving::model_directive_base_model(requested_model) else { - return false; - }; - if !contains_string(allowed_models, &base_model) { - return false; - } let Some(client_api_format) = decision .auth_endpoint_signature .as_deref() @@ -438,12 +466,12 @@ async fn model_directive_base_model_is_allowed_for_request( return false; }; for api_format in candidate_api_formats_for_model_resolution(&client_api_format) { - if crate::system_features::reasoning_model_directive_enabled_for_api_format_and_model( - state, - &api_format, - Some(requested_model), - ) - .await + let resolution = decision + .model_directive_policy + .resolve_reasoning(&api_format, Some(requested_model)); + if resolution + .base_model() + .is_some_and(|base_model| contains_string(allowed_models, base_model)) { return true; } @@ -467,13 +495,10 @@ async fn request_model_resolves_to_allowed_model( }; for api_format in candidate_api_formats_for_model_resolution(&client_api_format) { - let enable_model_directives = - crate::system_features::reasoning_model_directive_enabled_for_api_format_and_model( - state, - &api_format, - Some(requested_model), - ) - .await; + let resolution = decision + .model_directive_policy + .resolve_reasoning(&api_format, Some(requested_model)); + let routing_model = resolution.base_model().unwrap_or(requested_model); let rows = state .list_minimal_candidate_selection_rows_for_api_format(&api_format) .await?; @@ -482,18 +507,18 @@ async fn request_model_resolves_to_allowed_model( .filter(|row| { aether_scheduler_core::row_supports_requested_model_with_model_directives( row, - requested_model, + routing_model, &api_format, - enable_model_directives, + false, ) }) .collect::>(); let Some(resolved_global_model) = aether_scheduler_core::resolve_requested_global_model_name_with_model_directives( &matching_rows, - requested_model, + routing_model, &api_format, - enable_model_directives, + false, ) else { continue; @@ -525,9 +550,11 @@ fn push_unique_api_format(api_formats: &mut Vec, api_format: &str) { #[cfg(test)] mod tests { + use std::collections::BTreeMap; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; + use aether_contracts::{ExecutionPlan, RequestBody}; use aether_data::repository::candidate_selection::InMemoryMinimalCandidateSelectionReadRepository; use aether_data::repository::wallet::StoredWalletSnapshot; use aether_data_contracts::repository::billing::{ @@ -543,8 +570,9 @@ mod tests { use serde_json::json; use super::{ - estimate_cost_from_billing_context, request_model_local_rejection, - GatewayLocalAuthRejection, + execution_plan_balance_capacity_rejection, execution_plan_cost_upper_bound_cache_key, + max_output_tokens_from_request, openai_request_input_is_self_contained, + output_choice_count_upper_bound, request_model_local_rejection, GatewayLocalAuthRejection, }; use crate::control::{GatewayControlAuthContext, GatewayControlDecision}; use crate::data::GatewayDataState; @@ -668,6 +696,55 @@ mod tests { state_with_rows(vec![sample_row()]) } + fn execution_plan(body: serde_json::Value, api_format: &str) -> ExecutionPlan { + ExecutionPlan { + request_id: "request-1".to_string(), + candidate_id: Some("candidate-1".to_string()), + provider_name: Some("OpenAI".to_string()), + provider_id: "provider-1".to_string(), + endpoint_id: "endpoint-1".to_string(), + key_id: "key-1".to_string(), + method: "POST".to_string(), + url: "https://api.openai.com/v1/responses".to_string(), + headers: BTreeMap::new(), + content_type: Some("application/json".to_string()), + content_encoding: None, + body: RequestBody::from_json(body), + stream: false, + client_api_format: api_format.to_string(), + provider_api_format: api_format.to_string(), + model_name: Some("gpt-5".to_string()), + proxy: None, + transport_profile: None, + timeouts: None, + } + } + + fn billing_report_context() -> serde_json::Value { + json!({ + "model_id": "model-1", + "global_model_name": "gpt-5" + }) + } + + fn estimate_from_billing_context( + context: &StoredBillingModelContext, + api_format: &str, + input_tokens: i64, + max_output_tokens: Option, + ) -> Option { + let mut estimate = + aether_billing::BillingAuthorizationEstimateInput::new("chat", input_tokens); + estimate.api_format = Some(api_format.to_string()); + estimate.max_output_tokens = max_output_tokens; + aether_billing::BillingService::new() + .estimate_authorization_cost_upper_bound( + &aether_billing::BillingModelPricingSnapshot::from(context), + &estimate, + ) + .expect("estimate should calculate") + } + fn json_headers() -> HeaderMap { let mut headers = HeaderMap::new(); headers.insert( @@ -877,6 +954,58 @@ mod tests { ); } + #[tokio::test] + async fn model_rejection_reuses_request_policy_snapshot_for_directive_base_model() { + let state = state_with_rows(Vec::new()); + let mut decision = decision_with_allowed_models(vec!["gpt-5.6-sol".to_string()]); + decision.model_directive_policy = + crate::system_features::ModelDirectivePolicySnapshot::from_config_values( + Some(&json!(true)), + None, + ); + let uri: Uri = "/v1/chat/completions".parse().expect("uri should parse"); + let body = Bytes::from_static(br#"{"model":"gpt-5.6-sol-high","messages":[]}"#); + + let rejection = + request_model_local_rejection(&state, Some(&decision), &uri, &json_headers(), &body) + .await + .expect("model rejection should resolve"); + + assert_eq!(rejection, None); + } + + #[tokio::test] + async fn model_rejection_uses_custom_policy_suffix_for_base_model_authorization() { + let state = state_with_rows(Vec::new()); + let mut decision = decision_with_allowed_models(vec!["deployment-alias".to_string()]); + decision.model_directive_policy = + crate::system_features::ModelDirectivePolicySnapshot::from_config_values( + Some(&json!(true)), + Some(&json!({ + "reasoning_effort": { + "api_formats": { + "openai:chat": { + "suffixes": ["VendorFuture"], + "mappings": { + "VendorFuture": { "reasoning_effort": "high" } + } + } + } + } + })), + ); + let uri: Uri = "/v1/chat/completions".parse().expect("uri should parse"); + let body = + Bytes::from_static(br#"{"model":"deployment-alias-VendorFuture","messages":[]}"#); + + let rejection = + request_model_local_rejection(&state, Some(&decision), &uri, &json_headers(), &body) + .await + .expect("model rejection should resolve"); + + assert_eq!(rejection, None); + } + #[tokio::test] async fn positive_balance_allows_unbounded_output_request_without_cost_estimate() { let context = billing_context_with_pricing( @@ -897,17 +1026,20 @@ mod tests { context.clone(), ); let decision = decision_with_allowed_models(vec!["gpt-5".to_string()]); - let uri: Uri = "/v1/chat/completions".parse().expect("uri should parse"); - let body = Bytes::from_static( - br#"{"model":"gpt-5","messages":[{"role":"user","content":"hi"}],"stream":true}"#, + let plan = execution_plan( + json!({ + "model": "gpt-5", + "messages": [{"role": "user", "content": "hi"}], + "stream": true + }), + "openai:chat", ); - let rejection = request_model_local_rejection( + let rejection = execution_plan_balance_capacity_rejection( &state, - Some(&decision), - &uri, - &json_headers(), - &body, + &decision, + &plan, + Some(&billing_report_context()), ) .await .expect("quota rejection should resolve"); @@ -951,15 +1083,24 @@ mod tests { .with_data_state_for_tests(data) .with_auth_wallets_for_tests(vec![sample_wallet("user-1", 30.0)]); let decision = decision_with_allowed_models(vec!["gpt-5".to_string()]); - let uri: Uri = "/v1/chat/completions".parse().expect("uri should parse"); - let body = Bytes::from_static( - br#"{"model":"gpt-5","messages":[{"role":"user","content":"hi"}],"max_tokens":100000}"#, + let plan = execution_plan( + json!({ + "model": "gpt-5", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 100_000 + }), + "openai:chat", ); + let report_context = billing_report_context(); - let first = - request_model_local_rejection(&state, Some(&decision), &uri, &json_headers(), &body) - .await - .expect("first auth capacity check should resolve"); + let first = execution_plan_balance_capacity_rejection( + &state, + &decision, + &plan, + Some(&report_context), + ) + .await + .expect("first auth capacity check should resolve"); assert_eq!(first, None); assert_eq!(quota_calls.load(Ordering::Acquire), 1); @@ -978,10 +1119,14 @@ mod tests { wallet.gift_balance = 0.0; } - let cached = - request_model_local_rejection(&state, Some(&decision), &uri, &json_headers(), &body) - .await - .expect("cached auth capacity check should resolve"); + let cached = execution_plan_balance_capacity_rejection( + &state, + &decision, + &plan, + Some(&report_context), + ) + .await + .expect("cached auth capacity check should resolve"); assert_eq!(cached, None); assert_eq!(quota_calls.load(Ordering::Acquire), 1); @@ -989,10 +1134,14 @@ mod tests { state.invalidate_auth_context_cache(); - let refreshed = - request_model_local_rejection(&state, Some(&decision), &uri, &json_headers(), &body) - .await - .expect("refreshed auth capacity check should resolve"); + let refreshed = execution_plan_balance_capacity_rejection( + &state, + &decision, + &plan, + Some(&report_context), + ) + .await + .expect("refreshed auth capacity check should resolve"); assert_eq!( refreshed, @@ -1004,7 +1153,7 @@ mod tests { } #[tokio::test] - async fn admin_bypass_limits_does_not_skip_exhausted_daily_quota_capacity() { + async fn admin_bypass_limits_does_not_skip_unbounded_zero_balance_capacity() { let context = billing_context_with_pricing( Some(json!({ "tiers": [{ @@ -1022,15 +1171,24 @@ mod tests { if let Some(auth_context) = decision.auth_context.as_mut() { auth_context.admin_bypass_limits = true; } - let uri: Uri = "/v1/chat/completions".parse().expect("uri should parse"); - let body = Bytes::from_static( - br#"{"model":"gpt-5","messages":[{"role":"user","content":"hi"}],"stream":true}"#, + let plan = execution_plan( + json!({ + "model": "gpt-5", + "messages": [{"role": "user", "content": "hi"}], + "stream": true + }), + "openai:chat", ); + let report_context = billing_report_context(); - let rejection = - request_model_local_rejection(&state, Some(&decision), &uri, &json_headers(), &body) - .await - .expect("quota rejection should resolve"); + let rejection = execution_plan_balance_capacity_rejection( + &state, + &decision, + &plan, + Some(&report_context), + ) + .await + .expect("quota rejection should resolve"); assert_eq!( rejection, @@ -1041,13 +1199,51 @@ mod tests { } #[tokio::test] - async fn positive_balance_still_denies_known_cost_above_available_capacity() { + async fn zero_balance_allows_a_proven_free_tier_execution_plan() { + let context = billing_context_with_pricing( + Some(json!({ + "tiers": [{ + "up_to": null, + "input_price_per_1m": 100.0, + "output_price_per_1m": 100.0 + }] + })), + None, + None, + Some("free_tier"), + ); + let state = state_with_quota_and_wallet(quota_availability(0.0, false), context); + let decision = decision_with_allowed_models(vec!["gpt-5".to_string()]); + let plan = execution_plan( + json!({ + "model": "gpt-5", + "messages": [{"role": "user", "content": "hi"}], + "max_completion_tokens": 1_000_000 + }), + "openai:chat", + ); + let report_context = billing_report_context(); + + let rejection = execution_plan_balance_capacity_rejection( + &state, + &decision, + &plan, + Some(&report_context), + ) + .await + .expect("free tier capacity check should resolve"); + + assert_eq!(rejection, None); + } + + #[tokio::test] + async fn finalized_chat_output_fields_and_choice_count_bound_capacity() { let context = billing_context_with_pricing( Some(json!({ "tiers": [{ "up_to": null, "input_price_per_1m": 0.0, - "output_price_per_1m": 60.0 + "output_price_per_1m": 20.0 }] })), None, @@ -1056,15 +1252,25 @@ mod tests { ); let state = state_with_quota_and_wallet(quota_availability(50.0, false), context); let decision = decision_with_allowed_models(vec!["gpt-5".to_string()]); - let uri: Uri = "/v1/chat/completions".parse().expect("uri should parse"); - let body = Bytes::from_static( - br#"{"model":"gpt-5","messages":[{"role":"user","content":"hi"}],"max_tokens":1000000}"#, + let plan = execution_plan( + json!({ + "model": "gpt-5", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 1, + "max_completion_tokens": 1_000_000, + "n": 3 + }), + "openai:chat", ); - let rejection = - request_model_local_rejection(&state, Some(&decision), &uri, &json_headers(), &body) - .await - .expect("quota rejection should resolve"); + let rejection = execution_plan_balance_capacity_rejection( + &state, + &decision, + &plan, + Some(&billing_report_context()), + ) + .await + .expect("quota rejection should resolve"); assert_eq!( rejection, @@ -1074,6 +1280,45 @@ mod tests { ); } + #[tokio::test] + async fn stateful_responses_request_skips_unprovable_capacity_rejection() { + let context = billing_context_with_pricing( + Some(json!({ + "tiers": [{ + "up_to": null, + "input_price_per_1m": 100.0, + "output_price_per_1m": 100.0 + }] + })), + None, + None, + None, + ); + let state = state_with_quota_and_wallet(quota_availability(0.01, false), context); + let decision = decision_with_allowed_models(vec!["gpt-5".to_string()]); + let plan = execution_plan( + json!({ + "model": "gpt-5", + "input": "continue", + "previous_response_id": "resp_123", + "max_output_tokens": 1_000_000 + }), + "openai:responses", + ); + let report_context = billing_report_context(); + + let rejection = execution_plan_balance_capacity_rejection( + &state, + &decision, + &plan, + Some(&report_context), + ) + .await + .expect("stateful request capacity check should resolve"); + + assert_eq!(rejection, None); + } + #[tokio::test] async fn wallet_overage_policy_extends_known_cost_capacity_when_enabled() { let context = billing_context_with_pricing( @@ -1090,15 +1335,23 @@ mod tests { ); let state = state_with_quota_and_wallet(quota_availability(50.0, true), context); let decision = decision_with_allowed_models(vec!["gpt-5".to_string()]); - let uri: Uri = "/v1/chat/completions".parse().expect("uri should parse"); - let body = Bytes::from_static( - br#"{"model":"gpt-5","messages":[{"role":"user","content":"hi"}],"max_tokens":1000000}"#, + let plan = execution_plan( + json!({ + "model": "gpt-5", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 1_000_000 + }), + "openai:chat", ); - let rejection = - request_model_local_rejection(&state, Some(&decision), &uri, &json_headers(), &body) - .await - .expect("quota rejection should resolve"); + let rejection = execution_plan_balance_capacity_rejection( + &state, + &decision, + &plan, + Some(&billing_report_context()), + ) + .await + .expect("quota rejection should resolve"); assert_eq!(rejection, None); } @@ -1119,10 +1372,10 @@ mod tests { ); let estimate = - estimate_cost_from_billing_context(&context, "openai:chat", 1_000_000, Some(1_000_000)) - .expect("estimate should resolve"); + estimate_from_billing_context(&context, "openai:chat", 1_000_000, Some(1_000_000)) + .expect("estimate should be bounded"); - assert_eq!(estimate, 18.0); + assert_eq!(estimate, 18.75); } #[test] @@ -1141,10 +1394,10 @@ mod tests { ); let estimate = - estimate_cost_from_billing_context(&context, "openai:chat", 1_000_000, Some(1_000_000)) - .expect("estimate should resolve"); + estimate_from_billing_context(&context, "openai:chat", 1_000_000, Some(1_000_000)) + .expect("estimate should be bounded"); - assert_eq!(estimate, 6.0); + assert_eq!(estimate, 6.5); } #[test] @@ -1163,9 +1416,89 @@ mod tests { ); let estimate = - estimate_cost_from_billing_context(&context, "openai:chat", 1_000_000, Some(1_000_000)) - .expect("estimate should resolve"); + estimate_from_billing_context(&context, "openai:chat", 1_000_000, Some(1_000_000)) + .expect("estimate should be bounded"); assert_eq!(estimate, 0.0); } + + #[test] + fn output_bound_uses_largest_supported_field_and_chat_choice_count() { + let body = json!({ + "max_tokens": 1, + "max_completion_tokens": 100_000, + "max_output_tokens": 50_000, + "n": 3 + }); + + assert_eq!(max_output_tokens_from_request(&body), Some(100_000)); + assert_eq!(output_choice_count_upper_bound("openai:chat", &body), 3); + assert_eq!( + output_choice_count_upper_bound("openai:responses", &body), + 1 + ); + } + + #[test] + fn authorization_cache_key_includes_effective_cache_ttl() { + let plan = execution_plan(json!({"model": "gpt-5.6-sol"}), "openai:responses"); + let without_ttl = execution_plan_cost_upper_bound_cache_key( + &plan, + Some("model-1"), + Some("gpt-5.6-sol"), + "openai:responses", + 100, + Some(10), + Some("priority"), + None, + ); + let with_ttl = execution_plan_cost_upper_bound_cache_key( + &plan, + Some("model-1"), + Some("gpt-5.6-sol"), + "openai:responses", + 100, + Some(10), + Some("priority"), + Some(30), + ); + + assert_ne!(without_ttl, with_ttl); + } + + #[test] + fn indirect_request_inputs_are_not_treated_as_body_bounded() { + let self_contained = json!({ + "input": [{ + "role": "user", + "content": [{"type": "input_text", "text": "hello"}] + }], + "tools": [{"type": "function", "name": "lookup", "parameters": {}}] + }); + assert!(openai_request_input_is_self_contained( + "openai:responses", + &self_contained + )); + + for indirect in [ + json!({"input": "continue", "previous_response_id": "resp_123"}), + json!({"input": "hello", "conversation": "conv_123"}), + json!({"prompt": {"id": "pmpt_123", "variables": {}}}), + json!({"input": [{"type": "item_reference", "id": "item_123"}]}), + json!({"input": [{"type": "input_file", "file_id": "file_123"}]}), + json!({"input": [{"type": "input_image", "image_url": "https://example.test/a.png"}]}), + json!({"input": [{"type": "url", "url": "https://example.test/document"}]}), + json!({"input": [{"file_uri": "https://example.test/file"}]}), + json!({"input": "search", "tools": [{"type": "file_search", "vector_store_ids": ["vs_123"]}]}), + ] { + assert!(!openai_request_input_is_self_contained( + "openai:responses", + &indirect + )); + } + assert!(!openai_request_input_is_self_contained( + "claude:messages", + &self_contained + )); + } } diff --git a/apps/aether-gateway/src/control/auth/mod.rs b/apps/aether-gateway/src/control/auth/mod.rs index e445064c7..412832acf 100644 --- a/apps/aether-gateway/src/control/auth/mod.rs +++ b/apps/aether-gateway/src/control/auth/mod.rs @@ -6,8 +6,8 @@ mod types; pub(crate) use credentials::extract_requested_model; pub(crate) use gate::{ - request_model_local_rejection, should_buffer_request_for_local_auth, - trusted_auth_local_rejection, GatewayLocalAuthRejection, + execution_plan_balance_capacity_rejection, request_model_local_rejection, + should_buffer_request_for_local_auth, trusted_auth_local_rejection, GatewayLocalAuthRejection, }; pub(crate) use resolution::{ refresh_execution_runtime_auth_context, resolve_execution_runtime_auth_context, diff --git a/apps/aether-gateway/src/control/auth/resolution.rs b/apps/aether-gateway/src/control/auth/resolution.rs index 9e250124f..1515b09e4 100644 --- a/apps/aether-gateway/src/control/auth/resolution.rs +++ b/apps/aether-gateway/src/control/auth/resolution.rs @@ -1115,10 +1115,9 @@ fn normalize_api_format_alias(value: &str) -> String { fn auth_gate_api_format(auth_endpoint_signature: &str) -> String { let normalized = normalize_api_format_alias(auth_endpoint_signature); - if normalized == "antigravity:v1internal" { - "gemini:generate_content".to_string() - } else { - normalized + match normalized.as_str() { + "antigravity:v1internal" => "gemini:generate_content".to_string(), + _ => normalized, } } diff --git a/apps/aether-gateway/src/control/mod.rs b/apps/aether-gateway/src/control/mod.rs index 80b9d2ea2..edf284777 100644 --- a/apps/aether-gateway/src/control/mod.rs +++ b/apps/aether-gateway/src/control/mod.rs @@ -8,7 +8,8 @@ mod public; mod route; pub(crate) use auth::{ - extract_requested_model, refresh_execution_runtime_auth_context, request_model_local_rejection, + execution_plan_balance_capacity_rejection, extract_requested_model, + 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, diff --git a/apps/aether-gateway/src/control/route/ai.rs b/apps/aether-gateway/src/control/route/ai.rs index ec8ba8d96..65a4781ef 100644 --- a/apps/aether-gateway/src/control/route/ai.rs +++ b/apps/aether-gateway/src/control/route/ai.rs @@ -54,6 +54,14 @@ pub(super) fn classify_ai_public_route( true, )) } + } else if method == http::Method::POST && normalized_path == "/v1/alpha/search" { + Some(classified( + "ai_public", + "openai", + "search", + "openai:search", + true, + )) } else if method == http::Method::POST && matches!( normalized_path, diff --git a/apps/aether-gateway/src/control/route/mod.rs b/apps/aether-gateway/src/control/route/mod.rs index 0e12394db..bd1c513a5 100644 --- a/apps/aether-gateway/src/control/route/mod.rs +++ b/apps/aether-gateway/src/control/route/mod.rs @@ -25,6 +25,7 @@ pub(crate) struct GatewayControlDecision { pub(crate) auth_context: Option, pub(crate) admin_principal: Option, pub(crate) local_auth_rejection: Option, + pub(crate) model_directive_policy: crate::system_features::ModelDirectivePolicySnapshot, } impl GatewayControlDecision { @@ -47,6 +48,7 @@ impl GatewayControlDecision { auth_context: None, admin_principal: None, local_auth_rejection: None, + model_directive_policy: Default::default(), } } @@ -131,6 +133,7 @@ impl ClassifiedRoute { auth_context: None, admin_principal: None, local_auth_rejection: None, + model_directive_policy: Default::default(), } } } @@ -146,6 +149,10 @@ pub(crate) async fn resolve_control_route( return Ok(None); }; decision.public_query_string = uri.query().map(ToOwned::to_owned); + if decision.route_class.as_deref() == Some("ai_public") { + decision.model_directive_policy = + crate::system_features::ModelDirectivePolicySnapshot::load(state).await; + } match resolve_control_decision_auth(state, headers, uri, trace_id, decision).await? { ControlDecisionAuthResolution::Resolved(decision) => Ok(Some(decision)), @@ -197,6 +204,15 @@ pub(super) fn detect_public_models_auth_signature(uri: &Uri, headers: &http::Hea return "gemini:generate_content".to_string(); } + let has_codex_client_version = uri.path() == "/v1/models" + && uri.query().is_some_and(|query| { + url::form_urlencoded::parse(query.as_bytes()) + .any(|(key, value)| key == "client_version" && !value.trim().is_empty()) + }); + if has_codex_client_version { + return "openai:responses".to_string(); + } + if uri.path().starts_with("/v1beta/models") { return "gemini:generate_content".to_string(); } diff --git a/apps/aether-gateway/src/control/tests/ai.rs b/apps/aether-gateway/src/control/tests/ai.rs index 87c0097b5..c645cb54c 100644 --- a/apps/aether-gateway/src/control/tests/ai.rs +++ b/apps/aether-gateway/src/control/tests/ai.rs @@ -56,6 +56,28 @@ fn classifies_openai_rerank_as_rerank_not_chat() { assert!(decision.is_execution_runtime_candidate()); } +#[test] +fn classifies_openai_search_as_its_own_sync_endpoint() { + let headers = headers(&[("authorization", "Bearer sk-test")]); + let uri: Uri = "/v1/alpha/search".parse().expect("uri should parse"); + let decision = + classify_control_route(&http::Method::POST, &uri, &headers).expect("route should classify"); + + assert_eq!(decision.route_family.as_deref(), Some("openai")); + assert_eq!(decision.route_kind.as_deref(), Some("search")); + assert_eq!( + decision.auth_endpoint_signature.as_deref(), + Some("openai:search") + ); + assert!(decision.is_execution_runtime_candidate()); + assert!(classify_control_route(&http::Method::GET, &uri, &headers).is_none()); + + let upstream_uri: Uri = "/backend-api/codex/alpha/search" + .parse() + .expect("uri should parse"); + assert!(classify_control_route(&http::Method::POST, &upstream_uri, &headers).is_none()); +} + #[test] fn classifies_openai_chat_and_responses_separately_from_embedding() { let headers = headers(&[("authorization", "Bearer sk-test")]); diff --git a/apps/aether-gateway/src/control/tests/public_support.rs b/apps/aether-gateway/src/control/tests/public_support.rs index 593fdada2..5f1ea8c6e 100644 --- a/apps/aether-gateway/src/control/tests/public_support.rs +++ b/apps/aether-gateway/src/control/tests/public_support.rs @@ -20,6 +20,39 @@ fn classifies_models_list_as_public_support_route() { assert!(!decision.is_execution_runtime_candidate()); } +#[test] +fn classifies_codex_models_list_with_responses_auth_signature() { + let headers = headers(&[("authorization", "Bearer sk-test")]); + let uri: Uri = "/v1/models?client_version=0.144.1" + .parse() + .expect("uri should parse"); + let decision = + classify_control_route(&http::Method::GET, &uri, &headers).expect("route should classify"); + + assert_eq!(decision.route_class.as_deref(), Some("public_support")); + assert_eq!(decision.route_family.as_deref(), Some("models")); + assert_eq!(decision.route_kind.as_deref(), Some("list")); + assert_eq!( + decision.auth_endpoint_signature.as_deref(), + Some("openai:responses") + ); +} + +#[test] +fn empty_codex_client_version_keeps_standard_openai_models_signature() { + let headers = headers(&[("authorization", "Bearer sk-test")]); + let uri: Uri = "/v1/models?client_version=" + .parse() + .expect("uri should parse"); + let decision = + classify_control_route(&http::Method::GET, &uri, &headers).expect("route should classify"); + + assert_eq!( + decision.auth_endpoint_signature.as_deref(), + Some("openai:chat") + ); +} + #[test] fn classifies_v1beta_models_as_gemini_public_support_route() { let headers = headers(&[]); diff --git a/apps/aether-gateway/src/data/state/auth.rs b/apps/aether-gateway/src/data/state/auth.rs index b80148708..184b72315 100644 --- a/apps/aether-gateway/src/data/state/auth.rs +++ b/apps/aether-gateway/src/data/state/auth.rs @@ -12,10 +12,7 @@ use super::{ StoredWalletSnapshot, UpdateManagementTokenRecord, UpsertOAuthProviderConfigRecord, }; use crate::LocalMutationOutcome; -use aether_data::repository::auth::{ - read_resolved_auth_api_key_snapshot_by_key_hash, - read_resolved_auth_api_key_snapshot_by_user_api_key_ids, -}; +use aether_data::repository::auth::ResolvedAuthApiKeySnapshotReader; #[derive(Debug, Clone, Default)] pub(crate) struct GatewayUserEffectiveListPolicies { @@ -1751,15 +1748,14 @@ impl GatewayDataState { let snapshot = crate::request_diagnostics::observe_db_operation( "auth_api_key_snapshot", self.database_pool_summary(), - read_resolved_auth_api_key_snapshot_by_user_api_key_ids( - self, + self.find_stored_auth_api_key_snapshot(AuthApiKeyLookupKey::UserApiKeyIds { user_id, api_key_id, - now_unix_secs, - ), + }), ) .await?; - self.apply_user_group_effective_policies(snapshot).await + self.apply_user_group_effective_policies(snapshot, now_unix_secs) + .await } pub(crate) async fn read_auth_api_key_snapshot_by_key_hash( @@ -1770,25 +1766,33 @@ impl GatewayDataState { let snapshot = crate::request_diagnostics::observe_db_operation( "auth_api_key_snapshot_by_hash", self.database_pool_summary(), - read_resolved_auth_api_key_snapshot_by_key_hash(self, key_hash, now_unix_secs), + self.find_stored_auth_api_key_snapshot(AuthApiKeyLookupKey::KeyHash(key_hash)), ) .await?; - self.apply_user_group_effective_policies(snapshot).await + self.apply_user_group_effective_policies(snapshot, now_unix_secs) + .await } async fn apply_user_group_effective_policies( &self, - snapshot: Option, + snapshot: Option, + now_unix_secs: u64, ) -> Result, DataLayerError> { let Some(mut snapshot) = snapshot else { return Ok(None); }; if snapshot.user_role.eq_ignore_ascii_case("admin") && !snapshot.api_key_is_standalone { apply_admin_unrestricted_auth_snapshot(&mut snapshot); - return Ok(Some(snapshot)); + return Ok(Some(GatewayAuthApiKeySnapshot::from_stored( + snapshot, + now_unix_secs, + ))); } let Some(repository) = self.user_reader.as_ref() else { - return Ok(Some(snapshot)); + return Ok(Some(GatewayAuthApiKeySnapshot::from_stored( + snapshot, + now_unix_secs, + ))); }; let Some(user) = crate::request_diagnostics::observe_db_operation( "auth_user_policy", @@ -1797,57 +1801,50 @@ impl GatewayDataState { ) .await? else { - return Ok(Some(snapshot)); + return Ok(Some(GatewayAuthApiKeySnapshot::from_stored( + snapshot, + now_unix_secs, + ))); }; if user.role.eq_ignore_ascii_case("admin") && !snapshot.api_key_is_standalone { snapshot.user_role = user.role; apply_admin_unrestricted_auth_snapshot(&mut snapshot); - return Ok(Some(snapshot)); + return Ok(Some(GatewayAuthApiKeySnapshot::from_stored( + snapshot, + now_unix_secs, + ))); } let groups = self .effective_user_groups_for_user(&snapshot.user_id) .await?; - let mut allowed_providers = + let allowed_providers = resolve_effective_list_policy(None, "unrestricted", &groups, |group| { ( &group.allowed_providers_mode, group.allowed_providers.clone(), ) }); - let mut allowed_api_formats = - resolve_effective_list_policy(None, "unrestricted", &groups, |group| { + let allowed_api_formats = + resolve_effective_api_format_policy(None, "unrestricted", &groups, |group| { ( &group.allowed_api_formats_mode, group.allowed_api_formats.clone(), ) }); - let mut allowed_models = + let allowed_models = resolve_effective_list_policy(None, "unrestricted", &groups, |group| { (&group.allowed_models_mode, group.allowed_models.clone()) }); let user_rate_limit = resolve_effective_rate_limit_policy(None, "system", &groups); - if !snapshot.api_key_is_standalone { - constrain_api_key_list_policy_to_user_policy( - &mut allowed_providers, - &mut snapshot.api_key_allowed_providers, - ); - constrain_api_key_list_policy_to_user_policy( - &mut allowed_api_formats, - &mut snapshot.api_key_allowed_api_formats, - ); - constrain_api_key_list_policy_to_user_policy( - &mut allowed_models, - &mut snapshot.api_key_allowed_models, - ); - } - snapshot.apply_user_policy( - allowed_providers, - allowed_api_formats, - allowed_models, - user_rate_limit, - ); - Ok(Some(snapshot)) + snapshot.user_allowed_providers = allowed_providers; + snapshot.user_allowed_api_formats = allowed_api_formats; + snapshot.user_allowed_models = allowed_models; + snapshot.user_rate_limit = user_rate_limit; + Ok(Some(GatewayAuthApiKeySnapshot::from_stored( + snapshot, + now_unix_secs, + ))) } pub(crate) async fn resolve_user_effective_list_policies( @@ -1875,7 +1872,7 @@ impl GatewayDataState { ) }, ), - allowed_api_formats: resolve_effective_list_policy( + allowed_api_formats: resolve_effective_api_format_policy( user.allowed_api_formats.clone(), &user.allowed_api_formats_mode, &groups, @@ -1970,7 +1967,7 @@ impl GatewayDataState { } } -fn apply_admin_unrestricted_auth_snapshot(snapshot: &mut GatewayAuthApiKeySnapshot) { +fn apply_admin_unrestricted_auth_snapshot(snapshot: &mut StoredAuthApiKeySnapshot) { snapshot.user_allowed_providers = None; snapshot.user_allowed_api_formats = None; snapshot.user_allowed_models = None; @@ -1995,6 +1992,19 @@ fn resolve_effective_list_policy( intersect_list_policies(group_policy, user_policy) } +fn resolve_effective_api_format_policy( + user_values: Option>, + user_mode: &str, + groups: &[aether_data::repository::users::StoredUserGroup], + group_field: impl Fn( + &aether_data::repository::users::StoredUserGroup, + ) -> (&str, Option>), +) -> Option> { + let group_policy = union_group_list_policies(groups, group_field); + let user_policy = list_restriction_from_mode(user_mode, user_values); + intersect_api_format_list_policies(group_policy, user_policy) +} + fn union_group_list_policies( groups: &[aether_data::repository::users::StoredUserGroup], group_field: impl Fn( @@ -2089,6 +2099,19 @@ fn intersect_list_policies( } } +fn intersect_api_format_list_policies( + left: Option>, + right: Option>, +) -> Option> { + match (left, right) { + (None, None) => None, + (Some(values), None) | (None, Some(values)) => Some(values), + (Some(left_values), Some(right_values)) => Some( + crate::ai_serving::intersect_api_format_allowed_lists(&left_values, &right_values), + ), + } +} + fn intersect_rate_limit_policies( left: Option, right: Option, @@ -2133,22 +2156,6 @@ fn rate_limit_policy_value(policy: Option) -> Option } } -fn constrain_api_key_list_policy_to_user_policy( - user_policy: &mut Option>, - api_key_policy: &mut Option>, -) { - let Some(api_key_values) = api_key_policy.as_ref().filter(|values| !values.is_empty()) else { - return; - }; - let Some(user_values) = user_policy.clone() else { - return; - }; - let effective = intersect_list_policies(Some(api_key_values.to_vec()), Some(user_values)) - .unwrap_or_default(); - *user_policy = Some(effective.clone()); - *api_key_policy = Some(effective); -} - #[cfg(test)] mod tests { use std::sync::Arc; @@ -2279,6 +2286,41 @@ mod tests { ); } + #[test] + fn api_format_policy_intersection_preserves_search_companion_scope() { + let mut responses_group = + sample_group("responses", 10, None, "unrestricted", None, "system"); + responses_group.allowed_api_formats = Some(vec!["openai:responses".to_string()]); + responses_group.allowed_api_formats_mode = "specific".to_string(); + + let search_policy = resolve_effective_api_format_policy( + Some(vec!["openai:search".to_string()]), + "specific", + std::slice::from_ref(&responses_group), + |group| { + ( + &group.allowed_api_formats_mode, + group.allowed_api_formats.clone(), + ) + }, + ); + assert_eq!(search_policy, Some(vec!["openai:search".to_string()])); + + responses_group.allowed_api_formats = Some(vec!["openai:search".to_string()]); + let responses_policy = resolve_effective_api_format_policy( + Some(vec!["openai:responses".to_string()]), + "specific", + &[responses_group], + |group| { + ( + &group.allowed_api_formats_mode, + group.allowed_api_formats.clone(), + ) + }, + ); + assert_eq!(responses_policy, Some(vec!["openai:search".to_string()])); + } + #[test] fn list_policy_unions_multiple_group_restrictions_legacy_case() { let groups = vec![ @@ -2452,17 +2494,6 @@ mod tests { ); } - #[test] - fn api_key_specific_policy_cannot_expand_user_policy() { - let mut user_policy = Some(vec!["gpt-5".to_string()]); - let mut api_key_policy = Some(vec!["gpt-4.1".to_string()]); - - constrain_api_key_list_policy_to_user_policy(&mut user_policy, &mut api_key_policy); - - assert_eq!(user_policy, Some(Vec::::new())); - assert_eq!(api_key_policy, Some(Vec::::new())); - } - #[tokio::test] async fn admin_non_standalone_snapshot_bypasses_group_and_key_policies() { let mut snapshot = sample_snapshot_with_role("key-admin", "admin-1", "admin") @@ -2518,6 +2549,38 @@ mod tests { assert_eq!(resolved.api_key_concurrent_limit, None); } + #[tokio::test] + async fn current_admin_role_bypasses_stored_user_and_key_policies() { + let mut snapshot = sample_snapshot("key-admin", "admin-1"); + snapshot.api_key_allowed_providers = Some(vec!["anthropic".to_string()]); + snapshot.api_key_allowed_api_formats = Some(vec!["anthropic:messages".to_string()]); + snapshot.api_key_allowed_models = Some(vec!["claude-sonnet-4-5".to_string()]); + + let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( + Some("hash-admin".to_string()), + snapshot, + )])); + let user_repository = Arc::new(InMemoryUserReadRepository::seed_auth_users(vec![ + sample_auth_user("admin-1", "admin"), + ])); + let state = GatewayDataState::with_auth_api_key_reader_for_tests(auth_repository) + .with_user_reader(user_repository); + + let resolved = state + .read_auth_api_key_snapshot_by_key_hash("hash-admin", 100) + .await + .expect("snapshot should resolve") + .expect("snapshot should exist"); + + assert_eq!(resolved.user_role, "admin"); + assert_eq!(resolved.effective_allowed_providers(), None); + assert_eq!(resolved.effective_allowed_api_formats(), None); + assert_eq!(resolved.effective_allowed_models(), None); + assert_eq!(resolved.user_rate_limit, None); + assert_eq!(resolved.api_key_rate_limit, None); + assert_eq!(resolved.api_key_concurrent_limit, None); + } + #[tokio::test] async fn user_personal_policy_fields_are_ignored_when_groups_are_applied() { let mut snapshot = sample_snapshot("key-user", "user-1").with_user_rate_limit(Some(200)); @@ -2577,6 +2640,93 @@ mod tests { assert_eq!(resolved.user_rate_limit, Some(30)); } + #[tokio::test] + async fn group_responses_permission_and_key_search_scope_resolve_to_search() { + let mut snapshot = sample_snapshot("key-search", "user-search"); + snapshot.api_key_allowed_api_formats = Some(vec!["openai:search".to_string()]); + let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( + Some("hash-search".to_string()), + snapshot, + )])); + let user_repository = Arc::new(InMemoryUserReadRepository::seed_auth_users(vec![ + sample_auth_user("user-search", "user"), + ])); + let group = user_repository + .create_user_group(UpsertUserGroupRecord { + name: "Responses".to_string(), + description: None, + priority: 10, + allowed_providers: None, + allowed_providers_mode: "unrestricted".to_string(), + allowed_api_formats: Some(vec!["openai:responses".to_string()]), + allowed_api_formats_mode: "specific".to_string(), + allowed_models: None, + allowed_models_mode: "unrestricted".to_string(), + rate_limit: None, + rate_limit_mode: "system".to_string(), + }) + .await + .expect("group should create") + .expect("group should exist"); + user_repository + .add_user_to_group(&group.id, "user-search") + .await + .expect("group membership should create"); + + let state = GatewayDataState::with_auth_api_key_reader_for_tests(auth_repository) + .with_user_reader(user_repository); + let resolved = state + .read_auth_api_key_snapshot_by_key_hash("hash-search", 100) + .await + .expect("snapshot should resolve") + .expect("snapshot should exist"); + + assert_eq!( + resolved.effective_allowed_api_formats(), + Some(&["openai:search".to_string()][..]) + ); + } + + #[tokio::test] + async fn snapshot_without_user_reader_uses_stored_policy_intersection() { + let mut snapshot = sample_snapshot("key-search", "user-search"); + snapshot.api_key_allowed_api_formats = Some(vec!["openai:search".to_string()]); + let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( + Some("hash-search".to_string()), + snapshot, + )])); + let state = GatewayDataState::with_auth_api_key_reader_for_tests(auth_repository); + + let resolved = state + .read_auth_api_key_snapshot_by_key_hash("hash-search", 100) + .await + .expect("snapshot should resolve") + .expect("snapshot should exist"); + + assert_eq!(resolved.effective_allowed_api_formats(), Some(&[][..])); + } + + #[tokio::test] + async fn missing_current_user_uses_stored_policy_intersection() { + let mut snapshot = sample_snapshot("key-search", "missing-user"); + snapshot.api_key_allowed_api_formats = Some(vec!["openai:search".to_string()]); + let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( + Some("hash-search".to_string()), + snapshot, + )])); + let user_repository = Arc::new(InMemoryUserReadRepository::default()); + let state = GatewayDataState::with_auth_api_key_reader_for_tests(auth_repository) + .with_user_reader(user_repository); + + let resolved = state + .read_auth_api_key_snapshot_by_key_hash("hash-search", 100) + .await + .expect("snapshot should resolve") + .expect("snapshot should exist"); + + assert_eq!(resolved.effective_allowed_api_formats(), Some(&[][..])); + } + #[tokio::test] async fn data_state_lists_auth_api_key_export_records() { let repository = Arc::new( diff --git a/apps/aether-gateway/src/data/state/catalog.rs b/apps/aether-gateway/src/data/state/catalog.rs index 064d66fc5..183b9767d 100644 --- a/apps/aether-gateway/src/data/state/catalog.rs +++ b/apps/aether-gateway/src/data/state/catalog.rs @@ -532,6 +532,88 @@ impl GatewayDataState { Ok(updated) } + pub(crate) async fn upsert_provider_catalog_key_upstream_metadata_namespace( + &self, + key_id: &str, + namespace: &str, + value: &serde_json::Value, + updated_at_unix_secs: Option, + ) -> Result { + let updated = match &self.provider_catalog_writer { + Some(repository) => { + repository + .upsert_key_upstream_metadata_namespace( + key_id, + namespace, + value, + updated_at_unix_secs, + ) + .await + } + None => Ok(false), + }?; + if updated { + self.clear_provider_catalog_cache(); + } + Ok(updated) + } + + pub(crate) async fn update_provider_catalog_key_model_fetch_state( + &self, + key_id: &str, + allowed_models: Option<&serde_json::Value>, + last_models_fetch_at_unix_secs: Option, + last_models_fetch_error: Option<&str>, + updated_at_unix_secs: Option, + ) -> Result { + let updated = match &self.provider_catalog_writer { + Some(repository) => { + repository + .update_key_model_fetch_state( + key_id, + allowed_models, + last_models_fetch_at_unix_secs, + last_models_fetch_error, + updated_at_unix_secs, + ) + .await + } + None => Ok(false), + }?; + if updated { + self.clear_provider_catalog_cache(); + } + Ok(updated) + } + + pub(crate) async fn update_provider_catalog_key_model_fetch_success( + &self, + key_id: &str, + allowed_models: Option<&serde_json::Value>, + last_models_fetch_at_unix_secs: u64, + upstream_metadata_updates: &[aether_data_contracts::repository::provider_catalog::ProviderCatalogUpstreamMetadataNamespaceUpdate], + updated_at_unix_secs: Option, + ) -> Result { + let updated = match &self.provider_catalog_writer { + Some(repository) => { + repository + .update_key_model_fetch_success( + key_id, + allowed_models, + last_models_fetch_at_unix_secs, + upstream_metadata_updates, + updated_at_unix_secs, + ) + .await + } + None => Ok(false), + }?; + if updated { + self.clear_provider_catalog_cache(); + } + Ok(updated) + } + pub(crate) async fn delete_provider_catalog_key( &self, key_id: &str, diff --git a/apps/aether-gateway/src/dispatch/pool_scheduler.rs b/apps/aether-gateway/src/dispatch/pool_scheduler.rs index 11dbc7863..5c53d4f6c 100644 --- a/apps/aether-gateway/src/dispatch/pool_scheduler.rs +++ b/apps/aether-gateway/src/dispatch/pool_scheduler.rs @@ -4172,6 +4172,7 @@ mod tests { global_model_id: "global-model-1".to_string(), global_model_name: "gpt-5".to_string(), selected_provider_model_name: "gpt-5".to_string(), + supports_streaming: true, mapping_matched_model: None, }, provider_api_format: "openai:responses".to_string(), @@ -4285,6 +4286,7 @@ mod tests { global_model_id: "global-model-1".to_string(), global_model_name: "gpt-5".to_string(), selected_provider_model_name: "gpt-5".to_string(), + supports_streaming: true, mapping_matched_model: None, }, provider_api_format: "openai:chat".to_string(), diff --git a/apps/aether-gateway/src/dispatch/refs.rs b/apps/aether-gateway/src/dispatch/refs.rs index 48f847e45..e8acb0a1d 100644 --- a/apps/aether-gateway/src/dispatch/refs.rs +++ b/apps/aether-gateway/src/dispatch/refs.rs @@ -146,6 +146,7 @@ mod tests { global_model_id: "global-model-1".to_string(), global_model_name: "gpt-5".to_string(), selected_provider_model_name: "gpt-5".to_string(), + supports_streaming: true, mapping_matched_model: None, }, transport: Arc::new(crate::ai_serving::GatewayProviderTransportSnapshot { diff --git a/apps/aether-gateway/src/execution_runtime/fallback.rs b/apps/aether-gateway/src/execution_runtime/fallback.rs index a50e6ef30..bc4b811de 100644 --- a/apps/aether-gateway/src/execution_runtime/fallback.rs +++ b/apps/aether-gateway/src/execution_runtime/fallback.rs @@ -941,6 +941,7 @@ mod tests { success_failover_patterns: Vec::new(), error_stop_patterns: Vec::new(), stop_cyber_policy_errors: false, + retry_client_errors_by_default: true, } ); } diff --git a/apps/aether-gateway/src/execution_runtime/grok.rs b/apps/aether-gateway/src/execution_runtime/grok.rs index 937d2d248..5691bdb8b 100644 --- a/apps/aether-gateway/src/execution_runtime/grok.rs +++ b/apps/aether-gateway/src/execution_runtime/grok.rs @@ -1042,6 +1042,7 @@ fn grok_stream_terminal_summary( finish_reason: Some("stop".to_string()), response_id: None, model: plan.model_name.clone(), + provider_actual_service_tier: None, observed_finish: true, unknown_event_count: 0, parser_error: None, diff --git a/apps/aether-gateway/src/execution_runtime/server.rs b/apps/aether-gateway/src/execution_runtime/server.rs index 712b7e0c8..f2a8bf8c0 100644 --- a/apps/aether-gateway/src/execution_runtime/server.rs +++ b/apps/aether-gateway/src/execution_runtime/server.rs @@ -359,8 +359,7 @@ impl IntoResponse for ExecutionRuntimeAppError { return build_overloaded_response(&self.0.to_string()); } ExecutionRuntimeServerError::Transport( - ExecutionRuntimeTransportError::StreamUnsupported - | ExecutionRuntimeTransportError::RequestBodyRequired + ExecutionRuntimeTransportError::RequestBodyRequired | ExecutionRuntimeTransportError::BodyDecode(_) | ExecutionRuntimeTransportError::UnsupportedContentEncoding(_) | ExecutionRuntimeTransportError::ProxyUnsupported @@ -397,7 +396,9 @@ mod tests { build_execution_runtime_router_with_request_concurrency_limit, build_execution_runtime_router_with_request_gates, DISTRIBUTED_REQUEST_GATE_NAME, }; - use aether_contracts::{ExecutionPlan, ExecutionTimeouts, RequestBody}; + use aether_contracts::{ + ExecutionPlan, ExecutionTimeouts, RequestBody, StreamFrame, StreamFrameType, + }; use aether_runtime_state::{ MemoryRuntimeStateConfig, RuntimeSemaphore, RuntimeSemaphoreConfig, RuntimeState, }; @@ -459,6 +460,43 @@ mod tests { } } + #[tokio::test] + async fn execution_runtime_stream_endpoint_carries_non_stream_upstream_plan() { + let upstream = Router::new().route( + "/sync-json", + any(|| async { axum::Json(serde_json::json!({"ok": true})) }), + ); + let (upstream_url, upstream_handle) = start_server(upstream).await; + let runtime = build_execution_runtime_router_with_request_concurrency_limit(None); + let (runtime_url, runtime_handle) = start_server(runtime).await; + let mut plan = stream_plan(format!("{upstream_url}/sync-json")); + plan.stream = false; + + let response = reqwest::Client::new() + .post(format!("{runtime_url}/v1/execute/stream")) + .json(&plan) + .send() + .await + .expect("execution request should succeed"); + + assert_eq!(response.status(), StatusCode::OK); + let body = response.text().await.expect("frame body should read"); + let frame_types = body + .lines() + .map(|line| { + serde_json::from_str::(line) + .expect("execution runtime frame should decode") + .frame_type + }) + .collect::>(); + assert!(frame_types.contains(&StreamFrameType::Headers)); + assert!(frame_types.contains(&StreamFrameType::Data)); + assert!(frame_types.contains(&StreamFrameType::Eof)); + + runtime_handle.abort(); + upstream_handle.abort(); + } + #[tokio::test] async fn execution_runtime_rejects_second_in_flight_stream_request_with_overload() { let upstream_hits = Arc::new(AtomicUsize::new(0)); diff --git a/apps/aether-gateway/src/execution_runtime/stream/execution.rs b/apps/aether-gateway/src/execution_runtime/stream/execution.rs index b8b106906..890488de1 100644 --- a/apps/aether-gateway/src/execution_runtime/stream/execution.rs +++ b/apps/aether-gateway/src/execution_runtime/stream/execution.rs @@ -851,6 +851,9 @@ fn merge_stream_terminal_summary( if current_summary.model.is_none() { current_summary.model = observed.model; } + if observed.provider_actual_service_tier.is_some() { + current_summary.provider_actual_service_tier = observed.provider_actual_service_tier; + } current_summary.observed_finish |= observed.observed_finish; current_summary.unknown_event_count = current_summary .unknown_event_count @@ -5484,9 +5487,8 @@ async fn execute_stream_from_frame_stream( telemetry = Some(frame_telemetry); } StreamFramePayload::Eof { summary } => { - if summary.is_some() { - stream_terminal_summary = summary; - } + stream_terminal_summary = + merge_stream_terminal_summary(stream_terminal_summary.take(), summary); break; } StreamFramePayload::Error { error } => { @@ -6293,12 +6295,14 @@ mod tests { Some(ExecutionStreamTerminalSummary { standardized_usage: Some(runtime_usage), model: Some("gpt-5.5".to_string()), + provider_actual_service_tier: Some("priority".to_string()), unknown_event_count: 1, ..ExecutionStreamTerminalSummary::default() }), Some(ExecutionStreamTerminalSummary { standardized_usage: Some(observed_usage), response_id: Some("resp_123".to_string()), + provider_actual_service_tier: Some("default".to_string()), observed_finish: true, unknown_event_count: 2, ..ExecutionStreamTerminalSummary::default() @@ -6313,6 +6317,10 @@ mod tests { assert_eq!(usage.output_tokens, 137); assert_eq!(merged.model.as_deref(), Some("gpt-5.5")); assert_eq!(merged.response_id.as_deref(), Some("resp_123")); + assert_eq!( + merged.provider_actual_service_tier.as_deref(), + Some("default") + ); assert!(merged.observed_finish); assert_eq!(merged.unknown_event_count, 3); } diff --git a/apps/aether-gateway/src/execution_runtime/stream_pump.rs b/apps/aether-gateway/src/execution_runtime/stream_pump.rs index 1d71263b8..bfcc90889 100644 --- a/apps/aether-gateway/src/execution_runtime/stream_pump.rs +++ b/apps/aether-gateway/src/execution_runtime/stream_pump.rs @@ -676,6 +676,14 @@ fn should_buffer_non_stream_response( return false; } + if report_context + .get("upstream_is_stream") + .and_then(Value::as_bool) + == Some(false) + { + return true; + } + headers .get("content-length") .and_then(|value| value.trim().parse::().ok()) @@ -1134,22 +1142,36 @@ mod tests { } #[test] - fn buffers_non_sse_response_only_when_content_length_is_known() { - let report_context = serde_json::json!({ + fn buffers_declared_non_stream_responses_without_relying_on_content_length() { + let streaming_context = serde_json::json!({ "provider_api_format": "openai:chat", "client_api_format": "openai:chat", + "upstream_is_stream": true, + }); + let non_stream_context = serde_json::json!({ + "provider_api_format": "openai:image", + "client_api_format": "openai:responses", + "upstream_is_stream": false, }); assert!(!should_buffer_non_stream_response( &BTreeMap::from([("content-type".into(), "application/json".into())]), - &report_context + &streaming_context + )); + assert!(should_buffer_non_stream_response( + &BTreeMap::from([("content-type".into(), "application/json".into())]), + &non_stream_context )); assert!(should_buffer_non_stream_response( &BTreeMap::from([ ("content-type".into(), "application/json".into()), ("content-length".into(), "128".into()), ]), - &report_context + &streaming_context + )); + assert!(!should_buffer_non_stream_response( + &BTreeMap::from([("content-type".into(), "text/event-stream".into())]), + &non_stream_context )); } @@ -1567,33 +1589,48 @@ mod tests { .await .expect("listener should bind"); let addr = listener.local_addr().expect("local addr should resolve"); + let generated_image = "a".repeat(32 * 1024); + let expected_image = generated_image.clone(); let server = tokio::spawn(async move { let app = Router::new().route( - "/responses", - post(|| async { - let body = serde_json::json!({ - "created": 1776971267_u64, - "data": [{ - "b64_json": "aGVsbG8=" - }], - "usage": { - "total_tokens": 100, - "input_tokens": 50, - "output_tokens": 50, - "input_tokens_details": { - "text_tokens": 10, - "image_tokens": 40 + "/images/generations", + post(move || { + let generated_image = generated_image.clone(); + async move { + let body = serde_json::json!({ + "created": 1776971267_u64, + "data": [{ + "b64_json": generated_image + }], + "usage": { + "total_tokens": 100, + "input_tokens": 50, + "output_tokens": 50, + "input_tokens_details": { + "text_tokens": 10, + "image_tokens": 40 + } } - } - }); - let mut response = axum::http::Response::new(Body::from( - serde_json::to_vec(&body).expect("json should encode"), - )); - response.headers_mut().insert( - header::CONTENT_TYPE, - HeaderValue::from_static("application/json"), - ); - response + }); + let encoded = serde_json::to_vec(&body).expect("json should encode"); + let chunks = encoded + .chunks(4096) + .map(Bytes::copy_from_slice) + .collect::>(); + let chunked_body = stream! { + for chunk in chunks { + yield Ok::(chunk); + tokio::task::yield_now().await; + } + }; + let mut response = + axum::http::Response::new(Body::from_stream(chunked_body)); + response.headers_mut().insert( + header::CONTENT_TYPE, + HeaderValue::from_static("application/json"), + ); + response + } }), ); axum::serve(listener, app) @@ -1611,16 +1648,15 @@ mod tests { endpoint_id: "endpoint-1".to_string(), key_id: "key-1".to_string(), method: "POST".to_string(), - url: format!("http://{addr}/responses"), + url: format!("http://{addr}/images/generations"), headers: BTreeMap::new(), content_type: None, content_encoding: None, body: RequestBody::from_json(serde_json::json!({ "model": "gpt-image-1", - "prompt": "poster", - "stream": true + "prompt": "poster" })), - stream: true, + stream: false, client_api_format: "openai:image".to_string(), provider_api_format: "openai:image".to_string(), model_name: Some("gpt-image-1".into()), @@ -1673,7 +1709,8 @@ mod tests { let bridged_text = String::from_utf8(bridged_body).expect("bridged body should be utf8"); assert!(bridged_text.contains("event: image_generation.completed")); assert!(bridged_text.contains("\"type\":\"image_generation.completed\"")); - assert!(bridged_text.contains("\"b64_json\":\"aGVsbG8=\"")); + assert!(bridged_text.contains(&format!("\"b64_json\":\"{expected_image}\""))); + assert!(bridged_text.len() > 32 * 1024); assert!(bridged_text.contains("\"total_tokens\":100")); let eof_frame = frames diff --git a/apps/aether-gateway/src/execution_runtime/tests.rs b/apps/aether-gateway/src/execution_runtime/tests.rs index 7a0016cc9..5a862f4ce 100644 --- a/apps/aether-gateway/src/execution_runtime/tests.rs +++ b/apps/aether-gateway/src/execution_runtime/tests.rs @@ -43,6 +43,7 @@ fn missing_exact_provider_request_payload(decision_kind: &str) -> AiExecutionDec request_id: Some("req_123".to_string()), candidate_id: Some("cand_123".to_string()), provider_name: Some("provider".to_string()), + provider_type: None, provider_id: Some("provider_id".to_string()), endpoint_id: Some("endpoint_id".to_string()), key_id: Some("key_id".to_string()), diff --git a/apps/aether-gateway/src/execution_runtime/transport.rs b/apps/aether-gateway/src/execution_runtime/transport.rs index ab4e1b36e..231e33923 100644 --- a/apps/aether-gateway/src/execution_runtime/transport.rs +++ b/apps/aether-gateway/src/execution_runtime/transport.rs @@ -57,8 +57,8 @@ const TUNNEL_RELAY_PATH_PREFIX: &str = "/api/internal/tunnel/relay"; const DEFAULT_TUNNEL_TIMEOUT_MS: u64 = 60_000; const DEFAULT_STREAM_FIRST_BYTE_TIMEOUT_MS: u64 = 30_000; const DEFAULT_NON_STREAM_TOTAL_TIMEOUT_MS: u64 = 300_000; +const DEFAULT_CODEX_COMPACT_TOTAL_TIMEOUT_MS: u64 = 1_200_000; const MIN_TUNNEL_TIMEOUT_SECS: u64 = 1; -const MAX_TUNNEL_TIMEOUT_SECS: u64 = 300; const DIRECT_REQWEST_H2_CLIENT_SHARDS_ENV: &str = "AETHER_GATEWAY_DIRECT_REQWEST_H2_CLIENT_SHARDS"; const DIRECT_REQWEST_CLIENT_SHARDS_ENV: &str = "AETHER_GATEWAY_DIRECT_REQWEST_CLIENT_SHARDS"; const DIRECT_REQWEST_H2_TARGET_STREAMS_PER_CLIENT_ENV: &str = @@ -512,8 +512,6 @@ pub(crate) fn format_hyper_error_chain(err: &dyn std::error::Error) -> String { #[derive(Debug, Error)] pub(crate) enum ExecutionRuntimeTransportError { - #[error("stream execution is not supported for this plan")] - StreamUnsupported, #[error("request body must contain json_body or body_bytes_b64")] RequestBodyRequired, #[error("request body base64 is invalid: {0}")] @@ -681,10 +679,6 @@ impl DirectSyncExecutionRuntime { &self, plan: &ExecutionPlan, ) -> Result { - if !plan.stream { - return Err(ExecutionRuntimeTransportError::StreamUnsupported); - } - let build_body_started_at = Instant::now(); let body_bytes = build_request_body(plan)?; observe_gateway_stage_ms( @@ -835,6 +829,7 @@ fn build_stream_summary_report_context(plan: &ExecutionPlan) -> Value { "provider_api_format": plan.provider_api_format, "client_api_format": plan.client_api_format, "model": plan.model_name, + "upstream_is_stream": plan.stream, }) } @@ -2125,20 +2120,17 @@ pub(crate) fn build_request_body( Vec::new() }; - if should_gzip_request_body(plan) && plan.body.json_body.is_some() { - body_bytes = gzip_bytes(&body_bytes)?; + if plan.body.json_body.is_some() { + body_bytes = match normalize_content_encoding(plan.content_encoding.as_deref()).as_deref() { + Some("gzip") => gzip_bytes(&body_bytes)?, + Some("zstd") => zstd_bytes(&body_bytes)?, + _ => body_bytes, + }; } Ok(body_bytes) } -fn should_gzip_request_body(plan: &ExecutionPlan) -> bool { - matches!( - normalize_content_encoding(plan.content_encoding.as_deref()).as_deref(), - Some("gzip") - ) -} - fn normalize_content_encoding(value: Option<&str>) -> Option { value .map(str::trim) @@ -2156,6 +2148,11 @@ fn gzip_bytes(body_bytes: &[u8]) -> Result, ExecutionRuntimeTransportErr .map_err(|err| ExecutionRuntimeTransportError::RelayError(err.to_string())) } +fn zstd_bytes(body_bytes: &[u8]) -> Result, ExecutionRuntimeTransportError> { + zstd::stream::encode_all(std::io::Cursor::new(body_bytes), 3) + .map_err(|err| ExecutionRuntimeTransportError::RelayError(err.to_string())) +} + fn build_relay_client( timeouts: Option<&aether_contracts::ExecutionTimeouts>, ) -> Result { @@ -2218,28 +2215,49 @@ fn resolve_tunnel_first_byte_timeout(plan: &ExecutionPlan) -> Option { }) } -fn resolve_non_stream_total_timeout(plan: &ExecutionPlan) -> Option { - if plan.stream { +pub(crate) fn resolve_non_stream_total_timeout_for_request( + is_stream: bool, + provider_api_format: &str, + timeouts: Option<&aether_contracts::ExecutionTimeouts>, +) -> Option { + if is_stream { return None; } - let timeout_ms = plan - .timeouts - .as_ref() + let default_timeout_ms = + if crate::ai_serving::is_openai_responses_compact_format(provider_api_format) { + DEFAULT_CODEX_COMPACT_TOTAL_TIMEOUT_MS + } else { + DEFAULT_NON_STREAM_TOTAL_TIMEOUT_MS + }; + let timeout_ms = timeouts .and_then(|timeouts| timeouts.total_ms) - .unwrap_or(DEFAULT_NON_STREAM_TOTAL_TIMEOUT_MS); + .unwrap_or(default_timeout_ms); + Some(Duration::from_millis(timeout_ms.max(1))) +} + +fn resolve_non_stream_total_timeout(plan: &ExecutionPlan) -> Option { + resolve_non_stream_total_timeout_for_request( + plan.stream, + &plan.provider_api_format, + plan.timeouts.as_ref(), + ) +} + +pub(crate) fn resolve_stream_first_byte_timeout_for_request( + is_stream: bool, + timeouts: Option<&aether_contracts::ExecutionTimeouts>, +) -> Option { + if !is_stream { + return None; + } + let timeout_ms = timeouts + .and_then(|timeouts| timeouts.first_byte_ms) + .unwrap_or(DEFAULT_STREAM_FIRST_BYTE_TIMEOUT_MS); Some(Duration::from_millis(timeout_ms.max(1))) } pub(crate) fn resolve_stream_first_byte_timeout(plan: &ExecutionPlan) -> Option { - if !plan.stream { - return None; - } - let timeout_ms = plan - .timeouts - .as_ref() - .and_then(|timeouts| timeouts.first_byte_ms) - .unwrap_or(DEFAULT_STREAM_FIRST_BYTE_TIMEOUT_MS); - Some(Duration::from_millis(timeout_ms.max(1))) + resolve_stream_first_byte_timeout_for_request(plan.stream, plan.timeouts.as_ref()) } pub(crate) async fn with_non_stream_total_timeout( @@ -2359,7 +2377,10 @@ fn resolve_tunnel_timeout_metadata(plan: &ExecutionPlan) -> TunnelTimeoutMetadat fn timeout_ms_to_secs(ms: u64) -> u64 { let secs = ms.div_ceil(1_000); - secs.clamp(MIN_TUNNEL_TIMEOUT_SECS, MAX_TUNNEL_TIMEOUT_SECS) + secs.clamp( + MIN_TUNNEL_TIMEOUT_SECS, + aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_SECS, + ) } fn resolve_tunnel_node_id(proxy: Option<&ProxySnapshot>) -> Option { @@ -3552,7 +3573,7 @@ pub(crate) fn build_request_headers( let mut out = HeaderMap::new(); let normalized_content_encoding = normalize_content_encoding(content_encoding); if let Some(encoding) = normalized_content_encoding.as_deref() { - if encoding != "gzip" && !allow_passthrough_content_encoding { + if !matches!(encoding, "gzip" | "zstd") && !allow_passthrough_content_encoding { return Err(ExecutionRuntimeTransportError::UnsupportedContentEncoding( encoding.to_string(), )); @@ -4637,6 +4658,25 @@ mod tests { assert_eq!(timeout, std::time::Duration::from_secs(300)); } + #[test] + fn codex_compact_uses_the_full_unary_timeout_by_default() { + let mut plan = tunnel_timeout_plan(false); + plan.provider_api_format = "openai:responses:compact".to_string(); + plan.timeouts = None; + + let timeout = resolve_non_stream_total_timeout(&plan) + .expect("Codex Compact should have a total timeout"); + let meta = build_direct_tunnel_request_meta( + &plan, + &reqwest::header::HeaderMap::new(), + ExecutionTransportControls::default(), + ); + + assert_eq!(timeout, std::time::Duration::from_secs(1_200)); + assert_eq!(meta.request_timeout_ms, Some(1_200_000)); + assert_eq!(meta.timeout, 1_200); + } + #[test] fn tunnel_request_meta_uses_non_stream_default_instead_of_first_byte_default() { let mut plan = tunnel_timeout_plan(false); @@ -6323,13 +6363,21 @@ mod tests { .and_then(|value| value.to_str().ok()) .unwrap_or_default() .to_string(); - let mut decoder = flate2::read::GzDecoder::new(body.as_ref()); - let mut decoded = String::new(); - decoder - .read_to_string(&mut decoded) - .expect("gzip body should decode"); + let decoded = match header_encoding.as_str() { + "gzip" => { + let mut decoder = flate2::read::GzDecoder::new(body.as_ref()); + let mut decoded = Vec::new(); + decoder + .read_to_end(&mut decoded) + .expect("gzip body should decode"); + decoded + } + "zstd" => zstd::stream::decode_all(std::io::Cursor::new(body.as_ref())) + .expect("zstd body should decode"), + encoding => panic!("unexpected content encoding: {encoding}"), + }; let decoded_json: serde_json::Value = - serde_json::from_str(&decoded).expect("decoded json should parse"); + serde_json::from_slice(&decoded).expect("decoded json should parse"); ( axum::http::StatusCode::OK, Json(json!({ @@ -6346,45 +6394,47 @@ mod tests { }); let execution_runtime = DirectSyncExecutionRuntime::new(); - let result = execution_runtime - .execute_sync(&ExecutionPlan { - request_id: "req-gzip-1".into(), - candidate_id: Some("cand-1".into()), - provider_name: Some("openai".into()), - provider_id: "prov-1".into(), - endpoint_id: "ep-1".into(), - key_id: "key-1".into(), - method: "POST".into(), - url: format!("http://{addr}/chat"), - headers: BTreeMap::from([("content-type".into(), "application/json".into())]), - content_type: Some("application/json".into()), - content_encoding: Some("gzip".into()), - body: RequestBody::from_json(json!({"model": "gpt-4.1"})), - stream: false, - client_api_format: "openai:chat".into(), - provider_api_format: "openai:chat".into(), - model_name: Some("gpt-4.1".into()), - proxy: None, - transport_profile: None, - timeouts: Some(ExecutionTimeouts { - connect_ms: Some(5_000), - total_ms: Some(LOCAL_HTTP_SUCCESS_TIMEOUT_MS), - ..ExecutionTimeouts::default() - }), - }) - .await - .expect("gzip sync execution should succeed"); + for encoding in ["gzip", "zstd"] { + let result = execution_runtime + .execute_sync(&ExecutionPlan { + request_id: format!("req-{encoding}-1"), + candidate_id: Some("cand-1".into()), + provider_name: Some("openai".into()), + provider_id: "prov-1".into(), + endpoint_id: "ep-1".into(), + key_id: "key-1".into(), + method: "POST".into(), + url: format!("http://{addr}/chat"), + headers: BTreeMap::from([("content-type".into(), "application/json".into())]), + content_type: Some("application/json".into()), + content_encoding: Some(encoding.into()), + body: RequestBody::from_json(json!({"model": "gpt-4.1"})), + stream: false, + client_api_format: "openai:chat".into(), + provider_api_format: "openai:chat".into(), + model_name: Some("gpt-4.1".into()), + proxy: None, + transport_profile: None, + timeouts: Some(ExecutionTimeouts { + connect_ms: Some(5_000), + total_ms: Some(LOCAL_HTTP_SUCCESS_TIMEOUT_MS), + ..ExecutionTimeouts::default() + }), + }) + .await + .expect("compressed sync execution should succeed"); + + assert_eq!(result.status_code, 200); + assert_eq!( + result.body.and_then(|body| body.json_body), + Some(json!({ + "content_encoding": encoding, + "body": {"model": "gpt-4.1"}, + })) + ); + } server.abort(); - - assert_eq!(result.status_code, 200); - assert_eq!( - result.body.and_then(|body| body.json_body), - Some(json!({ - "content_encoding": "gzip", - "body": {"model": "gpt-4.1"}, - })) - ); } #[tokio::test] diff --git a/apps/aether-gateway/src/executor/candidate_loop.rs b/apps/aether-gateway/src/executor/candidate_loop.rs index 053d55044..94a9dbe16 100644 --- a/apps/aether-gateway/src/executor/candidate_loop.rs +++ b/apps/aether-gateway/src/executor/candidate_loop.rs @@ -166,18 +166,31 @@ where type Error = GatewayError; async fn execute_attempt(&self, attempt: &T) -> Result, Self::Error> { - prewarm_direct_reqwest_candidate_client(attempt.execution_plan()); + let plan = attempt.execution_plan(); + let report_context = attempt.report_context(); + if let Some(response) = execution_plan_balance_capacity_response( + self.state, + self.trace_id, + self.decision, + plan, + report_context.as_ref(), + ) + .await? + { + return Ok(Some(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( self.state, self.parts.uri.path(), - attempt.execution_plan().clone(), + plan.clone(), self.trace_id, self.decision, self.plan_kind, attempt.report_kind(), - attempt.report_context(), + report_context, ) .await?; observe_gateway_stage_ms( @@ -187,10 +200,7 @@ where .as_millis() as u64, ); if let Some(response) = response.as_mut() { - attach_redaction_execution_candidate( - response, - attempt.execution_plan().candidate_id.as_deref(), - ); + attach_redaction_execution_candidate(response, plan.candidate_id.as_deref()); } Ok(response) } @@ -355,7 +365,14 @@ where }; last_attempted = Some((attempt.execution_plan().clone(), attempt.report_context())); let execute_started_at = std::time::Instant::now(); - let response = port.execute_attempt(&attempt).await?; + let response = match port.execute_attempt(&attempt).await { + Ok(response) => response, + Err(err) => { + let remaining = source.drain_execution_attempts().await?; + port.mark_unused_attempts(remaining).await?; + return Err(err); + } + }; observe_gateway_stage_ms( "stream_candidate_execute", execute_started_at.elapsed().as_millis() as u64, @@ -450,6 +467,17 @@ where candidate_index = candidate_index.as_str(), "candidate loop attempting stream execution candidate" ); + if let Some(response) = execution_plan_balance_capacity_response( + self.state, + self.trace_id, + self.decision, + &plan, + report_context.as_ref(), + ) + .await? + { + return Ok(Some(response)); + } prewarm_direct_reqwest_candidate_client(&plan); let watchdog_plan = plan.clone(); let watchdog_report_context = report_context.clone(); @@ -523,45 +551,91 @@ fn prewarm_direct_reqwest_candidate_client(plan: &aether_contracts::ExecutionPla ); } +async fn execution_plan_balance_capacity_response( + state: &AppState, + trace_id: &str, + decision: &GatewayControlDecision, + plan: &aether_contracts::ExecutionPlan, + report_context: Option<&serde_json::Value>, +) -> Result>, GatewayError> { + let rejection = match crate::control::execution_plan_balance_capacity_rejection( + state, + decision, + plan, + report_context, + ) + .await + { + Ok(rejection) => rejection, + Err(err) => { + mark_unused_local_candidate(state, plan, report_context).await; + return Err(err); + } + }; + let Some(rejection) = rejection else { + return Ok(None); + }; + mark_unused_local_candidate(state, plan, report_context).await; + let mut response = crate::api::response::build_local_auth_rejection_response( + trace_id, + Some(decision), + &rejection, + )?; + attach_redaction_execution_candidate(&mut response, plan.candidate_id.as_deref()); + Ok(Some(response)) +} + pub(crate) async fn mark_unused_local_candidates(state: &AppState, remaining: Vec) where T: AiExecutionAttempt, { for plan_and_report in remaining { let report_context = plan_and_report.report_context(); - let metadata = - local_execution_candidate_metadata_from_report_context(report_context.as_ref()); - if let Some(lease) = metadata.pool_key_lease.as_ref() { - if let Err(err) = - release_admin_provider_pool_key_lease(state.runtime_state.as_ref(), lease).await - { - warn!( - error = ?err, - "gateway candidate loop: failed to release unused pool key lease" - ); - } - } - if should_skip_unused_persistence_from_metadata(&metadata) { - continue; - } - record_local_request_candidate_status( + mark_unused_local_candidate( state, plan_and_report.execution_plan(), report_context.as_ref(), - SchedulerRequestCandidateStatusUpdate { - status: RequestCandidateStatus::Unused, - status_code: None, - error_type: None, - error_message: None, - latency_ms: None, - started_at_unix_ms: None, - finished_at_unix_ms: None, - }, ) .await; } } +async fn mark_unused_local_candidate( + state: &AppState, + plan: &aether_contracts::ExecutionPlan, + report_context: Option<&serde_json::Value>, +) { + let metadata = local_execution_candidate_metadata_from_report_context(report_context); + if let Some(lease) = metadata.pool_key_lease.as_ref() { + if let Err(err) = + release_admin_provider_pool_key_lease(state.runtime_state.as_ref(), lease).await + { + warn!( + error = ?err, + "gateway candidate loop: failed to release unused pool key lease" + ); + } + } + if should_skip_unused_persistence_from_metadata(&metadata) { + return; + } + record_local_request_candidate_status( + state, + plan, + report_context, + SchedulerRequestCandidateStatusUpdate { + status: RequestCandidateStatus::Unused, + status_code: None, + error_type: None, + error_message: None, + latency_ms: None, + started_at_unix_ms: None, + finished_at_unix_ms: None, + }, + ) + .await; +} + fn should_skip_unused_persistence(report_context: Option<&serde_json::Value>) -> bool { let metadata = local_execution_candidate_metadata_from_report_context(report_context); should_skip_unused_persistence_from_metadata(&metadata) diff --git a/apps/aether-gateway/src/fallback_metrics.rs b/apps/aether-gateway/src/fallback_metrics.rs index 1ca546c76..94f614063 100644 --- a/apps/aether-gateway/src/fallback_metrics.rs +++ b/apps/aether-gateway/src/fallback_metrics.rs @@ -179,6 +179,7 @@ mod tests { auth_context: None, admin_principal: None, local_auth_rejection: None, + model_directive_policy: Default::default(), } } diff --git a/apps/aether-gateway/src/frontdoor_loop_guard.rs b/apps/aether-gateway/src/frontdoor_loop_guard.rs index 2b8f0f6e2..827faf7f9 100644 --- a/apps/aether-gateway/src/frontdoor_loop_guard.rs +++ b/apps/aether-gateway/src/frontdoor_loop_guard.rs @@ -41,6 +41,7 @@ pub(crate) fn frontdoor_self_loop_public_ai_path(path: &str) -> bool { | "/v1/rerank" | "/v1/responses" | "/v1/responses/compact" + | "/v1/alpha/search" | "/v1beta/files" | "/upload/v1beta/files" | "/v1beta/operations" diff --git a/apps/aether-gateway/src/handlers/admin/observability/monitoring/cache_affinity.rs b/apps/aether-gateway/src/handlers/admin/observability/monitoring/cache_affinity.rs index a1a569ce8..47c1d19bf 100644 --- a/apps/aether-gateway/src/handlers/admin/observability/monitoring/cache_affinity.rs +++ b/apps/aether-gateway/src/handlers/admin/observability/monitoring/cache_affinity.rs @@ -104,6 +104,7 @@ fn is_known_admin_monitoring_api_format(value: &str) -> bool { "openai:chat" | "openai:responses" | "openai:responses:compact" + | "openai:search" | "openai:image" | "openai:video" | "openai:embedding" @@ -492,6 +493,20 @@ mod tests { ); } + #[test] + fn parses_v2_scheduler_affinity_key_for_openai_search() { + let parsed = parse_admin_monitoring_scheduler_affinity_key( + "scheduler_affinity:v2:user-key-1:openai:search:gpt-5.6-sol:codex:sessionhash", + ) + .expect("Search scheduler key should parse"); + + assert_eq!(parsed.affinity_key, "user-key-1"); + assert_eq!(parsed.api_format, "openai:search"); + assert_eq!(parsed.model_name, "gpt-5.6-sol"); + assert_eq!(parsed.client_family.as_deref(), Some("codex")); + assert_eq!(parsed.session_hash.as_deref(), Some("sessionhash")); + } + #[test] fn parses_v2_scheduler_affinity_key_with_three_segment_api_format() { let parsed = parse_admin_monitoring_scheduler_affinity_key( diff --git a/apps/aether-gateway/src/handlers/admin/provider/crud/writes.rs b/apps/aether-gateway/src/handlers/admin/provider/crud/writes.rs index cbf44424d..fa4b96728 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/crud/writes.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/crud/writes.rs @@ -6,7 +6,8 @@ use crate::handlers::admin::provider::shared::payloads::{ AdminProviderCreateRequest, AdminProviderUpdatePatch, }; use crate::handlers::admin::provider::write::provider::{ - reconcile_admin_fixed_provider_template_endpoints, reconcile_admin_fixed_provider_template_keys, + reconcile_admin_fixed_provider_template_endpoints, + reconcile_admin_fixed_provider_template_endpoints_after_update, }; use crate::handlers::admin::request::{AdminAppState, AdminRequestContext}; use crate::handlers::admin::shared::attach_admin_audit_response; @@ -80,7 +81,6 @@ pub(crate) async fn maybe_build_local_admin_provider_writes_response( .is_some() { reconcile_admin_fixed_provider_template_endpoints(state, &created_provider).await?; - reconcile_admin_fixed_provider_template_keys(state, &created_provider).await?; } return Ok(Some(attach_admin_audit_response( Json(json!({ @@ -161,8 +161,12 @@ pub(crate) async fn maybe_build_local_admin_provider_writes_response( .fixed_provider_template(&updated_record.provider_type) .is_some() { - reconcile_admin_fixed_provider_template_endpoints(state, &updated_record).await?; - reconcile_admin_fixed_provider_template_keys(state, &updated_record).await?; + reconcile_admin_fixed_provider_template_endpoints_after_update( + state, + &existing_provider, + &updated_record, + ) + .await?; } return Ok(Some( match state diff --git a/apps/aether-gateway/src/handlers/admin/provider/query/models/mod.rs b/apps/aether-gateway/src/handlers/admin/provider/query/models/mod.rs index cf3cf883a..0cda52869 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/query/models/mod.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/query/models/mod.rs @@ -64,7 +64,8 @@ use aether_data_contracts::repository::provider_catalog::{ }; use aether_model_fetch::{ aggregate_models_for_cache, fetch_models_from_transports, json_string_list, - merge_upstream_metadata, preset_models_for_provider, selected_models_fetch_endpoints, + model_catalog_upstream_metadata, preset_models_for_provider, selected_models_fetch_endpoints, + upstream_metadata_namespace_updates, }; use axum::{ body::{to_bytes, Body}, @@ -281,6 +282,7 @@ fn provider_query_attach_model_test_capabilities( fn provider_query_codex_preset_fallback( provider: &StoredProviderCatalogProvider, + fetch_error: &str, ) -> Option { if !provider.provider_type.trim().eq_ignore_ascii_case("codex") { return None; @@ -289,12 +291,59 @@ fn provider_query_codex_preset_fallback( Some(ProviderQueryKeyFetchResult { models: aggregate_models_for_cache(&models), error: None, - warning: None, + warning: Some(format!( + "Codex 动态模型目录不可用，已使用内置模型卡：{fetch_error}" + )), from_cache: false, has_success: true, }) } +async fn provider_query_persist_preset_models( + state: &AdminAppState<'_>, + provider: &StoredProviderCatalogProvider, + key: &StoredProviderCatalogKey, + models: &[Value], +) -> Result<(), GatewayError> { + if models.is_empty() { + return Ok(()); + } + ::write_upstream_models_cache( + state.app(), + &provider.id, + &key.id, + models, + ) + .await; + if let Some(catalog_metadata) = model_catalog_upstream_metadata(&provider.provider_type, models) + { + provider_query_persist_upstream_metadata(state, key, &catalog_metadata).await?; + } + Ok(()) +} + +async fn provider_query_persist_upstream_metadata( + state: &AdminAppState<'_>, + key: &StoredProviderCatalogKey, + upstream_metadata: &Value, +) -> Result<(), GatewayError> { + let updated_at = current_unix_secs(); + for (namespace, value) in + upstream_metadata_namespace_updates(key.upstream_metadata.as_ref(), upstream_metadata) + { + state + .app() + .upsert_provider_catalog_key_upstream_metadata_namespace( + &key.id, + &namespace, + &value, + Some(updated_at), + ) + .await?; + } + Ok(()) +} + mod model_test; pub(crate) use self::model_test::{ @@ -439,11 +488,9 @@ async fn provider_query_fetch_models_for_key( let selected_endpoints = selected_models_fetch_endpoints(endpoints, key); if selected_endpoints.is_empty() { if let Some(models) = preset_models_for_provider(&provider.provider_type) { - let models = provider_query_filter_models_for_key( - provider, - key, - aggregate_models_for_cache(&models), - ); + let models = aggregate_models_for_cache(&models); + provider_query_persist_preset_models(state, provider, key, &models).await?; + let models = provider_query_filter_models_for_key(provider, key, models); return Ok(ProviderQueryKeyFetchResult { models, error: None, @@ -492,7 +539,11 @@ async fn provider_query_fetch_models_for_key( Ok(outcome) => outcome, Err(err) => { all_errors.push(err); - if let Some(fallback) = provider_query_codex_preset_fallback(provider) { + if let Some(fallback) = + provider_query_codex_preset_fallback(provider, &all_errors.join("; ")) + { + provider_query_persist_preset_models(state, provider, key, &fallback.models) + .await?; return Ok(fallback); } return Ok(ProviderQueryKeyFetchResult { @@ -517,20 +568,14 @@ async fn provider_query_fetch_models_for_key( .await; } if let Some(upstream_metadata) = outcome.upstream_metadata.as_ref() { - let merged_metadata = - merge_upstream_metadata(key.upstream_metadata.as_ref(), upstream_metadata); - state - .app() - .update_provider_catalog_key_upstream_metadata( - &key.id, - Some(&merged_metadata), - Some(current_unix_secs()), - ) - .await?; + provider_query_persist_upstream_metadata(state, key, upstream_metadata).await?; } if unique_models.is_empty() && !all_errors.is_empty() { - if let Some(fallback) = provider_query_codex_preset_fallback(provider) { + if let Some(fallback) = + provider_query_codex_preset_fallback(provider, &all_errors.join("; ")) + { + provider_query_persist_preset_models(state, provider, key, &fallback.models).await?; return Ok(fallback); } } 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 77d1455af..671f0bc19 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 @@ -29,10 +29,7 @@ use crate::handlers::shared::{ provider_key_status_snapshot_payload, }; use crate::model_fetch::ModelFetchRuntimeState; -use crate::provider_key_auth::{ - provider_key_auth_semantics, provider_key_configured_api_formats, - provider_key_inherits_provider_api_formats, -}; +use crate::provider_key_auth::provider_key_auth_semantics; use crate::provider_transport::antigravity::{ build_antigravity_safe_v1internal_request, build_antigravity_static_identity_headers, classify_local_antigravity_request_support, AntigravityEnvelopeRequestType, @@ -543,6 +540,22 @@ fn provider_query_build_test_request_body_for_api_format( model: &str, route_path: &str, client_api_format: &str, +) -> Value { + provider_query_build_test_request_body_for_api_format_with_search_session( + payload, + model, + route_path, + client_api_format, + None, + ) +} + +fn provider_query_build_test_request_body_for_api_format_with_search_session( + payload: &Value, + model: &str, + route_path: &str, + client_api_format: &str, + search_session_id: Option<&str>, ) -> Value { let client_api_format = provider_query_normalize_api_format_alias(client_api_format); let override_custom_model = route_path.ends_with("/test-model-failover") @@ -568,7 +581,7 @@ fn provider_query_build_test_request_body_for_api_format( ); } else if matches!( client_api_format.as_str(), - "openai:responses" | "openai:responses:compact" + "openai:responses" | "openai:responses:compact" | "openai:search" ) && !value_has_non_empty_text(object.get("input")) { if let Some(prompt) = object @@ -580,7 +593,7 @@ fn provider_query_build_test_request_body_for_api_format( } if matches!( client_api_format.as_str(), - "openai:responses" | "openai:responses:compact" + "openai:responses" | "openai:responses:compact" | "openai:search" ) && value_has_non_empty_text(object.get("input")) { object.remove("prompt"); @@ -590,6 +603,9 @@ fn provider_query_build_test_request_body_for_api_format( { object.remove("messages"); } + if client_api_format == "openai:search" { + provider_query_ensure_search_test_fields(object, payload, search_session_id); + } } return body; } @@ -620,6 +636,15 @@ fn provider_query_build_test_request_body_for_api_format( "temperature": 0.7, "stream": true, }), + "openai:search" => json!({ + "id": provider_query_search_test_session_id(search_session_id), + "model": model, + "input": message, + "commands": { + "search_query": [{"q": message}] + }, + "max_output_tokens": 256, + }), "claude:messages" => json!({ "model": model, "messages": [{ @@ -680,7 +705,7 @@ fn provider_query_insert_default_test_conversation( .entry("top_n".to_string()) .or_insert_with(|| Value::from(4_u64)); } - "openai:responses" | "openai:responses:compact" => { + "openai:responses" | "openai:responses:compact" | "openai:search" => { object.insert("input".to_string(), Value::String(message)); } "claude:messages" => { @@ -698,6 +723,48 @@ fn provider_query_insert_default_test_conversation( } } +fn provider_query_search_test_session_id(search_session_id: Option<&str>) -> String { + search_session_id + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(|value| format!("aether-model-test-{value}")) + .unwrap_or_else(|| format!("aether-model-test-{}", Uuid::new_v4().simple())) +} + +fn provider_query_ensure_search_test_fields( + object: &mut Map, + payload: &Value, + search_session_id: Option<&str>, +) { + let query = object + .get("input") + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) + .or_else(|| provider_query_extract_message(payload)) + .unwrap_or_else(|| DEFAULT_PROVIDER_QUERY_TEST_MESSAGE.to_string()); + let has_session_id = object + .get("id") + .and_then(Value::as_str) + .is_some_and(|value| !value.trim().is_empty()); + if !has_session_id { + object.insert( + "id".to_string(), + provider_query_search_test_session_id(search_session_id).into(), + ); + } + object + .entry("input".to_string()) + .or_insert_with(|| Value::String(query.clone())); + object + .entry("commands".to_string()) + .or_insert_with(|| json!({"search_query": [{"q": query}]})); + object + .entry("max_output_tokens".to_string()) + .or_insert_with(|| Value::from(256_u64)); +} + fn provider_query_grok_test_client_api_format(provider_api_format: &str) -> &'static str { match provider_query_normalize_api_format_alias(provider_api_format).as_str() { "openai:responses" | "openai:responses:compact" => "openai:responses", @@ -767,7 +834,7 @@ fn provider_query_request_body_has_conversation_for_api_format( client_api_format: &str, ) -> bool { match provider_query_normalize_api_format_alias(client_api_format).as_str() { - "openai:responses" | "openai:responses:compact" => { + "openai:responses" | "openai:responses:compact" | "openai:search" => { value_has_non_empty_text(body.get("input")) || value_has_non_empty_text(body.get("prompt")) } @@ -828,14 +895,12 @@ fn provider_query_resolve_standard_test_upstream_is_stream( provider_type: &str, provider_api_format: &str, ) -> bool { - let hard_requires_streaming = crate::ai_serving::force_upstream_streaming_for_provider( + crate::ai_serving::resolve_upstream_is_stream_for_provider( + endpoint_config, provider_type, provider_api_format, - ); - crate::ai_serving::resolve_upstream_is_stream_from_endpoint_config( - endpoint_config, false, - hard_requires_streaming, + false, ) } @@ -876,15 +941,11 @@ fn provider_query_key_supports_endpoint( provider_type: &str, endpoint_api_format: &str, ) -> bool { - if provider_key_inherits_provider_api_formats(key, provider_type) { - return true; - } - let formats = provider_key_configured_api_formats(key); - let endpoint_api_format = provider_query_normalize_api_format_alias(endpoint_api_format); - formats.is_empty() - || formats - .iter() - .any(|value| provider_query_normalize_api_format_alias(value) == endpoint_api_format) + crate::handlers::shared::provider_catalog_key_supports_format( + key, + provider_type, + endpoint_api_format, + ) } async fn provider_query_select_preferred_non_kiro_endpoint( @@ -1626,6 +1687,15 @@ fn provider_query_standard_execution_response_body( { return None; } + if result.status_code < 400 + && provider_query_normalize_api_format_alias(provider_api_format) == "openai:search" + && !body + .get("output") + .and_then(Value::as_str) + .is_some_and(|value| !value.trim().is_empty()) + { + return None; + } Some(body) } @@ -2104,8 +2174,18 @@ async fn provider_query_execute_openai_image_test_candidate( route_path, ); let incoming_request_headers = provider_query_extract_request_headers(payload); + let image_request_path = if request_body.get("image").is_some() + || request_body + .get("images") + .and_then(Value::as_array) + .is_some_and(|images| !images.is_empty()) + { + "/v1/images/edits" + } else { + "/v1/images/generations" + }; let mut synthetic_request = http::Request::builder() - .uri("/v1/images/generations") + .uri(image_request_path) .body(()) .map_err(|err| GatewayError::Internal(err.to_string()))?; *synthetic_request.headers_mut() = incoming_request_headers; @@ -2116,11 +2196,18 @@ async fn provider_query_execute_openai_image_test_candidate( &parts, &request_body, None, - provider_query_openai_image_normalize_options(provider_type), + provider_query_openai_image_normalize_options( + provider_type, + Some(candidate.effective_model.as_str()), + ), ) else { return Ok(provider_query_skipped_execution_outcome( request_body.clone(), - provider_query_openai_image_normalize_failure_message(provider_type, &request_body), + provider_query_openai_image_normalize_failure_message( + provider_type, + Some(candidate.effective_model.as_str()), + &request_body, + ), )); }; @@ -2139,29 +2226,50 @@ async fn provider_query_execute_openai_image_test_candidate( .provider_type .trim() .eq_ignore_ascii_case("codex"); - let mut provider_request_body = if is_chatgpt_web { - match crate::ai_serving::build_chatgpt_web_image_request_body(&parts, &request_body, None) { - Ok(body) => body, - Err(err) => err.to_error_json(), - } - } else if is_codex || is_grok { - crate::ai_serving::build_openai_image_provider_request_body(&normalized_request) + let upstream_is_stream = crate::ai_serving::resolve_upstream_is_stream_for_provider( + transport.endpoint.config.as_ref(), + transport.provider.provider_type.as_str(), + "openai:image", + request_body + .get("stream") + .and_then(Value::as_bool) + .unwrap_or(false), + false, + ); + let provider_request_body = if is_chatgpt_web { + Some( + match crate::ai_serving::build_chatgpt_web_image_request_body( + &parts, + &request_body, + None, + ) { + Ok(body) => body, + Err(err) => err.to_error_json(), + }, + ) + } else if is_codex { + crate::ai_serving::build_codex_openai_image_api_provider_request_body( + &normalized_request, + Some(candidate.effective_model.as_str()), + upstream_is_stream, + ) + } else if is_grok { + Some(crate::ai_serving::build_openai_image_provider_request_body( + &normalized_request, + )) } else { crate::ai_serving::build_openai_image_api_provider_request_body( &normalized_request, Some(candidate.effective_model.as_str()), + upstream_is_stream, ) }; - if !is_chatgpt_web { - crate::ai_serving::apply_codex_openai_responses_special_body_edits( - &mut provider_request_body, - transport.provider.provider_type.as_str(), - "openai:image", - transport.endpoint.body_rules.as_ref(), - Some(candidate.key.id.as_str()), - ); - } - + let Some(provider_request_body) = provider_request_body else { + return Ok(provider_query_skipped_execution_outcome( + request_body, + "Provider request is outside the Codex Images contract", + )); + }; let oauth_auth = state.resolve_local_oauth_header_auth(&transport).await?; let Some((auth_header, auth_value)) = crate::provider_transport::resolve_openai_image_auth(&transport).or(oauth_auth) @@ -2193,10 +2301,12 @@ async fn provider_query_execute_openai_image_test_candidate( headers: &parts.headers, auth_header: &auth_header, auth_value: &auth_value, - accept: if is_codex || is_chatgpt_web { - "text/event-stream" + accept: if is_codex { + None + } else if upstream_is_stream { + Some("text/event-stream") } else { - "application/json" + Some("application/json") }, header_rules: transport.endpoint.header_rules.as_ref(), provider_request_body: &provider_request_body, @@ -2221,7 +2331,7 @@ async fn provider_query_execute_openai_image_test_candidate( request_headers.insert("x-aether-chatgpt-web-image".to_string(), "1".to_string()); } else if is_grok { } else { - crate::ai_serving::apply_codex_openai_responses_special_headers( + crate::ai_serving::apply_codex_openai_special_headers( &mut request_headers, &provider_request_body, &parts.headers, @@ -2260,13 +2370,9 @@ async fn provider_query_execute_openai_image_test_candidate( }; let request_url = provider_query_openai_image_test_upstream_url( &transport, - Some(parts.uri.path()), + Some(image_request_path), parts.uri.query(), ); - let upstream_is_stream = provider_request_body - .get("stream") - .and_then(Value::as_bool) - .unwrap_or(true); let plan = ExecutionPlan { request_id: trace_id.to_string(), @@ -2842,12 +2948,14 @@ async fn provider_query_execute_standard_test_candidate( crate::ai_serving::normalize_api_format_alias(provider_api_format); let client_api_format = provider_query_standard_test_client_api_format(normalized_provider_api_format.as_str()); - let original_request_body = provider_query_build_test_request_body_for_api_format( - payload, - &candidate.effective_model, - route_path, - client_api_format, - ); + let original_request_body = + provider_query_build_test_request_body_for_api_format_with_search_session( + payload, + &candidate.effective_model, + route_path, + client_api_format, + Some(trace_id), + ); if crate::provider_transport::is_windsurf_provider_transport(&transport) && provider_query_normalize_api_format_alias(candidate.endpoint.api_format.as_str()) == "openai:chat" @@ -2992,6 +3100,45 @@ async fn provider_query_execute_standard_test_candidate( ); provider_request_body } + "openai:search" => { + let Some(mut provider_request_body) = + crate::provider_transport::build_same_format_provider_request_body( + crate::provider_transport::SameFormatProviderRequestBodyInput { + body_json: &request_body, + mapped_model: request_model, + client_api_format, + provider_api_format, + source_model: request_body.get("model").and_then(Value::as_str), + family: crate::provider_transport::SameFormatProviderFamily::Standard, + body_rules: transport.endpoint.body_rules.as_ref(), + request_headers: Some(&incoming_request_headers), + upstream_is_stream, + force_body_stream_field: require_body_stream_field, + kiro_auth_config: None, + is_claude_code: false, + enable_model_directives: false, + }, + ) + else { + return Ok(provider_query_skipped_execution_outcome( + request_body.clone(), + format!("Provider request body could not be built for {provider_api_format}"), + )); + }; + if let Err(err) = crate::provider_transport::apply_transport_request_body_semantics( + &mut provider_request_body, + &transport, + normalized_provider_api_format.as_str(), + ) { + return Ok(provider_query_skipped_execution_outcome( + provider_request_body, + format!( + "Provider request body is not compatible with transport semantics: {err}" + ), + )); + } + provider_request_body + } "openai:embedding" | "gemini:embedding" | "jina:embedding" @@ -3046,6 +3193,37 @@ async fn provider_query_execute_standard_test_candidate( upstream_is_stream, require_body_stream_field, ); + let source_model = provider_query_request_body_model(&request_body, request_model); + let codex_model_capabilities = crate::ai_serving::codex_model_capabilities_for_transport( + &transport, + provider_api_format, + request_model, + source_model, + ); + if matches!( + normalized_provider_api_format.as_str(), + "openai:chat" | "openai:responses" | "openai:responses:compact" | "openai:search" + ) && crate::ai_serving::finalize_openai_provider_request_with_codex_model_capabilities( + &mut provider_request_body, + crate::ai_serving::OpenAiProviderRequestFinalization { + source_api_format: client_api_format, + provider_api_format, + provider_type: transport.provider.provider_type.as_str(), + provider_model: request_model, + source_model, + body_rules: transport.endpoint.body_rules.as_ref(), + upstream_is_stream, + require_body_stream_field, + }, + codex_model_capabilities.as_ref(), + ) + .is_err() + { + return Ok(provider_query_skipped_execution_outcome( + provider_request_body, + "Provider request body violates the OpenAI provider contract", + )); + } if crate::provider_transport::is_gemini_cli_provider_transport(&transport) && normalized_provider_api_format == "gemini:generate_content" { @@ -3122,6 +3300,7 @@ async fn provider_query_execute_standard_test_candidate( "openai:chat" | "openai:responses" | "openai:responses:compact" + | "openai:search" | "claude:messages" | "gemini:generate_content" | "openai:embedding" @@ -3137,6 +3316,7 @@ async fn provider_query_execute_standard_test_candidate( "openai:chat" | "openai:responses" | "openai:responses:compact" + | "openai:search" | "openai:embedding" | "jina:embedding" | "doubao:embedding" @@ -3202,7 +3382,7 @@ async fn provider_query_execute_standard_test_candidate( &BTreeMap::new(), Some("application/json"), ), - "openai:responses" | "openai:responses:compact" => { + "openai:responses" | "openai:responses:compact" | "openai:search" => { crate::provider_transport::auth::build_complete_passthrough_headers_with_auth( &parts.headers, auth_header.as_deref().unwrap_or_default(), @@ -3268,8 +3448,10 @@ async fn provider_query_execute_standard_test_candidate( response_body: None, }); } - if crate::ai_serving::is_openai_responses_format(provider_api_format) { - crate::ai_serving::apply_codex_openai_responses_special_headers( + if crate::ai_serving::is_openai_responses_family_format(provider_api_format) + || crate::ai_serving::api_format_alias_matches(provider_api_format, "openai:search") + { + crate::ai_serving::apply_codex_openai_special_headers( &mut request_headers, &provider_request_body, &parts.headers, @@ -3278,9 +3460,17 @@ async fn provider_query_execute_standard_test_candidate( Some(trace_id), transport.key.decrypted_auth_config.as_deref(), ); - crate::provider_transport::apply_local_auth_config_header_overrides( + let final_provider_model = provider_request_body + .get("model") + .and_then(Value::as_str) + .unwrap_or(request_model); + crate::ai_serving::apply_codex_openai_responses_lite_header_with_capabilities( &mut request_headers, - transport.key.decrypted_auth_config.as_deref(), + transport.provider.provider_type.as_str(), + provider_api_format, + final_provider_model, + source_model, + codex_model_capabilities.as_ref(), ); } if !uses_vertex_query_auth { diff --git a/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/adapter.rs b/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/adapter.rs index 86333dc64..9bdd7802a 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/adapter.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/adapter.rs @@ -33,6 +33,8 @@ pub(super) fn provider_query_standard_test_client_api_format( let normalized_api_format = crate::ai_serving::normalize_api_format_alias(provider_api_format); if normalized_api_format == "openai:responses:compact" { "openai:responses:compact" + } else if normalized_api_format == "openai:search" { + "openai:search" } else if crate::ai_serving::is_embedding_api_format(&normalized_api_format) { "openai:embedding" } else if crate::ai_serving::is_rerank_api_format(&normalized_api_format) { @@ -71,6 +73,7 @@ pub(super) fn provider_query_standard_test_unsupported_reason( } "openai:responses" | "openai:responses:compact" + | "openai:search" | "claude:messages" | "openai:embedding" | "jina:embedding" @@ -261,6 +264,7 @@ pub(super) fn provider_query_test_adapter_for_provider_api_format( "openai:chat" | "openai:responses" | "openai:responses:compact" + | "openai:search" | "claude:messages" | "gemini:generate_content" | "gemini:interactions" @@ -365,6 +369,7 @@ pub(super) fn provider_query_transport_supports_model_test_execution( } "openai:responses" | "openai:responses:compact" + | "openai:search" | "openai:embedding" | "jina:embedding" | "doubao:embedding" diff --git a/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/capabilities.rs b/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/capabilities.rs index 156e23573..e2a38402e 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/capabilities.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/capabilities.rs @@ -9,16 +9,19 @@ pub(super) struct ProviderQueryOpenAiImageTestCapability(AdminProviderOpenAiImag pub(super) fn provider_query_openai_image_test_capability( provider_type: &str, + provider_model: Option<&str>, ) -> ProviderQueryOpenAiImageTestCapability { ProviderQueryOpenAiImageTestCapability(admin_provider_openai_image_test_capability( provider_type, + provider_model, )) } pub(super) fn provider_query_openai_image_normalize_options( provider_type: &str, + provider_model: Option<&str>, ) -> crate::ai_serving::OpenAiImageNormalizeOptions { - admin_provider_openai_image_normalize_options(provider_type) + admin_provider_openai_image_normalize_options(provider_type, provider_model) } pub(super) fn provider_query_openai_image_requested_count(request_body: &Value) -> Option { @@ -35,9 +38,10 @@ pub(super) fn provider_query_openai_image_requested_count(request_body: &Value) pub(super) fn provider_query_openai_image_normalize_failure_message( provider_type: &str, + provider_model: Option<&str>, request_body: &Value, ) -> String { - let capability = provider_query_openai_image_test_capability(provider_type); + let capability = provider_query_openai_image_test_capability(provider_type, provider_model); if provider_query_openai_image_requested_count(request_body) .is_some_and(|value| !capability.0.supports_generation_count(value)) { diff --git a/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/model_mapping.rs b/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/model_mapping.rs index 65bcdfa19..e4f670e40 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/model_mapping.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/model_mapping.rs @@ -101,8 +101,7 @@ fn provider_query_model_mapping_matches_endpoint( ) -> bool { let api_format_matches = mapping.api_formats.as_ref().is_none_or(|api_formats| { api_formats.iter().any(|value| { - aether_scheduler_core::normalize_api_format(value) - == aether_scheduler_core::normalize_api_format(&endpoint.api_format) + crate::ai_serving::api_format_permission_covers(value, &endpoint.api_format) }) }); if !api_format_matches { @@ -367,3 +366,41 @@ fn provider_query_parse_mapping_string_list_array( } Ok(parsed) } + +#[cfg(test)] +mod tests { + use super::*; + + fn sample_endpoint(api_format: &str) -> StoredProviderCatalogEndpoint { + StoredProviderCatalogEndpoint::new( + "endpoint-1".to_string(), + "provider-1".to_string(), + api_format.to_string(), + None, + None, + true, + ) + .expect("endpoint should build") + } + + fn sample_mapping(api_format: &str) -> StoredProviderModelMapping { + StoredProviderModelMapping { + name: "gpt-5.6-luna".to_string(), + priority: 1, + api_formats: Some(vec![api_format.to_string()]), + endpoint_ids: None, + } + } + + #[test] + fn responses_mapping_scope_covers_search_in_one_direction() { + assert!(provider_query_model_mapping_matches_endpoint( + &sample_mapping("openai:responses"), + &sample_endpoint("openai:search"), + )); + assert!(!provider_query_model_mapping_matches_endpoint( + &sample_mapping("openai:search"), + &sample_endpoint("openai:responses"), + )); + } +} diff --git a/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/tests.rs b/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/tests.rs index 4ea4bac28..dd7285a5f 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/tests.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/tests.rs @@ -337,6 +337,11 @@ fn provider_query_standard_test_resolves_codex_responses_upstream_streaming() { "codex", "openai:responses:compact", )); + assert!(!provider_query_resolve_standard_test_upstream_is_stream( + Some(&json!({"upstream_stream_policy": "force_stream"})), + "codex", + "openai:responses:compact", + )); assert!(!provider_query_resolve_standard_test_upstream_is_stream( None, "custom", @@ -470,6 +475,118 @@ fn provider_query_compact_test_request_body_defaults_to_responses_input() { assert!(body.get("messages").is_none()); } +#[test] +fn provider_query_search_test_request_body_defaults_to_typed_search_input() { + let payload = json!({"message": "find current documentation"}); + + let client_api_format = provider_query_standard_test_client_api_format("openai:search"); + let body = provider_query_build_test_request_body_for_api_format( + &payload, + "gpt-5.6-sol", + "/api/admin/provider-query/test-model", + client_api_format, + ); + + assert_eq!(client_api_format, "openai:search"); + assert!(body["id"] + .as_str() + .is_some_and(|value| value.starts_with("aether-model-test-"))); + assert_eq!(body["model"], json!("gpt-5.6-sol")); + assert_eq!(body["input"], json!("find current documentation")); + assert_eq!( + body["commands"]["search_query"][0]["q"], + json!("find current documentation") + ); + assert_eq!(body["max_output_tokens"], json!(256)); + assert!(body.get("messages").is_none()); + assert!(body.get("stream").is_none()); +} + +#[test] +fn provider_query_search_test_completes_missing_protocol_fields() { + let payload = json!({ + "request_body": { + "model": "gpt-5.6-luna", + "input": "find release notes" + } + }); + + let body = provider_query_build_test_request_body_for_api_format_with_search_session( + &payload, + "gpt-5.6-luna", + "/api/admin/provider-query/test-model", + "openai:search", + Some("trace-model-test-1"), + ); + + assert_eq!(body["id"], json!("aether-model-test-trace-model-test-1")); + assert_eq!( + body["commands"]["search_query"][0]["q"], + json!("find release notes") + ); + assert_eq!(body["max_output_tokens"], json!(256)); +} + +#[test] +fn provider_query_search_test_preserves_a_non_empty_client_session_id() { + let payload = json!({ + "request_body": { + "id": "client-search-session", + "model": "gpt-5.6-luna", + "input": "find release notes" + } + }); + + let body = provider_query_build_test_request_body_for_api_format_with_search_session( + &payload, + "gpt-5.6-luna", + "/api/admin/provider-query/test-model", + "openai:search", + Some("trace-model-test-1"), + ); + + assert_eq!(body["id"], json!("client-search-session")); +} + +#[test] +fn provider_query_search_success_requires_non_empty_output() { + fn result(body: Value) -> aether_contracts::ExecutionResult { + aether_contracts::ExecutionResult { + request_id: "provider-search-test".to_string(), + candidate_id: Some("candidate-0".to_string()), + status_code: 200, + headers: BTreeMap::new(), + body: Some(aether_contracts::ResponseBody { + json_body: Some(body), + body_bytes_b64: None, + }), + telemetry: None, + error: None, + } + } + + assert!(provider_query_standard_execution_response_body( + "openai:search", + &result(json!({})), + None, + ) + .is_none()); + assert!(provider_query_standard_execution_response_body( + "openai:search", + &result(json!({"output": " "})), + None, + ) + .is_none()); + assert_eq!( + provider_query_standard_execution_response_body( + "openai:search", + &result(json!({"output": "search result", "encrypted_output": "ciphertext"})), + None, + ), + Some(json!({"output": "search result", "encrypted_output": "ciphertext"})) + ); +} + #[test] fn provider_query_embedding_test_request_body_defaults_to_embedding_input() { let payload = json!({"message": "hello from embedding"}); @@ -632,6 +749,10 @@ fn provider_query_test_adapter_routes_fixed_provider_endpoint_types() { provider_query_test_adapter_for_provider_api_format("codex", "openai:responses:compact"), Some(ProviderQueryTestAdapter::Standard) ); + assert_eq!( + provider_query_test_adapter_for_provider_api_format("codex", "openai:search"), + Some(ProviderQueryTestAdapter::Standard) + ); assert_eq!( provider_query_test_adapter_for_provider_api_format("chatgpt_web", "openai:image"), Some(ProviderQueryTestAdapter::OpenAiImage) @@ -717,6 +838,10 @@ fn provider_query_endpoint_priority_prefers_text_before_cli_and_image() { provider_query_model_test_endpoint_priority("codex", "openai:responses:compact"), Some(1) ); + assert_eq!( + provider_query_model_test_endpoint_priority("codex", "openai:search"), + Some(1) + ); assert_eq!( provider_query_model_test_endpoint_priority("chatgpt_web", "openai:image"), Some(2) @@ -997,7 +1122,7 @@ fn provider_query_grok_image_test_allows_multi_generation_count() { &parts, &body, None, - provider_query_openai_image_normalize_options("grok"), + provider_query_openai_image_normalize_options("grok", Some("grok-imagine-image")), ) .expect("grok image model tests should allow multi-image generation"); let provider_body = crate::ai_serving::build_openai_image_provider_request_body(&normalized); @@ -1065,12 +1190,48 @@ fn provider_query_non_grok_image_test_keeps_single_generation_boundary() { &parts, &body, None, - provider_query_openai_image_normalize_options("chatgpt_web"), + provider_query_openai_image_normalize_options("chatgpt_web", Some("gpt-image-2")), ) .is_none() ); assert_eq!( - provider_query_openai_image_normalize_failure_message("chatgpt_web", &body), + provider_query_openai_image_normalize_failure_message( + "chatgpt_web", + Some("gpt-image-2"), + &body, + ), + "Provider request body could not be normalized for openai:image: selected provider supports n=1..1 for generation" + ); +} + +#[test] +fn provider_query_dall_e_3_image_test_keeps_single_generation_boundary() { + let request = http::Request::builder() + .uri("/v1/images/generations") + .body(()) + .expect("request should build"); + let (parts, _) = request.into_parts(); + let body = json!({ + "model": "dall-e-3", + "prompt": "draw", + "n": 2 + }); + + assert!( + crate::ai_serving::normalize_openai_image_request_with_options( + &parts, + &body, + None, + provider_query_openai_image_normalize_options("openai", Some("dall-e-3")), + ) + .is_none() + ); + assert_eq!( + provider_query_openai_image_normalize_failure_message( + "openai", + Some("dall-e-3"), + &body, + ), "Provider request body could not be normalized for openai:image: selected provider supports n=1..1 for generation" ); } diff --git a/apps/aether-gateway/src/handlers/admin/provider/shared/model_test_capabilities.rs b/apps/aether-gateway/src/handlers/admin/provider/shared/model_test_capabilities.rs index c025af8e4..88937a83e 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/shared/model_test_capabilities.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/shared/model_test_capabilities.rs @@ -1,5 +1,6 @@ use crate::image_capabilities::{ - openai_image_normalize_options_for_provider, openai_image_provider_max_generation_count, + openai_image_normalize_options_for_provider, + openai_image_provider_max_generation_count_for_model, }; use serde_json::{json, Value}; @@ -24,16 +25,21 @@ impl AdminProviderOpenAiImageTestCapability { pub(crate) fn admin_provider_openai_image_test_capability( provider_type: &str, + provider_model: Option<&str>, ) -> AdminProviderOpenAiImageTestCapability { AdminProviderOpenAiImageTestCapability { - max_generation_count: openai_image_provider_max_generation_count(provider_type), + max_generation_count: openai_image_provider_max_generation_count_for_model( + provider_type, + provider_model, + ), } } pub(crate) fn admin_provider_openai_image_normalize_options( provider_type: &str, + provider_model: Option<&str>, ) -> crate::ai_serving::OpenAiImageNormalizeOptions { - openai_image_normalize_options_for_provider(provider_type) + openai_image_normalize_options_for_provider(provider_type, provider_model) } pub(crate) fn admin_provider_model_test_capabilities_payload( @@ -47,7 +53,7 @@ pub(crate) fn admin_provider_model_test_capabilities_payload( provider_type.eq_ignore_ascii_case("grok") && model_id == GROK_IMAGE_EDIT_MODEL_ID; let openai_image = if supports_image_generation { Some(json!({ - "max_generation_count": admin_provider_openai_image_test_capability(provider_type).max_generation_count, + "max_generation_count": admin_provider_openai_image_test_capability(provider_type, Some(model_id)).max_generation_count, "supports_generation": !is_grok_image_edit, "supports_edit": is_grok_image_edit, })) @@ -105,6 +111,13 @@ mod tests { assert!(payload["openai:image"].is_null()); } + #[test] + fn dall_e_3_reports_its_model_specific_generation_limit() { + let payload = admin_provider_model_test_capabilities_payload("openai", "dall-e-3", true); + + assert_eq!(payload["openai:image"]["max_generation_count"], 1); + } + #[test] fn grok_image_support_uses_catalog_model_ids_not_global_fallback() { assert!(admin_provider_model_supports_image_generation( diff --git a/apps/aether-gateway/src/handlers/admin/provider/write/keys/create.rs b/apps/aether-gateway/src/handlers/admin/provider/write/keys/create.rs index e715043d9..9601c0eb5 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/write/keys/create.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/write/keys/create.rs @@ -2,7 +2,7 @@ use crate::handlers::admin::provider::shared::payloads::AdminProviderKeyCreateRe use crate::handlers::admin::provider::write::normalize::{ normalize_allow_auth_channel_mismatch_formats, normalize_api_format_json_object_keys, normalize_api_format_list, normalize_auth_type, normalize_auth_type_by_format, - normalize_max_probe_interval_minutes, validate_vertex_api_formats, + normalize_max_probe_interval_minutes, normalize_rate_multipliers, validate_vertex_api_formats, }; use crate::handlers::admin::request::AdminAppState; use crate::handlers::admin::shared::{ @@ -165,7 +165,7 @@ pub(crate) async fn build_admin_create_provider_key_record( }, encrypted_api_key, encrypted_auth_config, - normalize_api_format_json_object_keys(payload.rate_multipliers, "rate_multipliers")?, + normalize_rate_multipliers(payload.rate_multipliers)?, None, normalize_string_list(payload.allowed_models).map(|value| json!(value)), None, diff --git a/apps/aether-gateway/src/handlers/admin/provider/write/keys/update.rs b/apps/aether-gateway/src/handlers/admin/provider/write/keys/update.rs index f3893d3c9..2204bec34 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/write/keys/update.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/write/keys/update.rs @@ -2,7 +2,7 @@ use crate::handlers::admin::provider::shared::payloads::AdminProviderKeyUpdatePa use crate::handlers::admin::provider::write::normalize::{ normalize_allow_auth_channel_mismatch_formats, normalize_api_format_json_object_keys, normalize_api_format_list, normalize_auth_type, normalize_auth_type_by_format, - normalize_max_probe_interval_minutes, validate_vertex_api_formats, + normalize_max_probe_interval_minutes, normalize_rate_multipliers, validate_vertex_api_formats, }; use crate::handlers::admin::request::AdminAppState; use crate::handlers::admin::shared::{ @@ -260,8 +260,7 @@ pub(crate) async fn build_admin_update_provider_key_record( updated.name = trimmed.to_string(); } if fields.contains("rate_multipliers") { - updated.rate_multipliers = - normalize_api_format_json_object_keys(payload.rate_multipliers, "rate_multipliers")?; + updated.rate_multipliers = normalize_rate_multipliers(payload.rate_multipliers)?; } if let Some(internal_priority) = payload.internal_priority { updated.internal_priority = internal_priority; 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 436e54037..5f5f8c5f6 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/write/normalize.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/write/normalize.rs @@ -42,6 +42,31 @@ pub(crate) fn normalize_api_format_json_object_keys( Ok(Some(serde_json::Value::Object(normalized))) } +pub(crate) fn normalize_rate_multipliers( + value: Option, +) -> Result, String> { + let Some(value) = normalize_json_like_object(value, "rate_multipliers")? else { + return Ok(None); + }; + let serde_json::Value::Object(map) = value else { + return Ok(Some(value)); + }; + let mut normalized = serde_json::Map::new(); + for (key, value) in map { + let canonical = crate::ai_serving::normalize_api_format_alias(&key); + let multiplier = value + .as_f64() + .filter(|value| value.is_finite() && *value >= 0.0) + .ok_or_else(|| format!("rate_multipliers.{canonical} 必须是大于或等于 0 的有限数值"))?; + normalized.insert(canonical, serde_json::Value::from(multiplier)); + } + if normalized.is_empty() { + Ok(None) + } else { + Ok(Some(serde_json::Value::Object(normalized))) + } +} + pub(crate) fn normalize_auth_type_by_format( value: Option, field_name: &str, @@ -212,7 +237,7 @@ mod tests { normalize_allow_auth_channel_mismatch_formats, normalize_api_format_json_object_keys, normalize_api_format_list, normalize_auth_type, normalize_auth_type_by_format, normalize_chat_pii_redaction_config, normalize_pool_advanced_config, - normalize_provider_type_input, validate_vertex_api_formats, + normalize_provider_type_input, normalize_rate_multipliers, validate_vertex_api_formats, }; use serde_json::json; @@ -224,6 +249,21 @@ mod tests { ); } + #[test] + fn rate_multipliers_require_non_negative_finite_numbers() { + assert_eq!( + normalize_rate_multipliers(Some(json!({" OPENAI:RESPONSES ": 1.25}))) + .expect("valid multiplier should normalize"), + Some(json!({"openai:responses": 1.25})) + ); + for value in [ + json!({"openai:responses": -0.1}), + json!({"openai:responses": "1.0"}), + ] { + assert!(normalize_rate_multipliers(Some(value)).is_err()); + } + } + #[test] fn normalize_pool_advanced_rejects_legacy_booleans() { assert_eq!( diff --git a/apps/aether-gateway/src/handlers/admin/provider/write/provider.rs b/apps/aether-gateway/src/handlers/admin/provider/write/provider.rs index a9a727dbd..c8cea50e2 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/write/provider.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/write/provider.rs @@ -3,11 +3,79 @@ mod endpoint; mod template; mod update; +fn normalize_provider_request_timeout(value: Option) -> Result, String> { + let max = aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_SECS as f64; + match value { + Some(value) if (1.0..=max).contains(&value) => Ok(Some(value)), + Some(_) => Err(format!( + "request_timeout 必须是 1 到 {} 之间的数字", + aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_SECS + )), + None => Ok(None), + } +} + +fn normalize_provider_stream_first_byte_timeout(value: Option) -> Result, String> { + let max = aether_contracts::MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_SECS as f64; + match value { + Some(value) if (1.0..=max).contains(&value) => Ok(Some(value)), + Some(_) => Err(format!( + "stream_first_byte_timeout 必须是 1 到 {} 之间的数字", + aether_contracts::MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_SECS + )), + None => Ok(None), + } +} + pub(crate) use self::create::build_admin_create_provider_record; pub(crate) use self::endpoint::build_admin_fixed_provider_endpoint_record; pub(crate) use self::template::{ apply_admin_fixed_provider_endpoint_template_overrides, reconcile_admin_fixed_provider_template_endpoints, - reconcile_admin_fixed_provider_template_keys, + reconcile_admin_fixed_provider_template_endpoints_after_update, }; pub(crate) use self::update::build_admin_update_provider_record; + +#[cfg(test)] +mod tests { + use super::{normalize_provider_request_timeout, normalize_provider_stream_first_byte_timeout}; + + #[test] + fn provider_request_timeout_accepts_the_execution_protocol_range() { + assert_eq!(normalize_provider_request_timeout(Some(1.0)), Ok(Some(1.0))); + assert_eq!( + normalize_provider_request_timeout(Some( + aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_SECS as f64, + )), + Ok(Some( + aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_SECS as f64, + )) + ); + assert!(normalize_provider_request_timeout(Some(0.0)).is_err()); + assert!(normalize_provider_request_timeout(Some( + aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_SECS as f64 + 1.0, + )) + .is_err()); + assert!(normalize_provider_request_timeout(Some(f64::NAN)).is_err()); + } + + #[test] + fn provider_stream_first_byte_timeout_keeps_its_protocol_range() { + assert_eq!( + normalize_provider_stream_first_byte_timeout(Some(1.0)), + Ok(Some(1.0)) + ); + assert_eq!( + normalize_provider_stream_first_byte_timeout(Some( + aether_contracts::MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_SECS as f64, + )), + Ok(Some( + aether_contracts::MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_SECS as f64, + )) + ); + assert!(normalize_provider_stream_first_byte_timeout(Some( + aether_contracts::MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_SECS as f64 + 1.0, + )) + .is_err()); + } +} 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 b83c7638d..25a63e891 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 @@ -107,16 +107,9 @@ pub(crate) async fn build_admin_create_provider_record( None => Some(2), }; let proxy = normalize_json_object(payload.proxy, "proxy")?; - let stream_first_byte_timeout_secs = match payload.stream_first_byte_timeout { - Some(value) if (1.0..=300.0).contains(&value) => Some(value), - Some(_) => return Err("stream_first_byte_timeout 必须是 1 到 300 之间的数字".to_string()), - None => None, - }; - let request_timeout_secs = match payload.request_timeout { - Some(value) if (1.0..=600.0).contains(&value) => Some(value), - Some(_) => return Err("request_timeout 必须是 1 到 600 之间的数字".to_string()), - None => None, - }; + let stream_first_byte_timeout_secs = + super::normalize_provider_stream_first_byte_timeout(payload.stream_first_byte_timeout)?; + let request_timeout_secs = super::normalize_provider_request_timeout(payload.request_timeout)?; let mut config_map = normalize_json_object(payload.config, "config")? .and_then(|value| value.as_object().cloned()) diff --git a/apps/aether-gateway/src/handlers/admin/provider/write/provider/template.rs b/apps/aether-gateway/src/handlers/admin/provider/write/provider/template.rs index 7fa260e11..dc432355d 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/write/provider/template.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/write/provider/template.rs @@ -1,11 +1,11 @@ use super::endpoint::{ build_admin_fixed_provider_endpoint_defaults, build_admin_fixed_provider_endpoint_record, + AdminFixedProviderEndpointDefaults, }; use crate::handlers::admin::request::AdminAppState; -use crate::provider_key_auth::provider_key_is_oauth_managed; use crate::GatewayError; use aether_data_contracts::repository::provider_catalog::{ - StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider, + StoredProviderCatalogEndpoint, StoredProviderCatalogProvider, }; use aether_provider_transport::provider_types::{ fixed_provider_template, FixedProviderEndpointTemplate, FixedProviderTemplate, @@ -36,6 +36,30 @@ struct FixedProviderEndpointMetadata { pub(crate) async fn reconcile_admin_fixed_provider_template_endpoints( state: &AdminAppState<'_>, provider: &StoredProviderCatalogProvider, +) -> Result<(), GatewayError> { + reconcile_admin_fixed_provider_template_endpoints_with_adoption_provider( + state, provider, provider, + ) + .await +} + +pub(crate) async fn reconcile_admin_fixed_provider_template_endpoints_after_update( + state: &AdminAppState<'_>, + previous_provider: &StoredProviderCatalogProvider, + updated_provider: &StoredProviderCatalogProvider, +) -> Result<(), GatewayError> { + reconcile_admin_fixed_provider_template_endpoints_with_adoption_provider( + state, + updated_provider, + previous_provider, + ) + .await +} + +async fn reconcile_admin_fixed_provider_template_endpoints_with_adoption_provider( + state: &AdminAppState<'_>, + provider: &StoredProviderCatalogProvider, + adoption_provider: &StoredProviderCatalogProvider, ) -> Result<(), GatewayError> { let Some(template) = state.fixed_provider_template(&provider.provider_type) else { return Ok(()); @@ -53,8 +77,9 @@ pub(crate) async fn reconcile_admin_fixed_provider_template_endpoints( match existing_endpoint { Some(existing_endpoint) => { matched_endpoint_ids.insert(existing_endpoint.id.clone()); - let updated = reconcile_fixed_provider_endpoint( + let updated = reconcile_fixed_provider_endpoint_with_adoption_provider( provider, + adoption_provider, existing_endpoint, template, endpoint_template, @@ -115,31 +140,6 @@ pub(crate) async fn reconcile_admin_fixed_provider_template_endpoints( Ok(()) } -pub(crate) async fn reconcile_admin_fixed_provider_template_keys( - state: &AdminAppState<'_>, - provider: &StoredProviderCatalogProvider, -) -> Result<(), GatewayError> { - let Some(_) = state.fixed_provider_template(&provider.provider_type) else { - return Ok(()); - }; - - let existing_keys = state - .list_provider_catalog_keys_by_provider_ids(std::slice::from_ref(&provider.id)) - .await?; - for existing_key in existing_keys { - let Some(updated_key) = reconcile_fixed_provider_key(provider, &existing_key) else { - continue; - }; - let Some(_) = state.update_provider_catalog_key(&updated_key).await? else { - return Err(GatewayError::Internal( - "provider catalog key writer unavailable".to_string(), - )); - }; - } - - Ok(()) -} - pub(crate) fn apply_admin_fixed_provider_endpoint_template_overrides( provider: &StoredProviderCatalogProvider, existing_endpoint: &StoredProviderCatalogEndpoint, @@ -156,8 +156,14 @@ pub(crate) fn apply_admin_fixed_provider_endpoint_template_overrides( let defaults = build_admin_fixed_provider_endpoint_defaults(provider, template, endpoint_template)?; - let mut metadata = fixed_provider_endpoint_metadata(existing_endpoint) - .unwrap_or_else(|| managed_fixed_provider_endpoint_metadata(template, endpoint_template)); + let mut metadata = fixed_provider_endpoint_metadata(existing_endpoint).unwrap_or_else(|| { + adopt_fixed_provider_endpoint_metadata( + existing_endpoint, + &defaults, + template, + endpoint_template, + ) + }); let mut overrides = metadata.overrides.clone(); sync_override_if_changed( @@ -250,17 +256,48 @@ pub(crate) fn apply_admin_fixed_provider_endpoint_template_overrides( Ok(()) } +#[cfg(test)] fn reconcile_fixed_provider_endpoint( provider: &StoredProviderCatalogProvider, existing_endpoint: &StoredProviderCatalogEndpoint, template: &FixedProviderTemplate, endpoint_template: &FixedProviderEndpointTemplate, +) -> Result { + reconcile_fixed_provider_endpoint_with_adoption_provider( + provider, + provider, + existing_endpoint, + template, + endpoint_template, + ) +} + +fn reconcile_fixed_provider_endpoint_with_adoption_provider( + provider: &StoredProviderCatalogProvider, + adoption_provider: &StoredProviderCatalogProvider, + existing_endpoint: &StoredProviderCatalogEndpoint, + template: &FixedProviderTemplate, + endpoint_template: &FixedProviderEndpointTemplate, ) -> Result { let defaults = build_admin_fixed_provider_endpoint_defaults(provider, template, endpoint_template)?; let mut updated = existing_endpoint.clone(); - let metadata = fixed_provider_endpoint_metadata(existing_endpoint) - .unwrap_or_else(|| managed_fixed_provider_endpoint_metadata(template, endpoint_template)); + let metadata = match fixed_provider_endpoint_metadata(existing_endpoint) { + Some(metadata) => metadata, + None => { + let adoption_defaults = build_admin_fixed_provider_endpoint_defaults( + adoption_provider, + template, + endpoint_template, + )?; + adopt_fixed_provider_endpoint_metadata( + existing_endpoint, + &adoption_defaults, + template, + endpoint_template, + ) + } + }; updated.api_format = defaults.api_format.clone(); updated.api_family = Some(defaults.api_family.clone()); @@ -455,6 +492,77 @@ fn managed_fixed_provider_endpoint_metadata( } } +fn adopt_fixed_provider_endpoint_metadata( + existing_endpoint: &StoredProviderCatalogEndpoint, + defaults: &AdminFixedProviderEndpointDefaults, + template: &FixedProviderTemplate, + endpoint_template: &FixedProviderEndpointTemplate, +) -> FixedProviderEndpointMetadata { + let mut metadata = managed_fixed_provider_endpoint_metadata(template, endpoint_template); + sync_override( + &mut metadata.overrides, + OVERRIDE_BASE_URL, + &existing_endpoint.base_url, + &defaults.base_url, + ); + sync_override( + &mut metadata.overrides, + OVERRIDE_CUSTOM_PATH, + &existing_endpoint.custom_path, + &defaults.custom_path, + ); + sync_override( + &mut metadata.overrides, + OVERRIDE_HEADER_RULES, + &existing_endpoint.header_rules, + &defaults.header_rules, + ); + sync_override( + &mut metadata.overrides, + OVERRIDE_BODY_RULES, + &existing_endpoint.body_rules, + &defaults.body_rules, + ); + sync_override( + &mut metadata.overrides, + OVERRIDE_MAX_RETRIES, + &existing_endpoint.max_retries, + &defaults.max_retries, + ); + sync_override( + &mut metadata.overrides, + OVERRIDE_IS_ACTIVE, + &existing_endpoint.is_active, + &defaults.is_active, + ); + sync_override( + &mut metadata.overrides, + OVERRIDE_PROXY, + &existing_endpoint.proxy, + &defaults.proxy, + ); + sync_override( + &mut metadata.overrides, + OVERRIDE_FORMAT_ACCEPTANCE_CONFIG, + &existing_endpoint.format_acceptance_config, + &defaults.format_acceptance_config, + ); + + let existing_config = endpoint_config_without_metadata(existing_endpoint.config.as_ref()); + for (key, desired) in fixed_provider_endpoint_config_defaults(endpoint_template) { + let Some(actual) = existing_config.get(&key) else { + continue; + }; + sync_override( + &mut metadata.overrides, + &config_override_key(&key), + actual, + &desired, + ); + } + metadata +} + fn upsert_fixed_provider_endpoint_metadata( endpoint: &mut StoredProviderCatalogEndpoint, metadata: &FixedProviderEndpointMetadata, @@ -513,22 +621,6 @@ fn current_unix_secs() -> u64 { .unwrap_or(0) } -fn reconcile_fixed_provider_key( - provider: &StoredProviderCatalogProvider, - existing_key: &StoredProviderCatalogKey, -) -> Option { - if !provider_key_is_oauth_managed(existing_key, &provider.provider_type) - || existing_key.api_formats.is_none() - { - return None; - } - - let mut updated = existing_key.clone(); - updated.api_formats = None; - updated.updated_at_unix_secs = Some(current_unix_secs()); - Some(updated) -} - fn sync_override(overrides: &mut BTreeSet, key: &str, actual: &T, desired: &T) where T: PartialEq, @@ -624,4 +716,68 @@ mod tests { .expect("endpoint should reconcile"); assert_eq!(reconciled.base_url, "http://127.0.0.1:18181/v1"); } + + #[test] + fn fixed_provider_endpoint_reconcile_adopts_existing_customization() { + let provider = sample_codex_provider(); + let template = fixed_provider_template("codex").expect("codex template should exist"); + let endpoint_template = template + .endpoints + .iter() + .find(|endpoint| endpoint.api_format == "openai:responses") + .expect("responses endpoint template should exist"); + + let mut existing = sample_codex_endpoint("http://127.0.0.1:18181/backend-api/codex"); + existing.is_active = false; + existing.max_retries = Some(9); + existing.proxy = Some(serde_json::json!({"url": "http://proxy.internal:8080"})); + existing.config = Some(serde_json::json!({ + "upstream_stream_policy": "force_non_stream", + "custom_transport_option": true + })); + + let reconciled = + reconcile_fixed_provider_endpoint(&provider, &existing, template, endpoint_template) + .expect("endpoint should reconcile"); + assert_eq!( + reconciled.base_url, + "http://127.0.0.1:18181/backend-api/codex" + ); + assert!(!reconciled.is_active); + assert_eq!(reconciled.max_retries, Some(9)); + assert_eq!( + reconciled.proxy, + Some(serde_json::json!({"url": "http://proxy.internal:8080"})) + ); + assert_eq!( + reconciled + .config + .as_ref() + .and_then(|value| value.get("upstream_stream_policy")), + Some(&serde_json::json!("force_non_stream")) + ); + assert_eq!( + reconciled + .config + .as_ref() + .and_then(|value| value.get("custom_transport_option")), + Some(&serde_json::json!(true)) + ); + let metadata = fixed_provider_endpoint_metadata(&reconciled) + .expect("fixed provider metadata should exist"); + for key in [ + "base_url", + "is_active", + "max_retries", + "proxy", + "config.upstream_stream_policy", + ] { + assert!(metadata.overrides.contains(key), "missing override {key}"); + } + + let second = + reconcile_fixed_provider_endpoint(&provider, &reconciled, template, endpoint_template) + .expect("endpoint should reconcile idempotently"); + assert_eq!(second, reconciled); + } } 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 cd3ec295e..3c9bc1272 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 @@ -202,21 +202,13 @@ pub(crate) async fn build_admin_update_provider_record( } if fields.contains("stream_first_byte_timeout") { - updated.stream_first_byte_timeout_secs = match payload.stream_first_byte_timeout { - Some(value) if (1.0..=300.0).contains(&value) => Some(value), - Some(_) => { - return Err("stream_first_byte_timeout 必须是 1 到 300 之间的数字".to_string()); - } - None => None, - }; + updated.stream_first_byte_timeout_secs = + super::normalize_provider_stream_first_byte_timeout(payload.stream_first_byte_timeout)?; } if fields.contains("request_timeout") { - updated.request_timeout_secs = match payload.request_timeout { - Some(value) if (1.0..=600.0).contains(&value) => Some(value), - Some(_) => return Err("request_timeout 必须是 1 到 600 之间的数字".to_string()), - None => None, - }; + updated.request_timeout_secs = + super::normalize_provider_request_timeout(payload.request_timeout)?; } if fields.contains("enable_format_conversion") { 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 0f31ca7bb..5690d1c28 100644 --- a/apps/aether-gateway/src/handlers/admin/request/provider/builders.rs +++ b/apps/aether-gateway/src/handlers/admin/request/provider/builders.rs @@ -1,5 +1,26 @@ use super::*; +fn validate_admin_endpoint_stream_policy( + api_format: &str, + config: Option<&serde_json::Value>, +) -> Result<(), String> { + if !crate::ai_serving::api_format_alias_matches(api_format, "openai:search") { + return Ok(()); + } + let requested = config + .and_then(serde_json::Value::as_object) + .and_then(|config| { + config + .get("upstream_stream_policy") + .or_else(|| config.get("upstreamStreamPolicy")) + .or_else(|| config.get("upstream_stream")) + }); + if requested.is_some_and(crate::handlers::public::admin_requested_force_stream) { + return Err("OpenAI Search 端点仅支持非流式上游请求".to_string()); + } + Ok(()) +} + impl<'a> AdminAppState<'a> { pub(crate) async fn build_admin_keys_grouped_by_format_payload( &self, @@ -232,6 +253,7 @@ impl<'a> AdminAppState<'a> { let (normalized_api_format, api_family, endpoint_kind) = 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())?; let base_url = normalize_admin_base_url(&payload.base_url)?; let existing_endpoints = self @@ -337,6 +359,13 @@ impl<'a> AdminAppState<'a> { &update_fields, )?; + if fields.contains("config") { + validate_admin_endpoint_stream_policy( + existing_endpoint.api_format.as_str(), + updated.config.as_ref(), + )?; + } + if provider_type == "codex" && crate::ai_serving::is_openai_responses_format(&existing_endpoint.api_format) { @@ -389,3 +418,40 @@ impl<'a> AdminAppState<'a> { Ok(updated) } } + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::validate_admin_endpoint_stream_policy; + + #[test] + fn search_endpoint_rejects_explicit_streaming_policy_for_all_config_keys() { + for (api_format, key, value) in [ + ( + "openai:search", + "upstream_stream_policy", + json!("force_stream"), + ), + ("openai:search", "upstreamStreamPolicy", json!(true)), + ("/v1/alpha/search", "upstream_stream", json!("sse")), + ] { + let config = json!({(key): value}); + assert!(validate_admin_endpoint_stream_policy(api_format, Some(&config),).is_err()); + } + } + + #[test] + fn search_endpoint_accepts_non_streaming_and_unrelated_config() { + assert!(validate_admin_endpoint_stream_policy( + "openai:search", + Some(&json!({"upstream_stream_policy": "force_non_stream"})), + ) + .is_ok()); + assert!(validate_admin_endpoint_stream_policy( + "openai:responses", + Some(&json!({"upstream_stream_policy": "force_stream"})), + ) + .is_ok()); + } +} diff --git a/apps/aether-gateway/src/handlers/proxy/mod.rs b/apps/aether-gateway/src/handlers/proxy/mod.rs index 7c02ea797..19c797c8c 100644 --- a/apps/aether-gateway/src/handlers/proxy/mod.rs +++ b/apps/aether-gateway/src/handlers/proxy/mod.rs @@ -84,6 +84,8 @@ const OPENAI_RESPONSES_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL: &str = "当前 OpenAI Responses 请求无法在本地执行：没有匹配到可用的执行路径"; const OPENAI_RESPONSES_COMPACT_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL: &str = "当前 OpenAI Responses Compact 请求无法在本地执行：没有匹配到可用的执行路径"; +const OPENAI_SEARCH_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL: &str = + "当前 OpenAI Search 请求无法在本地执行：没有匹配到可用的执行路径"; const OPENAI_VIDEO_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL: &str = "当前 OpenAI Video 请求无法在本地执行：没有匹配到可用的执行路径"; const CLAUDE_MESSAGES_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL: &str = @@ -381,8 +383,9 @@ async fn maybe_forward_public_request_to_tunnel_owner( serde_json::from_slice::(body.as_ref()).ok() }); let client_session_affinity = - crate::client_session_affinity::client_session_affinity_from_parts( - parts, + crate::client_session_affinity::client_session_affinity_from_api_request( + api_format, + &parts.headers, body_json.as_ref(), ); let Some(target) = crate::scheduler::affinity::read_cached_scheduler_affinity_target( @@ -452,7 +455,27 @@ async fn maybe_forward_public_request_to_tunnel_owner( owner.relay_base_url.trim_end_matches('/'), request_context.request_path_and_query() ); - let mut upstream_request = state.client.request(parts.method.clone(), owner_url); + let is_stream = + owner_forward_request_is_stream(parts, decision, buffered_body.unwrap_or(&empty_body)); + let transport_timeouts = + crate::provider_transport::resolve_transport_execution_timeouts(&transport); + let non_stream_timeout = + crate::execution_runtime::transport::resolve_non_stream_total_timeout_for_request( + is_stream, + &transport.endpoint.api_format, + transport_timeouts.as_ref(), + ); + let stream_first_byte_timeout = + crate::execution_runtime::transport::resolve_stream_first_byte_timeout_for_request( + is_stream, + transport_timeouts.as_ref(), + ); + let mut upstream_request = state + .owner_forward_client + .request(parts.method.clone(), owner_url); + if let Some(timeout) = non_stream_timeout { + upstream_request = upstream_request.timeout(timeout); + } for (name, value) in &parts.headers { if should_skip_request_header(name.as_str()) || name == http::header::HOST { continue; @@ -498,14 +521,15 @@ async fn maybe_forward_public_request_to_tunnel_owner( upstream_request.header(TRUSTED_AUTH_BALANCE_HEADER, balance_remaining.to_string()); } - let upstream_response = upstream_request - .body(buffered_body.cloned().unwrap_or_default()) - .send() - .await - .map_err(|err| GatewayError::UpstreamUnavailable { - trace_id: request_context.trace_id.clone(), - message: format!("owner gateway affinity forward failed: {err}"), - })?; + let upstream_response = crate::tunnel::send_owner_forward_request( + upstream_request.body(buffered_body.cloned().unwrap_or_default()), + stream_first_byte_timeout, + ) + .await + .map_err(|message| GatewayError::UpstreamUnavailable { + trace_id: request_context.trace_id.clone(), + message: format!("owner gateway affinity forward failed: {message}"), + })?; let mut response = build_sync_aware_affinity_forward_response( request_context, @@ -523,6 +547,29 @@ async fn maybe_forward_public_request_to_tunnel_owner( Ok(Some(response)) } +fn owner_forward_request_is_stream( + parts: &http::request::Parts, + decision: &GatewayControlDecision, + body_bytes: &Bytes, +) -> bool { + let Some(plan_kind) = + crate::ai_serving::api::resolve_execution_runtime_stream_plan_kind(parts, decision) + else { + return false; + }; + let Some((body_json, body_base64)) = + crate::ai_serving::api::parse_direct_request_body(parts, body_bytes) + else { + return false; + }; + crate::ai_serving::api::is_matching_stream_request( + plan_kind, + parts, + &body_json, + body_base64.as_deref(), + ) +} + fn upstream_response_is_sse(headers: &reqwest::header::HeaderMap) -> bool { headers .get(http::header::CONTENT_TYPE) @@ -2196,6 +2243,7 @@ fn local_execution_runtime_miss_route_label( "/v1/chat/completions" => "OpenAI Chat Completions", "/v1/responses" => "OpenAI Responses", "/v1/responses/compact" => "OpenAI Responses Compact", + "/v1/alpha/search" => "OpenAI Search", "/v1/messages" => "Claude Messages", path if path.starts_with("/v1/videos") => "OpenAI Video", path if path.starts_with("/upload/v1beta/files") || path.starts_with("/v1beta/files") => { @@ -2238,6 +2286,7 @@ fn local_execution_runtime_miss_route_detail( "/v1/responses/compact" => { Some(OPENAI_RESPONSES_COMPACT_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL) } + "/v1/alpha/search" => Some(OPENAI_SEARCH_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL), "/v1/messages" => Some(CLAUDE_MESSAGES_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL), path if path.starts_with("/v1/videos") => { Some(OPENAI_VIDEO_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL) @@ -2262,15 +2311,129 @@ mod tests { use super::{ api_key_remote_ip_allowed, buffer_and_normalize_request_body, diagnostic_is_auth_api_key_concurrency_limited, local_execution_runtime_miss_detail, - restore_redacted_stream_execution_response, restore_redacted_sync_execution_response, - GatewayControlDecision, LocalExecutionRuntimeMissDiagnostic, RequestBodyBufferError, - RequestBodyBufferPolicy, + owner_forward_request_is_stream, restore_redacted_stream_execution_response, + restore_redacted_sync_execution_response, GatewayControlDecision, + LocalExecutionRuntimeMissDiagnostic, RequestBodyBufferError, RequestBodyBufferPolicy, }; use axum::body::{to_bytes, Body, Bytes}; use axum::http::{header, HeaderMap, HeaderValue, Method, Response}; use serde_json::json; use tokio::sync::Semaphore; + #[test] + fn owner_forward_uses_search_protocol_timeout_semantics() { + let request = http::Request::builder() + .method(Method::POST) + .uri("/v1/alpha/search") + .header(header::CONTENT_TYPE, "application/json") + .body(()) + .expect("request should build"); + let (parts, _) = request.into_parts(); + let decision = GatewayControlDecision::synthetic( + "/v1/alpha/search", + Some("ai_public".to_string()), + Some("openai".to_string()), + Some("search".to_string()), + Some("openai:search".to_string()), + ); + let body = + Bytes::from_static(br#"{"model":"gpt-5.6-sol","input":"find docs","stream":true}"#); + let is_stream = owner_forward_request_is_stream(&parts, &decision, &body); + let timeouts = aether_contracts::ExecutionTimeouts { + total_ms: Some(aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_MS), + first_byte_ms: Some(10), + ..aether_contracts::ExecutionTimeouts::default() + }; + + assert!(!is_stream); + assert_eq!( + crate::execution_runtime::transport::resolve_non_stream_total_timeout_for_request( + is_stream, + "openai:search", + Some(&timeouts), + ), + Some(Duration::from_millis( + aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_MS + )) + ); + assert_eq!( + crate::execution_runtime::transport::resolve_stream_first_byte_timeout_for_request( + is_stream, + Some(&timeouts), + ), + None + ); + } + + #[test] + fn owner_forward_keeps_streaming_for_stream_capable_protocols() { + let chat_request = http::Request::builder() + .method(Method::POST) + .uri("/v1/chat/completions") + .header(header::CONTENT_TYPE, "application/json") + .body(()) + .expect("request should build"); + let (chat_parts, _) = chat_request.into_parts(); + let chat_decision = GatewayControlDecision::synthetic( + "/v1/chat/completions", + Some("ai_public".to_string()), + Some("openai".to_string()), + Some("chat".to_string()), + Some("openai:chat".to_string()), + ); + + assert!(owner_forward_request_is_stream( + &chat_parts, + &chat_decision, + &Bytes::from_static(br#"{"model":"gpt-5.6-sol","stream":true}"#), + )); + assert!(!owner_forward_request_is_stream( + &chat_parts, + &chat_decision, + &Bytes::from_static(br#"{"model":"gpt-5.6-sol","stream":false}"#), + )); + + let image_request = http::Request::builder() + .method(Method::POST) + .uri("/v1/images/generations") + .header(header::CONTENT_TYPE, "application/json") + .body(()) + .expect("request should build"); + let (image_parts, _) = image_request.into_parts(); + let image_decision = GatewayControlDecision::synthetic( + "/v1/images/generations", + Some("ai_public".to_string()), + Some("openai".to_string()), + Some("image".to_string()), + Some("openai:image".to_string()), + ); + assert!(owner_forward_request_is_stream( + &image_parts, + &image_decision, + &Bytes::from_static(br#"{"model":"gpt-image-1","stream":true}"#), + )); + + let compact_request = http::Request::builder() + .method(Method::POST) + .uri("/v1/responses/compact") + .header(header::CONTENT_TYPE, "application/json") + .body(()) + .expect("request should build"); + let (compact_parts, _) = compact_request.into_parts(); + let compact_decision = GatewayControlDecision::synthetic( + "/v1/responses/compact", + Some("ai_public".to_string()), + Some("openai".to_string()), + Some("responses:compact".to_string()), + Some("openai:responses:compact".to_string()), + ); + assert!(!owner_forward_request_is_stream( + &compact_parts, + &compact_decision, + &Bytes::from_static(br#"{"model":"gpt-5.6-sol","stream":true}"#), + )); + } + #[test] fn api_key_remote_ip_allows_unrestricted_keys() { let remote_ip = "203.0.113.10".parse().expect("valid ip"); diff --git a/apps/aether-gateway/src/handlers/public/ai_public.rs b/apps/aether-gateway/src/handlers/public/ai_public.rs index 23a0e1aa4..8ca6fc401 100644 --- a/apps/aether-gateway/src/handlers/public/ai_public.rs +++ b/apps/aether-gateway/src/handlers/public/ai_public.rs @@ -1,9 +1,8 @@ +use crate::ai_serving::normalize_openai_image_quality; use crate::async_task::CancelVideoTaskError; use crate::control::GatewayControlDecision; use crate::control::GatewayPublicRequestContext; -use crate::image_capabilities::{ - openai_image_gateway_max_generation_count, openai_image_gateway_max_generation_count_for_model, -}; +use crate::image_capabilities::openai_image_gateway_max_generation_count; use crate::{AppState, GatewayError}; use aether_data_contracts::repository::video_tasks::{ StoredVideoTask, VideoTaskQueryFilter, VideoTaskStatus, @@ -26,7 +25,7 @@ const OPENAI_IMAGE_PARTIAL_IMAGES_DETAIL: &str = const OPENAI_IMAGE_STYLE_DETAIL: &str = "当前 Codex 图片反代暂不支持 style 参数"; const OPENAI_IMAGE_RESPONSE_FORMAT_DETAIL: &str = "response_format 仅支持 url 或 b64_json"; const OPENAI_IMAGE_OUTPUT_FORMAT_DETAIL: &str = "output_format 仅支持 png、jpeg 或 webp"; -const OPENAI_IMAGE_QUALITY_DETAIL: &str = "quality 仅支持 low、medium、high、standard 或 hd"; +const OPENAI_IMAGE_QUALITY_DETAIL: &str = "quality 仅支持 auto、low、medium、high、standard 或 hd"; const OPENAI_IMAGE_BACKGROUND_DETAIL: &str = "background 仅支持 auto、opaque 或 transparent"; const OPENAI_IMAGE_MODERATION_DETAIL: &str = "moderation 仅支持 auto 或 low"; const OPENAI_IMAGE_INPUT_FIDELITY_DETAIL: &str = "input_fidelity 仅支持 low 或 high"; @@ -303,7 +302,7 @@ fn maybe_build_local_openai_request_validation_response( if validation .quality .as_deref() - .is_some_and(|value| !matches!(value, "low" | "medium" | "high" | "standard" | "hd")) + .is_some_and(|value| normalize_openai_image_quality(value).is_none()) { return Some(build_ai_public_error_response( http::StatusCode::BAD_REQUEST, @@ -366,8 +365,7 @@ fn openai_image_n_detail(max_generation_count: u64) -> String { } fn validate_openai_image_n(validation: &OpenAiImageValidationInput) -> Option { - let max_generation_count = - openai_image_gateway_max_generation_count_for_model(validation.model.as_deref()); + let max_generation_count = openai_image_gateway_max_generation_count(); validation .n .is_some_and(|value| value == 0 || value > max_generation_count) @@ -1761,7 +1759,7 @@ mod tests { } #[test] - fn image_validation_restricts_multi_image_count_to_grok_models() { + fn image_validation_applies_the_global_count_limit_before_model_mapping() { let openai_body = Bytes::from_static(br#"{"model":"gpt-image-2","prompt":"draw","n":2}"#); let openai_validation = parse_openai_image_validation_input( OpenAiImageOperation::Generate, @@ -1770,10 +1768,7 @@ mod tests { ) .expect("valid image payload should parse"); - assert_eq!( - validate_openai_image_n(&openai_validation).as_deref(), - Some("当前图片模型仅支持 n=1..1") - ); + assert!(validate_openai_image_n(&openai_validation).is_none()); let grok_body = Bytes::from_static(br#"{"model":"grok-imagine-image-lite","prompt":"draw","n":4}"#); @@ -1785,5 +1780,28 @@ mod tests { .expect("valid grok image payload should parse"); assert!(validate_openai_image_n(&grok_validation).is_none()); + + let alias_body = + Bytes::from_static(br#"{"model":"production-image-alias","prompt":"draw","n":10}"#); + let alias_validation = parse_openai_image_validation_input( + OpenAiImageOperation::Generate, + Some("application/json"), + &alias_body, + ) + .expect("valid image alias payload should parse"); + assert!(validate_openai_image_n(&alias_validation).is_none()); + + let excessive_body = + Bytes::from_static(br#"{"model":"production-image-alias","prompt":"draw","n":11}"#); + let excessive_validation = parse_openai_image_validation_input( + OpenAiImageOperation::Generate, + Some("application/json"), + &excessive_body, + ) + .expect("image payload should parse before count validation"); + assert_eq!( + validate_openai_image_n(&excessive_validation).as_deref(), + Some("当前图片反代仅支持 n=1..10") + ); } } diff --git a/apps/aether-gateway/src/handlers/public/support/models/responses.rs b/apps/aether-gateway/src/handlers/public/support/models/responses.rs index f38e58bbe..cb470ba19 100644 --- a/apps/aether-gateway/src/handlers/public/support/models/responses.rs +++ b/apps/aether-gateway/src/handlers/public/support/models/responses.rs @@ -89,6 +89,7 @@ pub(super) fn build_models_not_found_response(model_id: &str, api_format: &str) pub(super) fn build_empty_models_list_response(api_format: &str) -> Response { match api_format { + "openai:responses" => Json(json!({ "models": [] })).into_response(), "claude:messages" => Json(json!({ "data": [], "has_more": false, @@ -101,6 +102,10 @@ pub(super) fn build_empty_models_list_response(api_format: &str) -> Response) -> Response { + Json(json!({ "models": models })).into_response() +} + pub(super) fn build_openai_models_list_response( rows: &[StoredMinimalCandidateSelectionRow], ) -> Response { diff --git a/apps/aether-gateway/src/handlers/public/support/models/route.rs b/apps/aether-gateway/src/handlers/public/support/models/route.rs index b71977b94..6588fc50b 100644 --- a/apps/aether-gateway/src/handlers/public/support/models/route.rs +++ b/apps/aether-gateway/src/handlers/public/support/models/route.rs @@ -1,21 +1,24 @@ +use std::collections::{BTreeMap, BTreeSet}; use std::fmt::Debug; use std::future::Future; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use aether_data_contracts::repository::candidate_selection::StoredMinimalCandidateSelectionRow; use axum::{body::Body, response::Response}; +use serde_json::Value; use tokio::time::timeout; use tracing::warn; use super::models_responses::{ build_claude_model_detail_response, build_claude_models_list_response, - build_empty_models_list_response, build_gemini_model_detail_response, - build_gemini_models_list_response, build_models_auth_error_response, - build_models_not_found_response, build_openai_model_detail_response, - build_openai_models_list_response, + build_codex_models_list_response, build_empty_models_list_response, + build_gemini_model_detail_response, build_gemini_models_list_response, + build_models_auth_error_response, build_models_not_found_response, + build_openai_model_detail_response, build_openai_models_list_response, }; use super::models_shared::{ - filter_rows_for_models, models_api_format, models_detail_id, models_query_api_formats, + filter_eligible_model_rows, filter_rows_for_models, models_api_format, models_detail_id, + models_query_api_formats, }; use super::{query_param_value, AppState, GatewayPublicRequestContext}; @@ -23,6 +26,7 @@ use super::{query_param_value, AppState, GatewayPublicRequestContext}; const MODELS_ROUTE_READ_TIMEOUT: Duration = Duration::from_secs(5); #[cfg(test)] const MODELS_ROUTE_READ_TIMEOUT: Duration = Duration::from_millis(50); +const CODEX_MODELS_QUERY_API_FORMATS: &[&str] = &["openai:responses"]; async fn await_models_route_read(operation: &'static str, future: Fut) -> Option where @@ -72,7 +76,7 @@ fn build_models_read_fallback_response( } } -fn sort_and_dedup_model_rows( +fn sort_model_rows( mut rows: Vec, ) -> Vec { rows.sort_by(|left, right| { @@ -85,9 +89,15 @@ fn sort_and_dedup_model_rows( .then(left.key_id.cmp(&right.key_id)) .then(left.model_id.cmp(&right.model_id)) }); + rows +} + +fn sort_and_dedup_model_rows( + rows: Vec, +) -> Vec { let mut deduped = Vec::with_capacity(rows.len()); let mut last_model_name: Option = None; - for row in rows { + for row in sort_model_rows(rows) { if last_model_name.as_deref() == Some(row.global_model_name.as_str()) { continue; } @@ -97,22 +107,157 @@ fn sort_and_dedup_model_rows( deduped } +fn is_codex_models_api_format(api_format: &str) -> bool { + crate::ai_serving::normalize_api_format_alias(api_format) == "openai:responses" +} + +fn is_codex_provider_row(row: &StoredMinimalCandidateSelectionRow) -> bool { + row.provider_type.trim().eq_ignore_ascii_case("codex") +} + +fn codex_model_card_is_complete(card: &serde_json::Map) -> bool { + card.get("slug").and_then(Value::as_str).is_some() + && card.get("display_name").and_then(Value::as_str).is_some() + && card + .get("supported_reasoning_levels") + .and_then(Value::as_array) + .is_some() + && card.get("shell_type").and_then(Value::as_str).is_some() + && card.get("visibility").and_then(Value::as_str).is_some() + && card + .get("supported_in_api") + .and_then(Value::as_bool) + .is_some() + && card.get("priority").and_then(Value::as_i64).is_some() + && card + .get("base_instructions") + .and_then(Value::as_str) + .is_some() + && card + .get("supports_reasoning_summary_parameter") + .is_none_or(Value::is_boolean) + && card + .get("support_verbosity") + .and_then(Value::as_bool) + .is_some() + && card + .get("truncation_policy") + .and_then(Value::as_object) + .is_some() + && card + .get("supports_parallel_tool_calls") + .and_then(Value::as_bool) + .is_some() + && card + .get("experimental_supported_tools") + .and_then(Value::as_array) + .is_some() +} + +fn project_codex_model_card( + cached_models: &[Value], + source_model: &str, + global_model: &str, +) -> Option { + let mut card = cached_models + .iter() + .find(|model| { + model.get("id").and_then(Value::as_str) == Some(source_model) + || model.get("slug").and_then(Value::as_str) == Some(source_model) + })? + .as_object()? + .clone(); + if !codex_model_card_is_complete(&card) { + return None; + } + + card.remove("id"); + card.remove("api_formats"); + card.insert("slug".to_string(), Value::String(global_model.to_string())); + Some(Value::Object(card)) +} + +async fn load_codex_model_cards( + state: &AppState, + rows: &[StoredMinimalCandidateSelectionRow], +) -> Vec { + let cache_keys = rows + .iter() + .filter(|row| is_codex_provider_row(row)) + .map(|row| format!("upstream_models:{}:{}", row.provider_id, row.key_id)) + .collect::>() + .into_iter() + .collect::>(); + let cached_values = await_models_route_read( + "codex_models_cache", + state.runtime_state.kv_get_many(&cache_keys), + ) + .await + .unwrap_or_default(); + let cached_models_by_key = cache_keys + .into_iter() + .zip(cached_values) + .filter_map(|(key, raw)| { + let models = serde_json::from_str::>(raw.as_deref()?).ok()?; + Some((key, models)) + }) + .collect::>(); + + let mut seen_global_models = BTreeSet::new(); + let mut cards = Vec::new(); + for row in rows.iter().filter(|row| is_codex_provider_row(row)) { + if seen_global_models.contains(&row.global_model_name) { + continue; + } + let cache_key = format!("upstream_models:{}:{}", row.provider_id, row.key_id); + let Some(cached_models) = cached_models_by_key.get(&cache_key) else { + continue; + }; + let source_model = + aether_scheduler_core::select_provider_model_name(row, "openai:responses"); + let Some(card) = project_codex_model_card( + cached_models, + source_model.as_str(), + row.global_model_name.as_str(), + ) else { + continue; + }; + seen_global_models.insert(row.global_model_name.clone()); + cards.push(card); + } + cards +} + async fn list_model_rows_for_client_format( state: &AppState, api_format: &str, auth_snapshot: Option<&crate::data::auth::GatewayAuthApiKeySnapshot>, ) -> Option> { let mut collected = Vec::new(); - for query_format in models_query_api_formats(api_format) { + let query_api_formats = if is_codex_models_api_format(api_format) { + CODEX_MODELS_QUERY_API_FORMATS + } else { + models_query_api_formats(api_format) + }; + for query_format in query_api_formats { let rows = await_models_route_read( "candidate_selection_by_api_format", state.list_minimal_candidate_selection_rows_for_api_format(query_format), ) .await?; - let mut filtered = filter_rows_for_models(rows, auth_snapshot, query_format); + let mut filtered = if is_codex_models_api_format(api_format) { + filter_eligible_model_rows(rows, auth_snapshot, query_format) + } else { + filter_rows_for_models(rows, auth_snapshot, query_format) + }; collected.append(&mut filtered); } - Some(sort_and_dedup_model_rows(collected)) + if is_codex_models_api_format(api_format) { + collected.retain(is_codex_provider_row); + Some(sort_model_rows(collected)) + } else { + Some(sort_and_dedup_model_rows(collected)) + } } async fn list_model_rows_for_client_format_and_global_model( @@ -190,6 +335,10 @@ pub(super) async fn maybe_build_local_models_route_response( if rows.is_empty() { return Some(build_empty_models_list_response(api_format)); } + if is_codex_models_api_format(api_format) { + let models = load_codex_model_cards(state, &rows).await; + return Some(build_codex_models_list_response(models)); + } let response = match api_format { "claude:messages" => { let before_id = query_param_value( diff --git a/apps/aether-gateway/src/handlers/public/support/models/shared.rs b/apps/aether-gateway/src/handlers/public/support/models/shared.rs index d57cf6500..1369c5a11 100644 --- a/apps/aether-gateway/src/handlers/public/support/models/shared.rs +++ b/apps/aether-gateway/src/handlers/public/support/models/shared.rs @@ -194,13 +194,12 @@ fn row_exposes_global_model_for_models( false } -pub(super) fn filter_rows_for_models( +pub(super) fn filter_eligible_model_rows( rows: Vec, auth_snapshot: Option<&crate::data::auth::GatewayAuthApiKeySnapshot>, api_format: &str, ) -> Vec { - let mut filtered = rows - .into_iter() + rows.into_iter() .filter(|row| { auth_snapshot_allows_provider_for_models( auth_snapshot, @@ -211,7 +210,15 @@ pub(super) fn filter_rows_for_models( }) .filter(|row| auth_snapshot_allows_model_for_models(auth_snapshot, &row.global_model_name)) .filter(|row| row_exposes_global_model_for_models(row, api_format)) - .collect::>(); + .collect() +} + +pub(super) fn filter_rows_for_models( + rows: Vec, + auth_snapshot: Option<&crate::data::auth::GatewayAuthApiKeySnapshot>, + api_format: &str, +) -> Vec { + let mut filtered = filter_eligible_model_rows(rows, auth_snapshot, api_format); filtered.sort_by(|left, right| left.global_model_name.cmp(&right.global_model_name)); let mut deduped = Vec::new(); let mut last_model_name: Option = None; diff --git a/apps/aether-gateway/src/handlers/public/support/user_me_usage.rs b/apps/aether-gateway/src/handlers/public/support/user_me_usage.rs index 0977b1f39..5c63cef02 100644 --- a/apps/aether-gateway/src/handlers/public/support/user_me_usage.rs +++ b/apps/aether-gateway/src/handlers/public/support/user_me_usage.rs @@ -220,14 +220,7 @@ fn users_me_usage_api_format_defaults_to_non_stream(item: &StoredRequestUsageAud let Some(value) = api_format else { return false; }; - matches!( - crate::ai_serving::normalize_api_format_alias(value).as_str(), - "openai:chat" - | "openai:responses" - | "openai:responses:compact" - | "openai:image" - | "claude:messages" - ) + crate::ai_serving::api_format_defaults_to_non_stream(value) } fn users_me_usage_request_body_implies_default_non_stream(item: &StoredRequestUsageAudit) -> bool { @@ -515,6 +508,9 @@ fn build_users_me_usage_record_payload( if let Some(service_tier) = item.provider_service_tier() { payload["service_tier"] = json!(service_tier); } + if let Some(actual_service_tier) = item.provider_actual_service_tier() { + payload["actual_service_tier"] = json!(actual_service_tier); + } if include_actual_cost { payload["actual_cost"] = json!(round_to(item.actual_total_cost_usd, 6)); payload["rate_multiplier"] = json!(rate_multiplier); @@ -582,6 +578,9 @@ fn build_users_me_usage_active_payload(item: &StoredRequestUsageAudit) -> serde_ if let Some(service_tier) = item.provider_service_tier() { payload["service_tier"] = json!(service_tier); } + if let Some(actual_service_tier) = item.provider_actual_service_tier() { + payload["actual_service_tier"] = json!(actual_service_tier); + } payload } @@ -1840,6 +1839,27 @@ mod tests { assert_eq!(active_payload["client_is_stream"], false); } + #[test] + fn user_usage_stream_defaults_to_non_stream_for_openai_search() { + let item = StoredRequestUsageAudit { + is_stream: false, + api_format: Some("openai:search".to_string()), + request_body: Some(json!({ + "id": "session-search-1", + "model": "gpt-5.6-sol", + "input": "current documentation" + })), + ..sample_usage("completed") + }; + + assert!(!users_me_usage_client_is_stream(&item)); + + let record_payload = + build_users_me_usage_record_payload(&item, false, &BTreeMap::new(), false); + assert_eq!(record_payload["client_requested_stream"], false); + assert_eq!(record_payload["client_is_stream"], false); + } + #[test] fn user_usage_upstream_stream_prefers_request_metadata_flag() { let item = StoredRequestUsageAudit { diff --git a/apps/aether-gateway/src/handlers/shared/catalog.rs b/apps/aether-gateway/src/handlers/shared/catalog.rs index d34e14848..8e21c272a 100644 --- a/apps/aether-gateway/src/handlers/shared/catalog.rs +++ b/apps/aether-gateway/src/handlers/shared/catalog.rs @@ -38,7 +38,7 @@ pub(crate) fn provider_catalog_key_supports_format( } formats .iter() - .any(|candidate| crate::ai_serving::api_format_alias_matches(candidate, api_format)) + .any(|candidate| crate::ai_serving::api_format_permission_covers(candidate, api_format)) } pub(crate) fn decrypt_catalog_secret_with_fallbacks( @@ -2732,6 +2732,25 @@ mod tests { .expect("key transport should build") } + #[test] + fn responses_key_scope_covers_search_in_one_direction() { + let mut responses_key = sample_catalog_key(); + responses_key.api_formats = Some(json!(["openai:responses"])); + assert!(provider_catalog_key_supports_format( + &responses_key, + "codex", + "openai:search", + )); + + let mut search_key = sample_catalog_key(); + search_key.api_formats = Some(json!(["openai:search"])); + assert!(!provider_catalog_key_supports_format( + &search_key, + "codex", + "openai:responses", + )); + } + #[test] fn masked_catalog_api_key_handles_unicode_plaintext_without_panicking() { let state = AppState::new().expect("gateway should build"); diff --git a/apps/aether-gateway/src/image_capabilities.rs b/apps/aether-gateway/src/image_capabilities.rs index bb1491260..21dc5f183 100644 --- a/apps/aether-gateway/src/image_capabilities.rs +++ b/apps/aether-gateway/src/image_capabilities.rs @@ -1,41 +1,51 @@ use crate::ai_serving::OpenAiImageNormalizeOptions; -const DEFAULT_OPENAI_IMAGE_MAX_GENERATION_COUNT: u64 = 1; +const DEFAULT_IMAGE_MAX_GENERATION_COUNT: u64 = 1; +const OPENAI_IMAGE_MAX_GENERATION_COUNT: u64 = 10; const GROK_OPENAI_IMAGE_MAX_GENERATION_COUNT: u64 = 4; pub(crate) fn openai_image_gateway_max_generation_count() -> u64 { - GROK_OPENAI_IMAGE_MAX_GENERATION_COUNT -} - -pub(crate) fn openai_image_gateway_max_generation_count_for_model(model: Option<&str>) -> u64 { - if model.is_some_and(is_grok_openai_image_model) { - GROK_OPENAI_IMAGE_MAX_GENERATION_COUNT - } else { - DEFAULT_OPENAI_IMAGE_MAX_GENERATION_COUNT - } + OPENAI_IMAGE_MAX_GENERATION_COUNT } pub(crate) fn openai_image_provider_max_generation_count(provider_type: &str) -> u64 { if provider_type.trim().eq_ignore_ascii_case("grok") { GROK_OPENAI_IMAGE_MAX_GENERATION_COUNT + } else if matches!( + provider_type.trim().to_ascii_lowercase().as_str(), + "openai" | "codex" + ) { + OPENAI_IMAGE_MAX_GENERATION_COUNT } else { - DEFAULT_OPENAI_IMAGE_MAX_GENERATION_COUNT + DEFAULT_IMAGE_MAX_GENERATION_COUNT } } +pub(crate) fn openai_image_provider_max_generation_count_for_model( + provider_type: &str, + provider_model: Option<&str>, +) -> u64 { + let provider_limit = openai_image_provider_max_generation_count(provider_type); + provider_model.map_or(provider_limit, |model| { + if is_dall_e_3_model(model) { + 1 + } else { + provider_limit + } + }) +} + pub(crate) fn openai_image_normalize_options_for_provider( provider_type: &str, + provider_model: Option<&str>, ) -> OpenAiImageNormalizeOptions { OpenAiImageNormalizeOptions::with_max_generation_count( - openai_image_provider_max_generation_count(provider_type), + openai_image_provider_max_generation_count_for_model(provider_type, provider_model), ) } -fn is_grok_openai_image_model(model: &str) -> bool { - model - .trim() - .to_ascii_lowercase() - .contains("grok-imagine-image") +fn is_dall_e_3_model(model: &str) -> bool { + model.trim().to_ascii_lowercase().starts_with("dall-e-3") } #[cfg(test)] @@ -43,18 +53,19 @@ mod tests { use super::*; #[test] - fn grok_owns_gateway_wide_image_generation_count_ceiling() { - assert_eq!(openai_image_gateway_max_generation_count(), 4); + fn image_count_capabilities_follow_provider_and_model_contracts() { + assert_eq!(openai_image_gateway_max_generation_count(), 10); + assert_eq!(openai_image_provider_max_generation_count("grok"), 4); + assert_eq!(openai_image_provider_max_generation_count("openai"), 10); + assert_eq!(openai_image_provider_max_generation_count("codex"), 10); + assert_eq!(openai_image_provider_max_generation_count("custom"), 1); assert_eq!( - openai_image_gateway_max_generation_count_for_model(Some("grok-imagine-image-lite")), - 4 - ); - assert_eq!( - openai_image_gateway_max_generation_count_for_model(Some("gpt-image-2")), + openai_image_provider_max_generation_count_for_model("openai", Some("dall-e-3")), 1 ); - assert_eq!(openai_image_gateway_max_generation_count_for_model(None), 1); - assert_eq!(openai_image_provider_max_generation_count("grok"), 4); - assert_eq!(openai_image_provider_max_generation_count("openai"), 1); + assert_eq!( + openai_image_normalize_options_for_provider("openai", Some("dall-e-3")), + OpenAiImageNormalizeOptions::with_max_generation_count(1) + ); } } diff --git a/apps/aether-gateway/src/maintenance/mod.rs b/apps/aether-gateway/src/maintenance/mod.rs index 98facca26..5cb50439b 100644 --- a/apps/aether-gateway/src/maintenance/mod.rs +++ b/apps/aether-gateway/src/maintenance/mod.rs @@ -12,8 +12,8 @@ pub(crate) use runtime::{ restore_proxy_upgrade_rollout_skipped_nodes, retry_proxy_upgrade_rollout_node, run_admin_system_cleanup_once, run_manual_usage_cleanup_once, skip_proxy_upgrade_rollout_node, spawn_account_self_check_worker, spawn_audit_cleanup_worker, spawn_db_maintenance_worker, - spawn_gemini_file_mapping_cleanup_worker, spawn_oauth_token_refresh_worker, - spawn_pending_cleanup_worker, spawn_pool_monitor_worker, + spawn_fixed_provider_reconciliation_task, spawn_gemini_file_mapping_cleanup_worker, + spawn_oauth_token_refresh_worker, spawn_pending_cleanup_worker, spawn_pool_monitor_worker, spawn_pool_quota_probe_replenish_for_request, spawn_pool_quota_probe_worker, spawn_pool_score_rebuild_worker, spawn_provider_checkin_worker, spawn_provider_quota_alert_worker, spawn_proxy_node_metrics_cleanup_worker, diff --git a/apps/aether-gateway/src/maintenance/runtime.rs b/apps/aether-gateway/src/maintenance/runtime.rs index f2dfd9c0b..b513c2a82 100644 --- a/apps/aether-gateway/src/maintenance/runtime.rs +++ b/apps/aether-gateway/src/maintenance/runtime.rs @@ -16,6 +16,8 @@ mod cleanup_runs; mod config; #[path = "runtime/db_maintenance.rs"] mod db_maintenance; +#[path = "runtime/fixed_provider_reconciliation.rs"] +mod fixed_provider_reconciliation; #[path = "runtime/oauth_token_refresh.rs"] mod oauth_token_refresh; #[path = "runtime/pending_cleanup.rs"] @@ -71,6 +73,9 @@ pub(crate) use cleanup_runs::{ }; use config::*; use db_maintenance::*; +pub(crate) use fixed_provider_reconciliation::{ + perform_fixed_provider_reconciliation_once, spawn_fixed_provider_reconciliation_task, +}; pub(crate) use oauth_token_refresh::{ perform_oauth_token_refresh_once, OAuthTokenRefreshRunSummary, }; diff --git a/apps/aether-gateway/src/maintenance/runtime/fixed_provider_reconciliation.rs b/apps/aether-gateway/src/maintenance/runtime/fixed_provider_reconciliation.rs new file mode 100644 index 000000000..9ebf13544 --- /dev/null +++ b/apps/aether-gateway/src/maintenance/runtime/fixed_provider_reconciliation.rs @@ -0,0 +1,301 @@ +use std::time::Duration; + +use tracing::warn; + +use crate::admin_api::{reconcile_admin_fixed_provider_template_endpoints, AdminAppState}; +use crate::task_runtime::{ + spawn_fire_and_forget, task_definition, TASK_KEY_FIXED_PROVIDER_RECONCILIATION, +}; +use crate::{AppState, GatewayError}; + +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, +) -> Result { + if !state.has_provider_catalog_data_reader() || !state.has_provider_catalog_data_writer() { + return Ok(false); + } + + let Some(lock) = state + .runtime_state + .lock_try_acquire( + FIXED_PROVIDER_RECONCILIATION_LOCK_KEY, + state.tunnel.local_instance_id(), + FIXED_PROVIDER_RECONCILIATION_LOCK_TTL, + ) + .await + .map_err(|error| GatewayError::Internal(error.to_string()))? + else { + return Ok(false); + }; + + let result = reconcile_fixed_provider_templates(state).await; + if let Err(error) = state.runtime_state.lock_release(&lock).await { + warn!( + event_name = "fixed_provider_reconciliation_lock_release_failed", + log_type = "ops", + error = ?error, + "gateway fixed provider reconciliation lock release failed" + ); + } + result.map(|()| true) +} + +async fn reconcile_fixed_provider_templates(state: &AppState) -> Result<(), GatewayError> { + let providers = state.list_provider_catalog_providers(false).await?; + 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() + { + continue; + } + if let Err(error) = + reconcile_admin_fixed_provider_template_endpoints(&admin_state, provider).await + { + failures.push(format!( + "provider {} endpoint reconciliation failed: {error:?}", + provider.id, + )); + continue; + } + } + if !failures.is_empty() { + return Err(GatewayError::Internal(failures.join("; "))); + } + Ok(()) +} + +async fn perform_fixed_provider_reconciliation_with_retry(state: &AppState) { + let max_attempts = task_definition(TASK_KEY_FIXED_PROVIDER_RECONCILIATION) + .map(|definition| definition.retry_policy.max_attempts) + .unwrap_or(1) + .max(1); + for attempt in 1..=max_attempts { + match perform_fixed_provider_reconciliation_once(state).await { + Ok(_) => return, + Err(error) if attempt < max_attempts => { + warn!( + event_name = "fixed_provider_reconciliation_retrying", + log_type = "ops", + attempt, + max_attempts, + error = ?error, + "gateway fixed provider reconciliation will retry" + ); + tokio::time::sleep(FIXED_PROVIDER_RECONCILIATION_RETRY_DELAY).await; + } + Err(error) => { + warn!( + event_name = "fixed_provider_reconciliation_failed", + log_type = "ops", + attempt, + max_attempts, + error = ?error, + "gateway fixed provider reconciliation failed" + ); + return; + } + } + } +} + +pub(crate) fn spawn_fixed_provider_reconciliation_task( + state: AppState, +) -> Option> { + if !state.has_provider_catalog_data_reader() || !state.has_provider_catalog_data_writer() { + return None; + } + + Some(spawn_fire_and_forget( + TASK_KEY_FIXED_PROVIDER_RECONCILIATION, + async move { + perform_fixed_provider_reconciliation_with_retry(&state).await; + }, + )) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository; + use aether_data_contracts::repository::provider_catalog::{ + ProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey, + StoredProviderCatalogProvider, + }; + use serde_json::json; + + use super::{ + perform_fixed_provider_reconciliation_once, FIXED_PROVIDER_RECONCILIATION_LOCK_KEY, + FIXED_PROVIDER_RECONCILIATION_LOCK_TTL, + }; + use crate::data::GatewayDataState; + use crate::AppState; + + #[tokio::test] + async fn fixed_provider_reconciliation_respects_runtime_singleton_lock() { + let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( + vec![], + vec![], + vec![], + )); + let state = AppState::new() + .expect("gateway state should build") + .with_data_state_for_tests( + GatewayDataState::with_provider_catalog_repository_for_tests(repository), + ); + let lock = state + .runtime_state + .lock_try_acquire( + FIXED_PROVIDER_RECONCILIATION_LOCK_KEY, + "another-gateway", + FIXED_PROVIDER_RECONCILIATION_LOCK_TTL, + ) + .await + .expect("runtime lock should be available") + .expect("runtime lock should be acquired"); + + assert!(!perform_fixed_provider_reconciliation_once(&state) + .await + .expect("locked reconciliation should skip")); + + assert!(state + .runtime_state + .lock_release(&lock) + .await + .expect("runtime lock should release")); + assert!(perform_fixed_provider_reconciliation_once(&state) + .await + .expect("unlocked reconciliation should run")); + } + + #[tokio::test] + async fn fixed_provider_reconciliation_preserves_existing_endpoint_and_is_idempotent() { + let mut provider = StoredProviderCatalogProvider::new( + "provider-codex".to_string(), + "Codex".to_string(), + None, + "codex".to_string(), + ) + .expect("provider should build"); + provider.is_active = false; + provider.max_retries = Some(2); + + let mut responses = StoredProviderCatalogEndpoint::new( + "endpoint-codex-responses".to_string(), + provider.id.clone(), + "openai:responses".to_string(), + Some("openai".to_string()), + Some("cli".to_string()), + false, + ) + .expect("endpoint should build") + .with_transport_fields( + "http://127.0.0.1:18181/backend-api/codex".to_string(), + None, + None, + Some(9), + None, + Some(json!({"upstream_stream_policy": "force_non_stream"})), + None, + Some(json!({"url": "http://proxy.internal:8080"})), + ) + .expect("endpoint transport should build"); + responses.updated_at_unix_secs = Some(100); + + let mut key = StoredProviderCatalogKey::new( + "key-codex".to_string(), + provider.id.clone(), + "oauth".to_string(), + "oauth".to_string(), + None, + true, + ) + .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(), + None, + "claude_code".to_string(), + ) + .expect("unrelated fixed provider should build"); + + let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( + vec![provider, unrelated_fixed_provider], + vec![responses], + vec![key], + )); + 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 first_endpoints = repository + .list_endpoints_by_provider_ids(&["provider-codex".to_string()]) + .await + .expect("endpoints should list"); + assert_eq!(first_endpoints.len(), 4); + let responses = first_endpoints + .iter() + .find(|endpoint| endpoint.api_format == "openai:responses") + .expect("responses endpoint should exist"); + assert_eq!( + responses.base_url, + "http://127.0.0.1:18181/backend-api/codex" + ); + assert!(!responses.is_active); + assert_eq!(responses.max_retries, Some(9)); + assert_eq!( + responses.proxy, + Some(json!({"url": "http://proxy.internal:8080"})) + ); + assert_eq!( + responses + .config + .as_ref() + .and_then(|value| value.get("upstream_stream_policy")), + Some(&json!("force_non_stream")) + ); + assert!(first_endpoints + .iter() + .any(|endpoint| endpoint.api_format == "openai:search")); + let keys = repository + .list_keys_by_provider_ids(&["provider-codex".to_string()]) + .await + .expect("keys should list"); + 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()]) + .await + .expect("unrelated endpoints should list") + .is_empty()); + + assert!(perform_fixed_provider_reconciliation_once(&state) + .await + .expect("second reconciliation should run")); + let second_endpoints = repository + .list_endpoints_by_provider_ids(&["provider-codex".to_string()]) + .await + .expect("endpoints should list again"); + assert_eq!(second_endpoints, first_endpoints); + } +} diff --git a/apps/aether-gateway/src/maintenance/runtime/tests.rs b/apps/aether-gateway/src/maintenance/runtime/tests.rs index 673a5d57d..70ee7e514 100644 --- a/apps/aether-gateway/src/maintenance/runtime/tests.rs +++ b/apps/aether-gateway/src/maintenance/runtime/tests.rs @@ -21,7 +21,8 @@ use super::{ pending_cleanup_timeout_minutes, plan_pending_cleanup_batch, provider_checkin_schedule, proxy_node_metrics_cleanup_settings, record_proxy_upgrade_traffic_success, run_db_maintenance_with, run_proxy_upgrade_rollout_once, spawn_account_self_check_worker, - spawn_audit_cleanup_worker, spawn_db_maintenance_worker, spawn_oauth_token_refresh_worker, + spawn_audit_cleanup_worker, spawn_db_maintenance_worker, + spawn_fixed_provider_reconciliation_task, spawn_oauth_token_refresh_worker, spawn_pending_cleanup_worker, spawn_pool_monitor_worker, spawn_pool_quota_probe_worker, spawn_provider_checkin_worker, spawn_proxy_node_stale_cleanup_worker, spawn_proxy_upgrade_rollout_worker, spawn_stats_aggregation_worker, @@ -74,6 +75,14 @@ async fn spawn_oauth_token_refresh_worker_skips_when_provider_catalog_unavailabl assert!(spawn_oauth_token_refresh_worker(state).is_none()); } +#[tokio::test] +async fn spawn_fixed_provider_reconciliation_task_skips_when_provider_catalog_unavailable() { + let state = AppState::new() + .expect("gateway state should build") + .with_data_state_for_tests(GatewayDataState::disabled()); + assert!(spawn_fixed_provider_reconciliation_task(state).is_none()); +} + #[tokio::test] async fn spawn_proxy_upgrade_rollout_worker_skips_when_system_config_unavailable() { let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![])); diff --git a/apps/aether-gateway/src/model_fetch/runtime.rs b/apps/aether-gateway/src/model_fetch/runtime.rs index e70848b42..2e7e9826b 100644 --- a/apps/aether-gateway/src/model_fetch/runtime.rs +++ b/apps/aether-gateway/src/model_fetch/runtime.rs @@ -3,13 +3,15 @@ use std::time::Duration; use std::time::{SystemTime, UNIX_EPOCH}; use aether_data_contracts::repository::provider_catalog::{ - StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider, + ProviderCatalogUpstreamMetadataNamespaceUpdate, StoredProviderCatalogEndpoint, + StoredProviderCatalogKey, StoredProviderCatalogProvider, }; use aether_model_fetch::{ - apply_model_filters, fetch_models_from_transports, json_string_list, merge_upstream_metadata, - model_fetch_interval_minutes, model_fetch_startup_delay_seconds, model_fetch_startup_enabled, - preset_models_for_provider, selected_models_fetch_endpoints, - sync_provider_model_whitelist_associations, ModelFetchAssociationStore, ModelFetchRunSummary, + apply_model_filters, fetch_models_from_transports, json_string_list, + model_catalog_upstream_metadata, model_fetch_interval_minutes, + model_fetch_startup_delay_seconds, model_fetch_startup_enabled, preset_models_for_provider, + selected_models_fetch_endpoints, sync_provider_model_whitelist_associations, + upstream_metadata_namespace_updates, ModelFetchAssociationStore, ModelFetchRunSummary, }; use serde_json::{json, Value}; use tracing::{debug, info, warn}; @@ -239,8 +241,16 @@ async fn fetch_and_persist_key_models( json_string_list(target.key.model_include_patterns.as_ref()), json_string_list(target.key.model_exclude_patterns.as_ref()), ); - persist_key_fetch_success(state, &target.key, now_unix_secs, &filtered_models, None) - .await?; + let upstream_metadata = + model_catalog_upstream_metadata(&target.provider.provider_type, &models); + persist_key_fetch_success( + state, + &target.key, + now_unix_secs, + &filtered_models, + upstream_metadata.as_ref(), + ) + .await?; state .write_upstream_models_cache(&target.provider.id, &target.key.id, &models) .await; @@ -328,7 +338,6 @@ async fn fetch_and_persist_key_models( json_string_list(target.key.model_include_patterns.as_ref()), json_string_list(target.key.model_exclude_patterns.as_ref()), ); - persist_key_fetch_success( state, &target.key, @@ -352,11 +361,15 @@ async fn persist_key_fetch_failure( now_unix_secs: u64, error: String, ) -> Result<(), GatewayError> { - let mut updated = key.clone(); - updated.last_models_fetch_at_unix_secs = Some(now_unix_secs); - updated.last_models_fetch_error = Some(error); - updated.updated_at_unix_secs = Some(now_unix_secs); - state.update_provider_catalog_key(&updated).await?; + state + .update_provider_catalog_key_model_fetch_state( + &key.id, + key.allowed_models.as_ref(), + Some(now_unix_secs), + Some(&error), + Some(now_unix_secs), + ) + .await?; Ok(()) } @@ -367,22 +380,33 @@ async fn persist_key_fetch_success( allowed_models: &[String], upstream_metadata: Option<&Value>, ) -> Result<(), GatewayError> { - let mut updated = key.clone(); - updated.allowed_models = if allowed_models.is_empty() { + let allowed_models = if allowed_models.is_empty() { None } else { Some(json!(allowed_models)) }; - if let Some(upstream_metadata) = upstream_metadata { - updated.upstream_metadata = Some(merge_upstream_metadata( - updated.upstream_metadata.as_ref(), - upstream_metadata, - )); - } - updated.last_models_fetch_at_unix_secs = Some(now_unix_secs); - updated.last_models_fetch_error = None; - updated.updated_at_unix_secs = Some(now_unix_secs); - state.update_provider_catalog_key(&updated).await?; + let upstream_metadata_updates = upstream_metadata + .map(|upstream_metadata| { + upstream_metadata_namespace_updates(key.upstream_metadata.as_ref(), upstream_metadata) + .into_iter() + .map( + |(namespace, value)| ProviderCatalogUpstreamMetadataNamespaceUpdate { + namespace, + value, + }, + ) + .collect::>() + }) + .unwrap_or_default(); + state + .update_provider_catalog_key_model_fetch_success( + &key.id, + allowed_models.as_ref(), + now_unix_secs, + &upstream_metadata_updates, + Some(now_unix_secs), + ) + .await?; Ok(()) } @@ -402,7 +426,8 @@ mod tests { StoredAdminProviderModel, UpsertAdminProviderModelRecord, }; use aether_data_contracts::repository::provider_catalog::{ - StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider, + ProviderCatalogUpstreamMetadataNamespaceUpdate, StoredProviderCatalogEndpoint, + StoredProviderCatalogKey, StoredProviderCatalogProvider, }; use aether_model_fetch::{ build_models_fetch_execution_plan, ModelFetchAssociationStore, ModelFetchTransportRuntime, @@ -428,6 +453,7 @@ mod tests { execution_results: Arc>>, executed_plans: Arc>>, cached_models: Arc>>>, + upstream_metadata_updates: Arc)>>>, } impl TestState { @@ -446,6 +472,7 @@ mod tests { execution_results: Arc::new(Mutex::new(VecDeque::from(execution_results))), executed_plans: Arc::new(Mutex::new(Vec::new())), cached_models: Arc::new(Mutex::new(HashMap::new())), + upstream_metadata_updates: Arc::new(Mutex::new(Vec::new())), } } @@ -615,15 +642,63 @@ mod tests { )) } - async fn update_provider_catalog_key( + async fn update_provider_catalog_key_model_fetch_state( &self, - key: &StoredProviderCatalogKey, + key_id: &str, + allowed_models: Option<&Value>, + last_models_fetch_at_unix_secs: Option, + last_models_fetch_error: Option<&str>, + updated_at_unix_secs: Option, ) -> Result<(), GatewayError> { let mut keys = self.keys.lock().expect("keys mutex"); - let Some(slot) = keys.iter_mut().find(|item| item.id == key.id) else { + let Some(key) = keys.iter_mut().find(|item| item.id == key_id) else { return Err(GatewayError::Internal("key not found".to_string())); }; - *slot = key.clone(); + key.allowed_models = allowed_models.cloned(); + key.last_models_fetch_at_unix_secs = last_models_fetch_at_unix_secs; + key.last_models_fetch_error = last_models_fetch_error.map(str::to_string); + key.updated_at_unix_secs = updated_at_unix_secs; + Ok(()) + } + + async fn update_provider_catalog_key_model_fetch_success( + &self, + key_id: &str, + allowed_models: Option<&Value>, + last_models_fetch_at_unix_secs: u64, + upstream_metadata_updates: &[ProviderCatalogUpstreamMetadataNamespaceUpdate], + updated_at_unix_secs: Option, + ) -> Result<(), GatewayError> { + let mut keys = self.keys.lock().expect("keys mutex"); + let Some(key) = keys.iter_mut().find(|key| key.id == key_id) else { + return Err(GatewayError::Internal("key not found".to_string())); + }; + key.allowed_models = allowed_models.cloned(); + key.last_models_fetch_at_unix_secs = Some(last_models_fetch_at_unix_secs); + key.last_models_fetch_error = None; + if !upstream_metadata_updates.is_empty() { + let metadata = key + .upstream_metadata + .get_or_insert_with(|| json!({})) + .as_object_mut() + .expect("upstream metadata object"); + for update in upstream_metadata_updates { + metadata.insert(update.namespace.clone(), update.value.clone()); + } + } + key.updated_at_unix_secs = updated_at_unix_secs; + drop(keys); + self.upstream_metadata_updates + .lock() + .expect("metadata updates mutex") + .extend(upstream_metadata_updates.iter().map(|update| { + ( + key_id.to_string(), + update.namespace.clone(), + update.value.clone(), + updated_at_unix_secs, + ) + })); Ok(()) } @@ -815,12 +890,19 @@ mod tests { #[tokio::test] async fn model_fetch_uses_preset_models_without_endpoint() { let provider = sample_provider("provider-codex", "codex"); - let key = sample_key( + let mut key = sample_key( "key-codex", "provider-codex", "api_key", &["openai:responses"], ); + key.upstream_metadata = Some(json!({ + "codex": { + "quota_by_model": { + "gpt-5.6-sol": {"remaining_fraction": 0.75} + } + } + })); let state = TestState::new(vec![provider], vec![], vec![key], HashMap::new(), vec![]); let summary = perform_model_fetch_once_with_state(&state) @@ -832,9 +914,47 @@ mod tests { let updated = state.key("key-codex"); let allowed_models = updated .allowed_models + .as_ref() .and_then(|value| value.as_array().cloned()) .expect("allowed_models should be set"); assert!(allowed_models.iter().any(|model| model == "gpt-5.4")); + let upstream_metadata = updated + .upstream_metadata + .as_ref() + .expect("Codex model catalog should be persisted"); + assert_eq!( + upstream_metadata["codex"]["quota_by_model"]["gpt-5.6-sol"]["remaining_fraction"], + 0.75 + ); + assert_eq!( + upstream_metadata["codex_models"]["cards"]["gpt-5.6-sol"]["multi_agent_version"], + "v2" + ); + let capabilities = crate::ai_serving::resolve_codex_responses_model_capabilities( + "gpt-5.6-sol", + "gpt-5.6-sol", + Some(upstream_metadata), + ); + assert!(capabilities.use_responses_lite); + assert_eq!( + capabilities.default_reasoning_effort.as_deref(), + Some("low") + ); + assert!(capabilities + .supported_reasoning_efforts + .iter() + .any(|effort| effort == "ultra")); + let metadata_updates = state + .upstream_metadata_updates + .lock() + .expect("metadata updates mutex"); + assert_eq!(metadata_updates.len(), 1); + assert_eq!(metadata_updates[0].0, "key-codex"); + assert_eq!(metadata_updates[0].1, "codex_models"); + assert_eq!( + metadata_updates[0].2["cards"]["gpt-5.6-sol"]["multi_agent_version"], + "v2" + ); assert!(state .cached_models .lock() diff --git a/apps/aether-gateway/src/model_fetch/runtime/state.rs b/apps/aether-gateway/src/model_fetch/runtime/state.rs index f09f232c7..81d23a5e0 100644 --- a/apps/aether-gateway/src/model_fetch/runtime/state.rs +++ b/apps/aether-gateway/src/model_fetch/runtime/state.rs @@ -4,7 +4,8 @@ use aether_data_contracts::repository::global_models::{ StoredAdminProviderModel, UpsertAdminProviderModelRecord, }; use aether_data_contracts::repository::provider_catalog::{ - StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider, + ProviderCatalogUpstreamMetadataNamespaceUpdate, StoredProviderCatalogEndpoint, + StoredProviderCatalogKey, StoredProviderCatalogProvider, }; use aether_model_fetch::{ModelFetchAssociationStore, ModelFetchTransportRuntime}; use async_trait::async_trait; @@ -42,9 +43,22 @@ pub(crate) trait ModelFetchRuntimeState: plan: &ExecutionPlan, ) -> Result; - async fn update_provider_catalog_key( + async fn update_provider_catalog_key_model_fetch_state( &self, - key: &StoredProviderCatalogKey, + key_id: &str, + allowed_models: Option<&Value>, + last_models_fetch_at_unix_secs: Option, + last_models_fetch_error: Option<&str>, + updated_at_unix_secs: Option, + ) -> Result<(), GatewayError>; + + async fn update_provider_catalog_key_model_fetch_success( + &self, + key_id: &str, + allowed_models: Option<&Value>, + last_models_fetch_at_unix_secs: u64, + upstream_metadata_updates: &[ProviderCatalogUpstreamMetadataNamespaceUpdate], + updated_at_unix_secs: Option, ) -> Result<(), GatewayError>; async fn write_upstream_models_cache( diff --git a/apps/aether-gateway/src/model_fetch/tests.rs b/apps/aether-gateway/src/model_fetch/tests.rs index 23a37821e..6c61411f8 100644 --- a/apps/aether-gateway/src/model_fetch/tests.rs +++ b/apps/aether-gateway/src/model_fetch/tests.rs @@ -310,6 +310,89 @@ async fn gateway_model_fetch_updates_key_and_syncs_provider_model_whitelist_asso execution_runtime_handle.abort(); } +#[tokio::test] +async fn codex_preset_model_fetch_associates_the_api_supported_review_model() { + let provider = StoredProviderCatalogProvider::new( + "provider-codex".to_string(), + "codex".to_string(), + Some("https://chatgpt.com".to_string()), + "codex".to_string(), + ) + .expect("provider should build") + .with_transport_fields(true, false, true, None, None, None, None, None, None); + let mut key = sample_key("provider-codex", "key-codex"); + key.locked_models = None; + key.model_include_patterns = None; + key.model_exclude_patterns = None; + let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( + vec![provider], + Vec::new(), + vec![key], + )); + let global_model_repository = Arc::new( + InMemoryGlobalModelReadRepository::seed(Vec::new()).with_admin_global_models(vec![ + sample_global_model( + "global-model-codex-auto-review", + "codex-auto-review", + &["codex-auto-review"], + ), + ]), + ); + let data_state = crate::data::GatewayDataState::disabled() + .attach_provider_catalog_repository_for_tests(Arc::clone(&provider_catalog_repository)) + .with_global_model_repository_for_tests(Arc::clone(&global_model_repository)) + .with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY); + let state = AppState::new() + .expect("gateway should build") + .with_data_state_for_tests(data_state); + + let summary = perform_model_fetch_once(&state) + .await + .expect("Codex preset model fetch should succeed"); + assert_eq!( + summary, + ModelFetchRunSummary { + attempted: 1, + succeeded: 1, + failed: 0, + skipped: 0, + } + ); + + let updated_key = provider_catalog_repository + .list_keys_by_ids(&["key-codex".to_string()]) + .await + .expect("keys should load") + .into_iter() + .next() + .expect("updated key should exist"); + assert!(updated_key + .allowed_models + .as_ref() + .and_then(serde_json::Value::as_array) + .is_some_and(|models| models.iter().any(|model| model == "codex-auto-review"))); + assert!(updated_key + .upstream_metadata + .as_ref() + .is_some_and(|metadata| { + metadata["codex_models"]["cards"]["codex-auto-review"]["supported_in_api"] == true + })); + + let provider_models = global_model_repository + .list_admin_provider_models(&AdminProviderModelListQuery { + provider_id: "provider-codex".to_string(), + is_active: None, + offset: 0, + limit: 10_000, + }) + .await + .expect("provider models should load"); + assert!(provider_models.iter().any(|model| { + model.provider_model_name == "codex-auto-review" + && model.global_model_id == "global-model-codex-auto-review" + })); +} + #[tokio::test] async fn gateway_model_fetch_updates_key_and_syncs_provider_model_whitelist_associations_without_execution_runtime_override( ) { diff --git a/apps/aether-gateway/src/orchestration/classifier.rs b/apps/aether-gateway/src/orchestration/classifier.rs index ffce19805..cafffafc6 100644 --- a/apps/aether-gateway/src/orchestration/classifier.rs +++ b/apps/aether-gateway/src/orchestration/classifier.rs @@ -92,7 +92,10 @@ pub(crate) fn classify_local_failover( return LocalFailoverClassification::RetryStatusCode; } - if should_failover_local_upstream_status(input.status_code) { + if should_failover_local_upstream_status( + input.status_code, + policy.retry_client_errors_by_default, + ) { return LocalFailoverClassification::RetryUpstreamFailure; } @@ -109,8 +112,11 @@ pub(crate) fn local_failover_error_message(response_text: Option<&str>) -> Optio .filter(|value| !value.is_empty()) } -fn should_failover_local_upstream_status(status_code: u16) -> bool { - status_code >= 400 +fn should_failover_local_upstream_status( + status_code: u16, + retry_client_errors_by_default: bool, +) -> bool { + status_code >= 500 || status_code >= 400 && retry_client_errors_by_default } fn local_error_response_has_cyber_policy_code(response_text: Option<&str>) -> bool { @@ -474,6 +480,39 @@ mod tests { } } + #[test] + fn classifier_passes_through_client_errors_when_protocol_default_disables_failover() { + let policy = LocalFailoverPolicy { + retry_client_errors_by_default: false, + ..LocalFailoverPolicy::default() + }; + + for status_code in [400, 401, 429, 499] { + assert_eq!( + classify_local_failover(&policy, LocalFailoverInput::new(status_code, None)), + LocalFailoverClassification::UseDefault + ); + } + assert_eq!( + classify_local_failover(&policy, LocalFailoverInput::new(500, None)), + LocalFailoverClassification::RetryUpstreamFailure + ); + } + + #[test] + fn classifier_explicit_continue_rule_overrides_protocol_client_error_default() { + let policy = LocalFailoverPolicy { + continue_status_codes: [429].into_iter().collect(), + retry_client_errors_by_default: false, + ..LocalFailoverPolicy::default() + }; + + assert_eq!( + classify_local_failover(&policy, LocalFailoverInput::new(429, None)), + LocalFailoverClassification::RetryStatusCode + ); + } + #[test] fn classifier_keeps_embedded_rate_limit_error_in_success_response_on_default_path() { assert_eq!( diff --git a/apps/aether-gateway/src/orchestration/effects.rs b/apps/aether-gateway/src/orchestration/effects.rs index b1c81b7d2..241649344 100644 --- a/apps/aether-gateway/src/orchestration/effects.rs +++ b/apps/aether-gateway/src/orchestration/effects.rs @@ -243,7 +243,11 @@ fn local_client_session_affinity(report_context: Option<&Value>) -> Option) -> http::HeaderMap { diff --git a/apps/aether-gateway/src/orchestration/policy.rs b/apps/aether-gateway/src/orchestration/policy.rs index 6ae4028dd..23ca33252 100644 --- a/apps/aether-gateway/src/orchestration/policy.rs +++ b/apps/aether-gateway/src/orchestration/policy.rs @@ -7,7 +7,7 @@ use tracing::debug; use crate::provider_transport::GatewayProviderTransportSnapshot; use crate::AppState; -#[derive(Debug, Clone, Default, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct LocalFailoverPolicy { pub(crate) max_retries: Option, pub(crate) stop_status_codes: BTreeSet, @@ -15,6 +15,21 @@ pub(crate) struct LocalFailoverPolicy { pub(crate) success_failover_patterns: Vec, pub(crate) error_stop_patterns: Vec, pub(crate) stop_cyber_policy_errors: bool, + pub(crate) retry_client_errors_by_default: bool, +} + +impl Default for LocalFailoverPolicy { + fn default() -> Self { + Self { + max_retries: None, + stop_status_codes: BTreeSet::new(), + continue_status_codes: BTreeSet::new(), + success_failover_patterns: Vec::new(), + error_stop_patterns: Vec::new(), + stop_cyber_policy_errors: false, + retry_client_errors_by_default: true, + } + } } #[derive(Debug, Clone, PartialEq, Eq)] @@ -81,6 +96,10 @@ pub(crate) fn local_failover_policy_from_transport( LocalFailoverPolicy { max_retries, + retry_client_errors_by_default: + crate::ai_serving::api_format_defaults_to_client_error_failover( + &transport.endpoint.api_format, + ), stop_cyber_policy_errors: codex_cyber_flag_passthrough_enabled( &transport.provider.provider_type, transport.provider.config.as_ref(), @@ -144,6 +163,10 @@ pub(crate) fn local_failover_policy_from_report_context( .get("stop_cyber_policy_errors") .and_then(Value::as_bool) .unwrap_or(false), + retry_client_errors_by_default: object + .get("retry_client_errors_by_default") + .and_then(Value::as_bool) + .unwrap_or(true), }) } @@ -178,6 +201,7 @@ fn local_failover_policy_to_value(policy: &LocalFailoverPolicy) -> Value { "success_failover_patterns": policy.success_failover_patterns.iter().map(local_failover_regex_rule_to_value).collect::>(), "error_stop_patterns": policy.error_stop_patterns.iter().map(local_failover_regex_rule_to_value).collect::>(), "stop_cyber_policy_errors": policy.stop_cyber_policy_errors, + "retry_client_errors_by_default": policy.retry_client_errors_by_default, }) } @@ -375,10 +399,26 @@ mod tests { status_codes: [422].into_iter().collect(), }], stop_cyber_policy_errors: false, + retry_client_errors_by_default: true, }) ); } + #[test] + fn search_transport_disables_default_client_error_failover() { + let mut transport = sample_transport(None, None, None); + transport.endpoint.api_format = "openai:search".to_string(); + let policy = local_failover_policy_from_transport(&transport); + + assert!(!policy.retry_client_errors_by_default); + let report_context = append_local_failover_policy_to_value(json!({}), &transport); + assert_eq!( + local_failover_policy_from_report_context(Some(&report_context)) + .map(|policy| policy.retry_client_errors_by_default), + Some(false) + ); + } + #[test] fn codex_cyber_policy_passthrough_defaults_on_and_can_be_disabled() { let mut transport = sample_transport(None, None, None); diff --git a/apps/aether-gateway/src/privacy/mod.rs b/apps/aether-gateway/src/privacy/mod.rs index 42c2caad4..f2c91e03a 100644 --- a/apps/aether-gateway/src/privacy/mod.rs +++ b/apps/aether-gateway/src/privacy/mod.rs @@ -1063,6 +1063,7 @@ impl MaskChatRequestOptions { pub(crate) enum ChatPiiRedactionRequestFormat { OpenAiChat, OpenAiResponses, + OpenAiSearch, ClaudeMessages, } @@ -1071,6 +1072,7 @@ impl ChatPiiRedactionRequestFormat { match api_format.trim().to_ascii_lowercase().as_str() { "openai:chat" => Some(Self::OpenAiChat), "openai:responses" | "openai:responses:compact" => Some(Self::OpenAiResponses), + "openai:search" => Some(Self::OpenAiSearch), "claude:messages" => Some(Self::ClaudeMessages), _ => None, } @@ -1604,6 +1606,7 @@ fn request_collision_corpus(format: ChatPiiRedactionRequestFormat, value: &Value .map(|messages| chat_message_collision_corpus(messages)) .unwrap_or_default(), ChatPiiRedactionRequestFormat::OpenAiResponses => openai_responses_collision_corpus(value), + ChatPiiRedactionRequestFormat::OpenAiSearch => openai_search_collision_corpus(value), ChatPiiRedactionRequestFormat::ClaudeMessages => claude_messages_collision_corpus(value), } } @@ -1622,6 +1625,9 @@ fn mask_request_value( ChatPiiRedactionRequestFormat::OpenAiResponses => { mask_openai_responses_request_value(value, session, scan_state) } + ChatPiiRedactionRequestFormat::OpenAiSearch => { + mask_openai_search_request_value(value, session, scan_state) + } ChatPiiRedactionRequestFormat::ClaudeMessages => { mask_claude_messages_request_value(value, session, scan_state) } @@ -1643,6 +1649,9 @@ async fn mask_request_value_async( ChatPiiRedactionRequestFormat::OpenAiResponses => { mask_openai_responses_request_value_async(value, session, scan_state, cache).await } + ChatPiiRedactionRequestFormat::OpenAiSearch => { + mask_openai_search_request_value_async(value, session, scan_state, cache).await + } ChatPiiRedactionRequestFormat::ClaudeMessages => { mask_claude_messages_request_value_async(value, session, scan_state, cache).await } @@ -1810,6 +1819,31 @@ fn openai_responses_collision_corpus(value: &Value) -> Vec { corpus } +const OPENAI_SEARCH_COMMAND_TEXT_FIELDS: [(&str, &str); 4] = [ + ("search_query", "q"), + ("image_query", "q"), + ("find", "pattern"), + ("weather", "location"), +]; + +fn openai_search_collision_corpus(value: &Value) -> Vec { + let mut corpus = openai_responses_collision_corpus(value); + let Some(commands) = value.get("commands").and_then(Value::as_object) else { + return corpus; + }; + for (command, field) in OPENAI_SEARCH_COMMAND_TEXT_FIELDS { + let Some(entries) = commands.get(command).and_then(Value::as_array) else { + continue; + }; + for entry in entries { + if let Some(text) = entry.get(field).and_then(Value::as_str) { + corpus.push(text.to_string()); + } + } + } + corpus +} + fn collect_openai_responses_input_collision_text(value: &Value, corpus: &mut Vec) { match value { Value::String(text) => corpus.push(text.clone()), @@ -2025,6 +2059,31 @@ fn mask_openai_responses_request_value( Ok(redacted) } +fn mask_openai_search_request_value( + value: &mut Value, + session: &mut RedactionSession, + scan_state: &mut RedactionScanState, +) -> Result { + let mut redacted = false; + if let Some(input) = value.get_mut("input") { + redacted |= mask_openai_responses_input_value(input, session, scan_state)?; + } + let Some(commands) = value.get_mut("commands").and_then(Value::as_object_mut) else { + return Ok(redacted); + }; + for (command, field) in OPENAI_SEARCH_COMMAND_TEXT_FIELDS { + let Some(entries) = commands.get_mut(command).and_then(Value::as_array_mut) else { + continue; + }; + for entry in entries { + if let Some(Value::String(text)) = entry.get_mut(field) { + redacted |= mask_json_string(text, session, scan_state)?; + } + } + } + Ok(redacted) +} + fn mask_openai_responses_input_value( value: &mut Value, session: &mut RedactionSession, @@ -2310,6 +2369,33 @@ async fn mask_openai_responses_request_value_async( Ok(redacted) } +async fn mask_openai_search_request_value_async( + value: &mut Value, + session: &mut RedactionSession, + scan_state: &mut RedactionScanState, + cache: Option<&RedisRedactionMappingCache<'_>>, +) -> Result { + let mut redacted = false; + if let Some(input) = value.get_mut("input") { + redacted |= + mask_openai_responses_input_value_async(input, session, scan_state, cache).await?; + } + let Some(commands) = value.get_mut("commands").and_then(Value::as_object_mut) else { + return Ok(redacted); + }; + for (command, field) in OPENAI_SEARCH_COMMAND_TEXT_FIELDS { + let Some(entries) = commands.get_mut(command).and_then(Value::as_array_mut) else { + continue; + }; + for entry in entries { + if let Some(Value::String(text)) = entry.get_mut(field) { + redacted |= mask_json_string_async(text, session, scan_state, cache).await?; + } + } + } + Ok(redacted) +} + async fn mask_openai_responses_input_value_async( value: &mut Value, session: &mut RedactionSession, @@ -4072,6 +4158,7 @@ mod tests { build_redaction_session_config, detect_candidates_with_probe, mask_chat_request_json, mask_chat_request_json_with_options, parse_chat_pii_redaction_rules, restore_sync_response_body, try_mask_chat_pii_request_json_with_options, + try_mask_chat_pii_request_value_with_cache_options, try_mask_chat_request_json_with_cache_options, try_mask_chat_request_json_with_options, ChatPiiRedactionRequestFormat, ChatPiiRedactionRuntimeConfig, DetectorProbe, MappingKey, MaskChatRequestOptions, RedactionKind, RedactionLimitError, RedactionMapping, @@ -4083,7 +4170,7 @@ mod tests { use aether_runtime_state::{RedisClientConfig, RuntimeState}; use aether_testkit::ManagedRedisServer; - use serde_json::json; + use serde_json::{json, Value}; fn assert_debug_surface_hides_values(debug: &str, originals: &[&str], sentinels: &[String]) { for original in originals { @@ -4734,6 +4821,98 @@ mod tests { .contains("secretValueABCDEF1234567890abcdef")); } + #[test] + fn pii_redaction_request_masks_openai_search_text_fields() { + let request = json!({ + "id": "session-1", + "model": "gpt-5.6", + "input": "Find alice@example.com", + "commands": { + "search_query": [{"q": "Phone +14155552671"}], + "image_query": [{"q": "Image for bob@example.com"}], + "find": [{"ref_id": "https://example.com/alice@example.com", "pattern": "secret_key=secretValueABCDEF1234567890abcdef"}], + "weather": [{"location": "Contact carol@example.com"}], + "open": [{"ref_id": "https://example.com/alice@example.com"}] + } + }); + let raw = serde_json::to_vec(&request).expect("request should serialize"); + + let masked = try_mask_chat_pii_request_json_with_options( + &raw, + ChatPiiRedactionRequestFormat::OpenAiSearch, + test_config(), + MaskChatRequestOptions::runtime(), + ) + .expect("search request should mask"); + let masked_json: Value = + serde_json::from_slice(&masked.body).expect("masked request should parse"); + + assert!(masked.redacted); + assert!(!masked_json["input"] + .as_str() + .unwrap() + .contains("alice@example.com")); + assert!(!masked_json["commands"]["search_query"][0]["q"] + .as_str() + .unwrap() + .contains("+14155552671")); + assert!(!masked_json["commands"]["image_query"][0]["q"] + .as_str() + .unwrap() + .contains("bob@example.com")); + assert!(!masked_json["commands"]["find"][0]["pattern"] + .as_str() + .unwrap() + .contains("secretValueABCDEF1234567890abcdef")); + assert!(!masked_json["commands"]["weather"][0]["location"] + .as_str() + .unwrap() + .contains("carol@example.com")); + assert_eq!( + masked_json["commands"]["open"][0]["ref_id"], + "https://example.com/alice@example.com" + ); + } + + #[tokio::test] + async fn pii_redaction_async_request_masks_openai_search_text_fields() { + let request = json!({ + "id": "session-1", + "model": "gpt-5.6", + "input": "Find alice@example.com", + "commands": { + "search_query": [{"q": "Phone +14155552671"}], + "find": [{"ref_id": "turn0search0", "pattern": "bob@example.com"}] + } + }); + + let masked = try_mask_chat_pii_request_value_with_cache_options( + &request, + ChatPiiRedactionRequestFormat::OpenAiSearch, + test_config(), + MaskChatRequestOptions::runtime(), + None, + ) + .await + .expect("search request should mask"); + let masked_json = masked.body_json.expect("masked body should be present"); + + assert!(masked.redacted); + assert!(!masked_json["input"] + .as_str() + .unwrap() + .contains("alice@example.com")); + assert!(!masked_json["commands"]["search_query"][0]["q"] + .as_str() + .unwrap() + .contains("+14155552671")); + assert!(!masked_json["commands"]["find"][0]["pattern"] + .as_str() + .unwrap() + .contains("bob@example.com")); + assert_eq!(masked_json["commands"]["find"][0]["ref_id"], "turn0search0"); + } + #[test] fn pii_redaction_request_avoids_cross_message_and_tool_argument_sentinel_collisions() { let mut probe = session_at(600); diff --git a/apps/aether-gateway/src/rate_limit.rs b/apps/aether-gateway/src/rate_limit.rs index 7bcb237ac..39f7541d7 100644 --- a/apps/aether-gateway/src/rate_limit.rs +++ b/apps/aether-gateway/src/rate_limit.rs @@ -508,6 +508,7 @@ mod tests { auth_context: Some(auth_context), admin_principal: None, local_auth_rejection: None, + model_directive_policy: Default::default(), } } diff --git a/apps/aether-gateway/src/request_candidate_runtime.rs b/apps/aether-gateway/src/request_candidate_runtime.rs index 6077db0ea..af7f2ad09 100644 --- a/apps/aether-gateway/src/request_candidate_runtime.rs +++ b/apps/aether-gateway/src/request_candidate_runtime.rs @@ -133,7 +133,7 @@ pub(crate) async fn resolve_request_candidate_required_capabilities( api_key_id: &str, requested_model: Option<&str>, explicit_required_capabilities: Option<&Value>, - enable_model_directives: bool, + model_directive_base_model: Option<&str>, ) -> Option { let mut merged = serde_json::Map::new(); @@ -146,7 +146,7 @@ pub(crate) async fn resolve_request_candidate_required_capabilities( select_requested_model_capabilities( settings.as_ref(), requested_model, - enable_model_directives, + model_directive_base_model, ), ), Err(error) => { @@ -199,7 +199,7 @@ fn merge_capability_object(target: &mut serde_json::Map, source: fn select_requested_model_capabilities<'a>( settings: Option<&'a Value>, requested_model: Option<&str>, - enable_model_directives: bool, + model_directive_base_model: Option<&str>, ) -> Option<&'a Value> { let requested_model = requested_model .map(str::trim) @@ -207,10 +207,9 @@ fn select_requested_model_capabilities<'a>( let settings = settings?.as_object()?; find_model_capabilities(settings, requested_model).or_else(|| { - enable_model_directives - .then(|| crate::ai_serving::model_directive_base_model(requested_model)) - .flatten() - .as_deref() + model_directive_base_model + .map(str::trim) + .filter(|base_model| !base_model.is_empty() && *base_model != requested_model) .and_then(|base_model| find_model_capabilities(settings, base_model)) }) } @@ -924,7 +923,7 @@ mod tests { use super::{ ensure_execution_request_candidate_slot, persist_available_local_candidate, record_report_request_candidate_status, resolve_request_candidate_required_capabilities, - SchedulerRequestCandidateStatusUpdate, + select_requested_model_capabilities, SchedulerRequestCandidateStatusUpdate, }; use crate::data::GatewayDataState; use crate::AppState; @@ -997,6 +996,7 @@ mod tests { global_model_id: "global-model-1".to_string(), global_model_name: "gpt-5".to_string(), selected_provider_model_name: "gpt-5".to_string(), + supports_streaming: true, mapping_matched_model: None, } } @@ -1230,7 +1230,7 @@ mod tests { "api-key-1", Some("gpt-5"), Some(&explicit_required_capabilities), - false, + None, ) .await .expect("required capabilities should resolve"); @@ -1240,6 +1240,48 @@ mod tests { assert_eq!(required_capabilities["gemini_files"], json!(true)); } + #[test] + fn requested_model_capabilities_use_the_policy_resolved_base_model() { + let base_only = json!({ + "deployment-alias": { + "context_1m": true + } + }); + assert_eq!( + select_requested_model_capabilities( + Some(&base_only), + Some("deployment-alias-VendorFuture"), + Some("deployment-alias"), + ), + Some(&base_only["deployment-alias"]) + ); + assert_eq!( + select_requested_model_capabilities( + Some(&base_only), + Some("deployment-alias-VendorFuture"), + None, + ), + None + ); + + let exact_and_base = json!({ + "deployment-alias-VendorFuture": { + "cache_1h": true + }, + "deployment-alias": { + "context_1m": true + } + }); + assert_eq!( + select_requested_model_capabilities( + Some(&exact_and_base), + Some("deployment-alias-VendorFuture"), + Some("deployment-alias"), + ), + Some(&exact_and_base["deployment-alias-VendorFuture"]) + ); + } + #[tokio::test] async fn persists_request_required_capabilities_instead_of_provider_key_capabilities() { let repository = Arc::new(InMemoryRequestCandidateRepository::default()); diff --git a/apps/aether-gateway/src/scheduler/candidate/tests/model.rs b/apps/aether-gateway/src/scheduler/candidate/tests/model.rs index d96454012..94b79ed2c 100644 --- a/apps/aether-gateway/src/scheduler/candidate/tests/model.rs +++ b/apps/aether-gateway/src/scheduler/candidate/tests/model.rs @@ -138,6 +138,7 @@ fn scheduler_candidate_is_serializable() { global_model_id: "global-model-1".to_string(), global_model_name: "gpt-4.1".to_string(), selected_provider_model_name: "gpt-4.1-canary".to_string(), + supports_streaming: true, mapping_matched_model: Some("gpt-4.1-canary".to_string()), }; diff --git a/apps/aether-gateway/src/state/app.rs b/apps/aether-gateway/src/state/app.rs index 71e1796fd..2c67c12fc 100644 --- a/apps/aether-gateway/src/state/app.rs +++ b/apps/aether-gateway/src/state/app.rs @@ -377,6 +377,7 @@ pub struct AppState { pub(crate) upstream_target_admission: Arc, pub(crate) distributed_request_gate: Option>, pub(crate) client: reqwest::Client, + pub(crate) owner_forward_client: reqwest::Client, pub(crate) auth_context_cache: Arc, pub(crate) auth_snapshot_cache: Arc, pub(crate) admin_security_blacklist_cache: Arc>, diff --git a/apps/aether-gateway/src/state/catalog.rs b/apps/aether-gateway/src/state/catalog.rs index b05712cee..1e3d5a823 100644 --- a/apps/aether-gateway/src/state/catalog.rs +++ b/apps/aether-gateway/src/state/catalog.rs @@ -721,6 +721,79 @@ impl AppState { Ok(updated) } + pub(crate) async fn upsert_provider_catalog_key_upstream_metadata_namespace( + &self, + key_id: &str, + namespace: &str, + value: &serde_json::Value, + updated_at_unix_secs: Option, + ) -> Result { + let updated = self + .data + .upsert_provider_catalog_key_upstream_metadata_namespace( + key_id, + namespace, + value, + updated_at_unix_secs, + ) + .await + .map_err(|err| GatewayError::Internal(err.to_string()))?; + if updated { + self.invalidate_provider_routing_caches(); + } + Ok(updated) + } + + pub(crate) async fn update_provider_catalog_key_model_fetch_state( + &self, + key_id: &str, + allowed_models: Option<&serde_json::Value>, + last_models_fetch_at_unix_secs: Option, + last_models_fetch_error: Option<&str>, + updated_at_unix_secs: Option, + ) -> Result { + let updated = self + .data + .update_provider_catalog_key_model_fetch_state( + key_id, + allowed_models, + last_models_fetch_at_unix_secs, + last_models_fetch_error, + updated_at_unix_secs, + ) + .await + .map_err(|err| GatewayError::Internal(err.to_string()))?; + if updated { + self.invalidate_provider_routing_caches(); + } + Ok(updated) + } + + pub(crate) async fn update_provider_catalog_key_model_fetch_success( + &self, + key_id: &str, + allowed_models: Option<&serde_json::Value>, + last_models_fetch_at_unix_secs: u64, + upstream_metadata_updates: &[provider_catalog::ProviderCatalogUpstreamMetadataNamespaceUpdate], + updated_at_unix_secs: Option, + ) -> Result { + let updated = self + .data + .update_provider_catalog_key_model_fetch_success( + key_id, + allowed_models, + last_models_fetch_at_unix_secs, + upstream_metadata_updates, + updated_at_unix_secs, + ) + .await + .map_err(|err| GatewayError::Internal(err.to_string()))?; + if updated { + self.invalidate_provider_routing_caches(); + } + Ok(updated) + } + pub(crate) async fn delete_provider_catalog_key( &self, key_id: &str, @@ -1341,6 +1414,7 @@ mod tests { "fixed_order", true, None, + "", ); state.candidate_page_cache.insert( cache_key.clone(), diff --git a/apps/aether-gateway/src/state/core.rs b/apps/aether-gateway/src/state/core.rs index c2f9d4501..ec7aa5dad 100644 --- a/apps/aether-gateway/src/state/core.rs +++ b/apps/aether-gateway/src/state/core.rs @@ -51,6 +51,7 @@ use super::super::{provider_transport, usage}; use crate::maintenance::spawn_account_self_check_worker; use crate::maintenance::spawn_audit_cleanup_worker; use crate::maintenance::spawn_db_maintenance_worker; +use crate::maintenance::spawn_fixed_provider_reconciliation_task; use crate::maintenance::spawn_gemini_file_mapping_cleanup_worker; use crate::maintenance::spawn_oauth_token_refresh_worker; use crate::maintenance::spawn_pending_cleanup_worker; @@ -250,6 +251,11 @@ impl AppState { http2_adaptive_window: true, ..HttpClientConfig::default() })?; + let owner_forward_client = build_http_client(&HttpClientConfig { + connect_timeout_ms: Some(10_000), + http2_adaptive_window: true, + ..HttpClientConfig::default() + })?; let frontdoor_runtime_guards = Arc::new(FrontdoorRuntimeGuardConfig::from_env()); Ok(Self { #[cfg(test)] @@ -287,6 +293,7 @@ impl AppState { ), distributed_request_gate: None, client, + owner_forward_client, auth_context_cache: Arc::new(AuthContextCache::default()), auth_snapshot_cache: Arc::new(AuthSnapshotCache::default()), admin_security_blacklist_cache: Arc::new(ValueCache::default()), @@ -716,6 +723,7 @@ impl AppState { self.data.clear_minimal_candidate_selection_cache(); self.data.clear_routing_group_cache(); self.data.clear_provider_catalog_cache(); + self.auth_request_cost_upper_bound_cache.clear(); self.routing_group_selection_cache.clear(); self.candidate_row_page_cache.clear(); self.candidate_page_cache.clear(); @@ -1501,6 +1509,13 @@ impl AppState { record_boot(crate::task_runtime::TASK_KEY_USAGE_QUEUE_WORKER); } + if let Some(handle) = spawn_fixed_provider_reconciliation_task(self.clone()) { + supervisor.supervise_handle( + crate::task_runtime::TASK_KEY_FIXED_PROVIDER_RECONCILIATION, + handle, + ); + } + let mut supervise_worker = |task_key: &'static str, handle: Option>| { if let Some(handle) = handle { diff --git a/apps/aether-gateway/src/state/integrations.rs b/apps/aether-gateway/src/state/integrations.rs index 195aeb88b..89b68f231 100644 --- a/apps/aether-gateway/src/state/integrations.rs +++ b/apps/aether-gateway/src/state/integrations.rs @@ -9,7 +9,8 @@ use aether_data_contracts::repository::global_models::{ StoredAdminProviderModel, UpsertAdminProviderModelRecord, }; use aether_data_contracts::repository::provider_catalog::{ - StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider, + ProviderCatalogUpstreamMetadataNamespaceUpdate, StoredProviderCatalogEndpoint, + StoredProviderCatalogKey, StoredProviderCatalogProvider, }; use aether_data_contracts::repository::quota::StoredProviderQuotaSnapshot; use aether_model_fetch::{ @@ -319,11 +320,43 @@ impl ModelFetchRuntimeState for AppState { execution_runtime::execute_execution_runtime_sync_plan(self, None, plan).await } - async fn update_provider_catalog_key( + async fn update_provider_catalog_key_model_fetch_state( &self, - key: &StoredProviderCatalogKey, + key_id: &str, + allowed_models: Option<&Value>, + last_models_fetch_at_unix_secs: Option, + last_models_fetch_error: Option<&str>, + updated_at_unix_secs: Option, ) -> Result<(), GatewayError> { - AppState::update_provider_catalog_key(self, key).await?; + AppState::update_provider_catalog_key_model_fetch_state( + self, + key_id, + allowed_models, + last_models_fetch_at_unix_secs, + last_models_fetch_error, + updated_at_unix_secs, + ) + .await?; + Ok(()) + } + + async fn update_provider_catalog_key_model_fetch_success( + &self, + key_id: &str, + allowed_models: Option<&Value>, + last_models_fetch_at_unix_secs: u64, + upstream_metadata_updates: &[ProviderCatalogUpstreamMetadataNamespaceUpdate], + updated_at_unix_secs: Option, + ) -> Result<(), GatewayError> { + AppState::update_provider_catalog_key_model_fetch_success( + self, + key_id, + allowed_models, + last_models_fetch_at_unix_secs, + upstream_metadata_updates, + updated_at_unix_secs, + ) + .await?; Ok(()) } diff --git a/apps/aether-gateway/src/system_features.rs b/apps/aether-gateway/src/system_features.rs index 78e10ff79..726a42f44 100644 --- a/apps/aether-gateway/src/system_features.rs +++ b/apps/aether-gateway/src/system_features.rs @@ -1,5 +1,16 @@ +use std::collections::BTreeSet; +use std::sync::Arc; + +use sha2::Digest as _; use tracing::warn; +use crate::ai_serving::{ + default_model_directive_mapping_patch, default_model_directive_suffixes, + model_directive_builtin_suffix_supported_for_source_model, + model_directive_suffix_has_builtin_mapping, parse_model_directive, + parse_model_directive_with_suffixes, ReasoningEffort, ServiceTier, MODEL_DIRECTIVE_API_FORMATS, +}; + use crate::handlers::shared::system_config_bool; use crate::state::AppState; @@ -7,129 +18,256 @@ pub(crate) const ENABLE_MODEL_DIRECTIVES_CONFIG_KEY: &str = "enable_model_direct pub(crate) const MODEL_DIRECTIVES_CONFIG_KEY: &str = "model_directives"; const REASONING_EFFORT_DIRECTIVE_KEY: &str = "reasoning_effort"; -pub(crate) async fn model_directives_enabled(state: &AppState) -> bool { - match state - .read_system_config_json_value(ENABLE_MODEL_DIRECTIVES_CONFIG_KEY) - .await - { - Ok(value) => system_config_bool(value.as_ref(), false), - Err(error) => { +#[derive(Debug, Clone)] +pub(crate) struct ModelDirectivePolicySnapshot { + policy: Arc, +} + +#[derive(Debug)] +struct ModelDirectivePolicy { + directives_enabled: bool, + reasoning_settings: Option, + cache_key: String, +} + +impl Default for ModelDirectivePolicySnapshot { + fn default() -> Self { + Self::from_config_values(None, None) + } +} + +impl ModelDirectivePolicySnapshot { + pub(crate) async fn load(state: &AppState) -> Self { + let (enabled, settings) = tokio::join!( + state.read_system_config_json_value(ENABLE_MODEL_DIRECTIVES_CONFIG_KEY), + state.read_system_config_json_value(MODEL_DIRECTIVES_CONFIG_KEY), + ); + if let Err(error) = &enabled { warn!( error = ?error, "gateway model directives config lookup failed" ); - false + } + if let Err(error) = &settings { + warn!( + error = ?error, + "gateway model directives detail config lookup failed" + ); + } + model_directive_policy_snapshot_from_config_reads(enabled, settings) + } + + pub(crate) fn from_config_values( + enabled: Option<&serde_json::Value>, + settings: Option<&serde_json::Value>, + ) -> Self { + let directives_enabled = system_config_bool(enabled, false); + let reasoning_settings = parse_reasoning_model_directive_settings(settings); + let cache_key = + model_directive_policy_cache_key(directives_enabled, reasoning_settings.as_ref()); + Self { + policy: Arc::new(ModelDirectivePolicy { + directives_enabled, + reasoning_settings, + cache_key, + }), + } + } + + pub(crate) fn cache_key(&self) -> &str { + self.policy.cache_key.as_str() + } + + pub(crate) fn reasoning_enabled(&self) -> bool { + self.policy.directives_enabled + && self + .policy + .reasoning_settings + .as_ref() + .map(ReasoningModelDirectiveSettings::enabled) + .unwrap_or(true) + } + + pub(crate) fn resolve_reasoning( + &self, + api_format: &str, + requested_model: Option<&str>, + ) -> ReasoningModelDirectiveResolution { + if !self.reasoning_enabled() { + return ReasoningModelDirectiveResolution::default(); + } + + let api_format = crate::ai_serving::normalize_api_format_alias(api_format); + if api_format.is_empty() + || !self + .policy + .reasoning_settings + .as_ref() + .and_then(|settings| settings.api_format_enabled(&api_format)) + .unwrap_or(true) + { + return ReasoningModelDirectiveResolution::default(); + } + + let enabled_suffixes = self + .policy + .reasoning_settings + .as_ref() + .and_then(|settings| settings.api_format_suffixes(&api_format)); + let mappings = self + .policy + .reasoning_settings + .as_ref() + .and_then(|settings| settings.api_format_mappings(&api_format)); + let default_suffixes; + let suffixes = match enabled_suffixes { + Some(suffixes) => suffixes + .iter() + .filter(|suffix| { + model_directive_suffix_has_builtin_mapping(suffix) + || mappings.is_some_and(|mappings| mappings.contains_key(*suffix)) + }) + .map(String::as_str) + .collect::>(), + None => { + default_suffixes = default_model_directive_suffixes(&api_format); + default_suffixes.to_vec() + } + }; + let Some(directive) = requested_model + .and_then(|model| parse_model_directive_with_suffixes(model, suffixes.iter().copied())) + else { + return ReasoningModelDirectiveResolution::default(); + }; + if directive.suffixes.iter().any(|suffix| { + model_directive_suffix_has_builtin_mapping(suffix) + && !model_directive_builtin_suffix_supported_for_source_model( + suffix, + &directive.base_model, + ) + }) { + return ReasoningModelDirectiveResolution::default(); + } + + ReasoningModelDirectiveResolution { + enabled: true, + api_format, + base_model: directive.base_model, + suffixes: directive.suffixes.clone(), + custom_mapping: model_directive_mapping_for_suffixes(&directive.suffixes, mappings), } } } -pub(crate) async fn reasoning_model_directive_enabled(state: &AppState) -> bool { - model_directives_enabled(state).await - && read_reasoning_model_directive_settings(state) - .await - .map(|settings| settings.enabled()) - .unwrap_or(true) +fn model_directive_policy_snapshot_from_config_reads( + enabled: Result, E>, + settings: Result, E>, +) -> ModelDirectivePolicySnapshot { + match (enabled, settings) { + (Ok(enabled), Ok(settings)) => { + ModelDirectivePolicySnapshot::from_config_values(enabled.as_ref(), settings.as_ref()) + } + _ => ModelDirectivePolicySnapshot::default(), + } } -pub(crate) async fn reasoning_model_directive_enabled_for_api_format( - state: &AppState, - api_format: &str, -) -> bool { - if !model_directives_enabled(state).await { - return false; - } - let settings = read_reasoning_model_directive_settings(state).await; - let enabled = settings - .as_ref() - .map(|settings| settings.enabled()) - .unwrap_or(true); - if !enabled { - return false; - } - - let api_format = crate::ai_serving::normalize_api_format_alias(api_format); - if api_format.is_empty() { - return false; - } - - settings - .as_ref() - .and_then(|settings| settings.api_format_enabled(&api_format)) - .unwrap_or(true) +fn model_directive_policy_cache_key( + directives_enabled: bool, + settings: Option<&ReasoningModelDirectiveSettings>, +) -> String { + let mut hasher = sha2::Sha256::new(); + let policy_value = serde_json::to_value((directives_enabled, settings)) + .expect("parsed model directive policy should always serialize"); + let serialized = serde_json::to_vec(&canonicalize_policy_json(policy_value)) + .expect("canonical model directive policy should always serialize"); + hasher.update(serialized); + format!("{:x}", hasher.finalize()) } -pub(crate) async fn reasoning_model_directive_enabled_for_api_format_and_model( - state: &AppState, - api_format: &str, - requested_model: Option<&str>, -) -> bool { - if !model_directives_enabled(state).await { - return false; +fn canonicalize_policy_json(value: serde_json::Value) -> serde_json::Value { + match value { + serde_json::Value::Object(object) => { + let mut entries = object.into_iter().collect::>(); + entries.sort_by(|left, right| left.0.cmp(&right.0)); + serde_json::Value::Object( + entries + .into_iter() + .map(|(key, value)| (key, canonicalize_policy_json(value))) + .collect(), + ) + } + serde_json::Value::Array(items) => { + serde_json::Value::Array(items.into_iter().map(canonicalize_policy_json).collect()) + } + value => value, } - let settings = read_reasoning_model_directive_settings(state).await; - let enabled = settings - .as_ref() - .map(|settings| settings.enabled()) - .unwrap_or(true); - if !enabled { - return false; - } - - let api_format = crate::ai_serving::normalize_api_format_alias(api_format); - if api_format.is_empty() { - return false; - } - - let api_format_enabled = settings - .as_ref() - .and_then(|settings| settings.api_format_enabled(&api_format)) - .unwrap_or(true); - if !api_format_enabled { - return false; - } - - let Some(suffixes) = requested_model.and_then(model_directive_suffixes_from_model) else { - return false; - }; - - let mappings = settings - .as_ref() - .and_then(|settings| settings.api_format_mappings(&api_format)); - suffixes - .iter() - .all(|suffix| suffix_supported_for_api_format(&api_format, suffix, mappings.as_ref())) } -pub(crate) async fn reasoning_model_directive_mapping_for_api_format_and_model( - state: &AppState, - api_format: &str, - requested_model: Option<&str>, -) -> Option { - if !reasoning_model_directive_enabled_for_api_format_and_model( - state, - api_format, - requested_model, - ) - .await - { - return None; - } - let suffixes = requested_model.and_then(model_directive_suffixes_from_model)?; - let api_format = crate::ai_serving::normalize_api_format_alias(api_format); - let settings = read_reasoning_model_directive_settings(state).await; - let mappings = settings - .as_ref() - .and_then(|settings| settings.api_format_mappings(&api_format)); - model_directive_mapping_for_suffixes(&api_format, &suffixes, mappings.as_ref()) -} - -const DEFAULT_MODEL_DIRECTIVE_SUFFIXES: &[&str] = - &["low", "medium", "high", "xhigh", "max", "fast"]; - #[derive(Debug, Clone, Default)] +pub(crate) struct ReasoningModelDirectiveResolution { + enabled: bool, + api_format: String, + base_model: String, + suffixes: Vec, + custom_mapping: Option, +} + +impl ReasoningModelDirectiveResolution { + pub(crate) fn enabled(&self) -> bool { + self.enabled + } + + pub(crate) fn base_model(&self) -> Option<&str> { + self.enabled.then_some(self.base_model.as_str()) + } + + pub(crate) fn mapping_patch_for_mapped_model( + &self, + mapped_model: &str, + ) -> Result, &'static str> { + if !self.enabled { + return Ok(None); + } + + let mut patch = serde_json::json!({}); + let mut has_patch = false; + for suffix in &self.suffixes { + if model_directive_suffix_has_builtin_mapping(suffix) { + let Some(builtin_patch) = default_model_directive_mapping_patch( + &self.api_format, + mapped_model, + &self.base_model, + suffix, + ) else { + return Err("model_directive_target_unsupported"); + }; + deep_merge_json(&mut patch, &builtin_patch); + has_patch = true; + } + } + if let Some(custom_mapping) = &self.custom_mapping { + deep_merge_json(&mut patch, custom_mapping); + has_patch = true; + } + if !has_patch { + return Err("model_directive_mapping_missing"); + } + Ok(Some(patch)) + } +} + +#[derive(Debug, Clone, Default, serde::Serialize)] struct ReasoningModelDirectiveSettings { enabled: Option, - api_formats: Option>, + api_formats: Vec, +} + +#[derive(Debug, Clone, serde::Serialize)] +struct ReasoningApiFormatSettings { + api_format: String, + enabled: bool, + suffixes: Option>, + mappings: Option>, } impl ReasoningModelDirectiveSettings { @@ -138,198 +276,93 @@ impl ReasoningModelDirectiveSettings { } fn api_format_enabled(&self, api_format: &str) -> Option { - let api_formats = self.api_formats.as_ref()?; - api_formats.iter().find_map(|(key, value)| { - if crate::ai_serving::normalize_api_format_alias(key) != api_format { - return None; - } - Some(match value { - serde_json::Value::Object(object) => object - .get("enabled") - .map(|value| system_config_bool(Some(value), true)) - .unwrap_or(true), - _ => system_config_bool(Some(value), true), - }) - }) + self.api_format_settings(api_format) + .map(|settings| settings.enabled) } fn api_format_mappings( &self, api_format: &str, - ) -> Option> { - let api_formats = self.api_formats.as_ref()?; - api_formats.iter().find_map(|(key, value)| { - if crate::ai_serving::normalize_api_format_alias(key) != api_format { - return None; - } - let object = value.as_object()?; - if let Some(mappings) = object - .get("mappings") - .and_then(serde_json::Value::as_object) - { - return Some(normalize_reasoning_mappings(mappings)); - } - let mappings = object - .get("suffixes")? - .as_array()? - .iter() - .filter_map(|value| value.as_str()) - .filter_map(normalize_reasoning_suffix) - .filter_map(|suffix| { - default_reasoning_mapping(api_format, &suffix).map(|mapping| (suffix, mapping)) - }) - .collect::>(); - Some(mappings) - }) + ) -> Option<&serde_json::Map> { + self.api_format_settings(api_format)?.mappings.as_ref() + } + + fn api_format_suffixes(&self, api_format: &str) -> Option<&BTreeSet> { + self.api_format_settings(api_format)?.suffixes.as_ref() + } + + fn api_format_settings(&self, api_format: &str) -> Option<&ReasoningApiFormatSettings> { + self.api_formats + .iter() + .find(|settings| settings.api_format == api_format) } } fn model_directive_suffixes_from_model(model: &str) -> Option> { - let mut base_model = model.trim(); - let mut suffixes = Vec::new(); - let mut has_reasoning_effort = false; - let mut has_service_tier = false; - while let Some((candidate_base, suffix)) = base_model.rsplit_once('-') { - let Some(suffix) = normalize_reasoning_suffix(suffix) else { - break; - }; - match model_directive_suffix_kind(&suffix)? { - ModelDirectiveSuffixKind::ReasoningEffort => { - if has_reasoning_effort { - return None; - } - has_reasoning_effort = true; - } - ModelDirectiveSuffixKind::ServiceTier => { - if has_service_tier { - return None; - } - has_service_tier = true; - } - } - suffixes.push(suffix); - base_model = candidate_base.trim(); - } - if base_model.is_empty() || suffixes.is_empty() { + Some( + parse_model_directive(model)? + .overrides + .iter() + .map(|override_item| override_item.suffix().to_string()) + .collect(), + ) +} + +fn normalize_model_directive_suffix(suffix: &str) -> Option { + let suffix = suffix.trim(); + if suffix.is_empty() { return None; } - suffixes.sort_by_key(|suffix| { - DEFAULT_MODEL_DIRECTIVE_SUFFIXES - .iter() - .position(|value| value == suffix) - .unwrap_or(usize::MAX) - }); - Some(suffixes) + ReasoningEffort::parse(suffix) + .map(|effort| effort.as_str().to_string()) + .or_else(|| ServiceTier::parse(suffix).map(|tier| tier.as_directive_suffix().to_string())) + .or_else(|| Some(suffix.to_string())) } -fn normalize_reasoning_suffix(suffix: &str) -> Option { - let normalized = suffix.trim().to_ascii_lowercase(); - DEFAULT_MODEL_DIRECTIVE_SUFFIXES - .contains(&normalized.as_str()) - .then_some(normalized) -} - -enum ModelDirectiveSuffixKind { - ReasoningEffort, - ServiceTier, -} - -fn model_directive_suffix_kind(suffix: &str) -> Option { - match suffix { - "low" | "medium" | "high" | "xhigh" | "max" => { - Some(ModelDirectiveSuffixKind::ReasoningEffort) - } - "fast" => Some(ModelDirectiveSuffixKind::ServiceTier), - _ => None, - } +fn normalize_configured_suffixes<'a>(suffixes: impl Iterator) -> BTreeSet { + suffixes + .filter_map(normalize_model_directive_suffix) + .collect() } fn normalize_reasoning_mappings( mappings: &serde_json::Map, ) -> serde_json::Map { - mappings - .iter() - .filter_map(|(suffix, mapping)| { - normalize_reasoning_suffix(suffix).map(|suffix| (suffix, mapping.clone())) - }) - .collect() + let mut normalized = serde_json::Map::new(); + for (suffix, mapping) in mappings { + let Some(suffix) = normalize_model_directive_suffix(suffix) else { + continue; + }; + normalized.insert(suffix, mapping.clone()); + } + normalized } fn default_reasoning_mapping(api_format: &str, suffix: &str) -> Option { - match api_format { - "openai:chat" if suffix == "fast" => { - Some(serde_json::json!({ "service_tier": "priority" })) - } - "openai:chat" => { - let effort = openai_reasoning_effort_value(suffix)?; - Some(serde_json::json!({ "reasoning_effort": effort })) - } - "openai:responses" | "openai:responses:compact" if suffix == "fast" => { - Some(serde_json::json!({ "service_tier": "priority" })) - } - "openai:responses" | "openai:responses:compact" => { - let effort = openai_reasoning_effort_value(suffix)?; - Some(serde_json::json!({ "reasoning": { "effort": effort } })) - } - "claude:messages" => Some(serde_json::json!({ - "thinking": { - "type": "enabled", - "budget_tokens": match suffix { - "low" => 1024, - "medium" => 4096, - "high" => 8192, - "xhigh" => 16384, - "max" => 32768, - _ => return None, - } - } - })), - "gemini:generate_content" => Some(serde_json::json!({ - "generationConfig": { - "thinkingConfig": { - "thinkingBudget": match suffix { - "low" => 1024, - "medium" => 4096, - "high" => 8192, - "xhigh" => 16384, - "max" => -1, - _ => return None, - } - } - } - })), - _ => None, - } + default_reasoning_mapping_for_model(api_format, "", suffix) } -fn suffix_supported_for_api_format( +fn default_reasoning_mapping_for_model( api_format: &str, + model: &str, suffix: &str, - mappings: Option<&serde_json::Map>, -) -> bool { - mappings - .map(|mappings| mappings.contains_key(suffix)) - .unwrap_or_else(|| default_reasoning_mapping(api_format, suffix).is_some()) +) -> Option { + default_model_directive_mapping_patch(api_format, model, model, suffix) } fn model_directive_mapping_for_suffixes( - api_format: &str, suffixes: &[String], mappings: Option<&serde_json::Map>, ) -> Option { let mut combined = serde_json::json!({}); + let mut has_custom_mapping = false; for suffix in suffixes { - let mapping = mappings - .and_then(|mappings| mappings.get(suffix).cloned()) - .or_else(|| { - mappings - .is_none() - .then(|| default_reasoning_mapping(api_format, suffix)) - .flatten() - })?; - deep_merge_json(&mut combined, &mapping); + if let Some(mapping) = mappings.and_then(|mappings| mappings.get(suffix)) { + deep_merge_json(&mut combined, mapping); + has_custom_mapping = true; + } } - Some(combined) + has_custom_mapping.then_some(combined) } fn deep_merge_json(target: &mut serde_json::Value, patch: &serde_json::Value) { @@ -350,60 +383,165 @@ fn deep_merge_json(target: &mut serde_json::Value, patch: &serde_json::Value) { } } -fn openai_reasoning_effort_value(suffix: &str) -> Option<&'static str> { - match suffix { - "low" => Some("low"), - "medium" => Some("medium"), - "high" => Some("high"), - "xhigh" => Some("xhigh"), - "max" => Some("max"), - _ => None, - } -} - -async fn read_reasoning_model_directive_settings( - state: &AppState, -) -> Option { - match state - .read_system_config_json_value(MODEL_DIRECTIVES_CONFIG_KEY) - .await - { - Ok(value) => parse_reasoning_model_directive_settings(value.as_ref()), - Err(error) => { - warn!( - error = ?error, - "gateway model directives detail config lookup failed" - ); - None - } - } -} - fn parse_reasoning_model_directive_settings( value: Option<&serde_json::Value>, ) -> Option { let root = value?.as_object()?; let reasoning = root.get(REASONING_EFFORT_DIRECTIVE_KEY)?.as_object()?; + let mut api_formats = reasoning + .get("api_formats") + .and_then(serde_json::Value::as_object) + .map(|formats| { + formats + .iter() + .filter_map(|(api_format, value)| { + let api_format = crate::ai_serving::normalize_api_format_alias(api_format); + if !MODEL_DIRECTIVE_API_FORMATS.contains(&api_format.as_str()) { + return None; + } + let Some(object) = value.as_object() else { + return Some(ReasoningApiFormatSettings { + api_format, + enabled: system_config_bool(Some(value), true), + suffixes: None, + mappings: None, + }); + }; + let configured_mappings = object + .get("mappings") + .and_then(serde_json::Value::as_object); + let configured_suffixes = + object.get("suffixes").and_then(serde_json::Value::as_array); + let suffixes = configured_suffixes + .map(|suffixes| { + normalize_configured_suffixes( + suffixes.iter().filter_map(serde_json::Value::as_str), + ) + }) + .or_else(|| { + configured_mappings.map(|mappings| { + normalize_configured_suffixes(mappings.keys().map(String::as_str)) + }) + }); + let mappings = configured_mappings + .map(normalize_reasoning_mappings) + .or_else(|| configured_suffixes.map(|_| serde_json::Map::new())); + Some(ReasoningApiFormatSettings { + api_format, + enabled: object + .get("enabled") + .map(|value| system_config_bool(Some(value), true)) + .unwrap_or(true), + suffixes, + mappings, + }) + }) + .collect::>() + }) + .unwrap_or_default(); + api_formats.sort_by(|left, right| left.api_format.cmp(&right.api_format)); Some(ReasoningModelDirectiveSettings { enabled: reasoning .get("enabled") .map(|value| system_config_bool(Some(value), true)), - api_formats: reasoning - .get("api_formats") - .and_then(serde_json::Value::as_object) - .cloned(), + api_formats, }) } #[cfg(test)] mod tests { use super::{ - default_reasoning_mapping, model_directive_mapping_for_suffixes, + default_reasoning_mapping, default_reasoning_mapping_for_model, + model_directive_mapping_for_suffixes, model_directive_policy_snapshot_from_config_reads, model_directive_suffixes_from_model, parse_reasoning_model_directive_settings, - suffix_supported_for_api_format, + ModelDirectivePolicySnapshot, }; use serde_json::json; + #[test] + fn model_directive_policy_snapshot_is_immutable_after_source_values_change() { + let enabled = json!(true); + let mut settings = json!({ + "reasoning_effort": { + "enabled": true, + "api_formats": { + "openai:responses": { + "suffixes": ["low"] + } + } + } + }); + let snapshot = + ModelDirectivePolicySnapshot::from_config_values(Some(&enabled), Some(&settings)); + let cloned_snapshot = snapshot.clone(); + assert!(std::sync::Arc::ptr_eq( + &snapshot.policy, + &cloned_snapshot.policy, + )); + let equivalent_snapshot = ModelDirectivePolicySnapshot::from_config_values( + Some(&json!("true")), + Some(&json!({ + "reasoning_effort": { + "enabled": "true", + "api_formats": { + "OPENAI:RESPONSES": { + "suffixes": ["LOW"] + } + } + } + })), + ); + assert_eq!(snapshot.cache_key(), equivalent_snapshot.cache_key()); + + settings["reasoning_effort"]["api_formats"]["openai:responses"]["suffixes"] = + json!(["high"]); + let updated_snapshot = + ModelDirectivePolicySnapshot::from_config_values(Some(&enabled), Some(&settings)); + assert_ne!(snapshot.cache_key(), updated_snapshot.cache_key()); + + assert!(snapshot + .resolve_reasoning("openai:responses", Some("gpt-5.6-sol-low")) + .enabled()); + assert!(!snapshot + .resolve_reasoning("openai:responses", Some("gpt-5.6-sol-high")) + .enabled()); + } + + #[test] + fn model_directive_policy_config_read_errors_fail_closed_but_absence_keeps_defaults() { + let missing_details = model_directive_policy_snapshot_from_config_reads::<&str>( + Ok(Some(json!(true))), + Ok(None), + ); + assert!(missing_details.reasoning_enabled()); + assert!(missing_details + .resolve_reasoning("openai:responses", Some("gpt-5.6-sol-low")) + .enabled()); + + let detail_error = model_directive_policy_snapshot_from_config_reads( + Ok(Some(json!(true))), + Err("detail config unavailable"), + ); + assert!(!detail_error.reasoning_enabled()); + assert!(!detail_error + .resolve_reasoning("openai:responses", Some("gpt-5.6-sol-low")) + .enabled()); + assert_eq!( + detail_error.cache_key(), + ModelDirectivePolicySnapshot::default().cache_key() + ); + + let master_error = model_directive_policy_snapshot_from_config_reads( + Err("master config unavailable"), + Ok(Some(json!({ + "reasoning_effort": { + "enabled": true + } + }))), + ); + assert!(!master_error.reasoning_enabled()); + } + #[test] fn reasoning_model_directive_settings_parse_endpoint_flags() { let value = json!({ @@ -415,7 +553,7 @@ mod tests { "enabled": true, "mappings": { "high": { "thinking": { "type": "enabled", "budget_tokens": 8192 } }, - "max": { "thinking": { "type": "enabled", "budget_tokens": 32768 } } + "max": { "thinking": { "type": "enabled", "budget_tokens": 65536 } } } } } @@ -432,11 +570,71 @@ mod tests { settings .api_format_mappings("claude:messages") .and_then(|mappings| mappings.get("max").cloned()), - Some(json!({ "thinking": { "type": "enabled", "budget_tokens": 32768 } })) + Some(json!({ "thinking": { "type": "enabled", "budget_tokens": 65536 } })) ); assert_eq!(settings.api_format_enabled("gemini:generate_content"), None); } + #[test] + fn configured_mappings_remain_authoritative() { + let value = json!({ + "reasoning_effort": { + "api_formats": { + "openai:chat": { + "suffixes": ["low", "max"] + }, + "openai:responses": { + "mappings": { + "low": { "reasoning": { "effort": "low" } }, + "max": { "reasoning": { "effort": "xhigh" } } + } + }, + "claude:messages": { + "mappings": { + "low": { "thinking": { "type": "enabled", "budget_tokens": 1024 } }, + "high": { "thinking": { "type": "enabled", "budget_tokens": 7777 } } + } + } + } + } + }); + + let settings = + parse_reasoning_model_directive_settings(Some(&value)).expect("settings should parse"); + assert!(settings + .api_format_mappings("openai:chat") + .expect("chat mappings") + .is_empty()); + assert_eq!( + settings + .api_format_suffixes("openai:chat") + .expect("chat suffixes") + .into_iter() + .collect::>(), + vec!["low", "max"] + ); + let responses = settings + .api_format_mappings("openai:responses") + .expect("responses mappings"); + assert_eq!( + responses["low"], + json!({ "reasoning": { "effort": "low" } }) + ); + assert_eq!( + responses["max"], + json!({ "reasoning": { "effort": "xhigh" } }) + ); + + let claude = settings + .api_format_mappings("claude:messages") + .expect("Claude mappings"); + assert_eq!( + claude["low"], + json!({ "thinking": { "type": "enabled", "budget_tokens": 1024 } }) + ); + assert_eq!(claude["high"]["thinking"]["budget_tokens"], 7777); + } + #[test] fn default_fast_suffix_maps_to_openai_priority_service_tier() { assert_eq!( @@ -459,7 +657,17 @@ mod tests { } #[test] - fn default_max_suffix_is_preserved_for_openai_formats() { + fn openai_reasoning_defaults_cover_the_complete_effort_contract() { + for effort in ["none", "minimal", "low", "medium", "high", "xhigh"] { + assert_eq!( + default_reasoning_mapping("openai:chat", effort), + Some(json!({ "reasoning_effort": effort })) + ); + assert_eq!( + default_reasoning_mapping("openai:responses", effort), + Some(json!({ "reasoning": { "effort": effort } })) + ); + } assert_eq!( default_reasoning_mapping("openai:chat", "max"), Some(json!({ "reasoning_effort": "max" })) @@ -469,9 +677,27 @@ mod tests { Some(json!({ "reasoning": { "effort": "max" } })) ); assert_eq!( - default_reasoning_mapping("openai:responses:compact", "max"), + default_reasoning_mapping_for_model("openai:chat", "gpt-5.6-sol", "max"), + Some(json!({ "reasoning_effort": "max" })) + ); + assert_eq!( + default_reasoning_mapping_for_model("openai:responses", "gpt-5.6-sol", "max"), Some(json!({ "reasoning": { "effort": "max" } })) ); + assert_eq!( + default_reasoning_mapping_for_model("openai:responses", "gpt-5.4", "max"), + None + ); + } + + #[test] + fn parses_global_reasoning_suffix_vocabulary() { + for effort in ["none", "minimal", "low", "medium", "high", "xhigh", "max"] { + assert_eq!( + model_directive_suffixes_from_model(&format!("gpt-5.6-sol-{effort}")), + Some(vec![effort.to_string()]) + ); + } } #[test] @@ -487,29 +713,212 @@ mod tests { ); assert_eq!( model_directive_mapping_for_suffixes( - "openai:chat", expected.as_ref().expect("suffixes should parse"), None, ), + None + ); + + let custom = json!({ + "xhigh": { "reasoning_effort": "high" }, + "fast": { "service_tier": "default" } + }); + assert_eq!( + model_directive_mapping_for_suffixes( + expected.as_ref().expect("suffixes should parse"), + custom.as_object(), + ), Some(json!({ - "reasoning_effort": "xhigh", - "service_tier": "priority" + "reasoning_effort": "high", + "service_tier": "default" })) ); } #[test] - fn combined_suffix_requires_each_suffix_to_support_api_format() { - let suffixes = model_directive_suffixes_from_model("gpt-5.4-xhigh-fast").expect("suffixes"); - assert!(suffixes - .iter() - .all(|suffix| { suffix_supported_for_api_format("openai:responses", suffix, None) })); - assert!(!suffixes - .iter() - .all(|suffix| { suffix_supported_for_api_format("claude:messages", suffix, None) })); + fn mapped_model_controls_builtin_reasoning_capability() { + let snapshot = ModelDirectivePolicySnapshot::from_config_values(Some(&json!(true)), None); + let resolution = + snapshot.resolve_reasoning("openai:responses", Some("deployment-alias-max")); + + assert_eq!(resolution.base_model(), Some("deployment-alias")); assert_eq!( - model_directive_mapping_for_suffixes("claude:messages", &suffixes, None), - None + resolution.mapping_patch_for_mapped_model("gpt-5.6-sol"), + Ok(Some(json!({ "reasoning": { "effort": "max" } }))) + ); + assert_eq!( + resolution.mapping_patch_for_mapped_model("gpt-5.4"), + Err("model_directive_target_unsupported") + ); + + let known_source = snapshot.resolve_reasoning("openai:responses", Some("gpt-5.6-sol-max")); + assert_eq!( + known_source.mapping_patch_for_mapped_model("azure-production"), + Ok(Some(json!({ "reasoning": { "effort": "max" } }))) + ); + + let unsupported_source = + snapshot.resolve_reasoning("openai:responses", Some("gpt-5.4-max")); + assert_eq!( + unsupported_source.mapping_patch_for_mapped_model("azure-production"), + Err("model_directive_target_unsupported") ); } + + #[test] + fn codex_ultra_uses_the_builtin_policy_path_for_default_and_explicit_settings() { + let snapshots = [ + ModelDirectivePolicySnapshot::from_config_values(Some(&json!(true)), None), + ModelDirectivePolicySnapshot::from_config_values( + Some(&json!(true)), + Some(&json!({ + "reasoning_effort": { + "enabled": true, + "api_formats": { + "openai:responses": { + "enabled": true, + "suffixes": ["ultra"], + "mappings": {} + } + } + } + })), + ), + ]; + + for snapshot in snapshots { + for model in ["gpt-5.6-sol", "gpt-5.6-terra"] { + let requested_model = format!("{model}-ultra"); + let resolution = + snapshot.resolve_reasoning("openai:responses", Some(&requested_model)); + assert_eq!(resolution.base_model(), Some(model)); + assert_eq!( + resolution.mapping_patch_for_mapped_model(model), + Ok(Some(json!({ "reasoning": { "effort": "ultra" } }))) + ); + } + + for requested_model in [ + "gpt-5.6-ultra", + "gpt-5.6-luna-ultra", + "gpt-5.4-ultra", + "other-model-ultra", + ] { + assert!(!snapshot + .resolve_reasoning("openai:responses", Some(requested_model)) + .enabled()); + } + } + } + + #[test] + fn custom_suffix_mapping_defines_an_executable_policy_directive() { + let snapshot = ModelDirectivePolicySnapshot::from_config_values( + Some(&json!(true)), + Some(&json!({ + "reasoning_effort": { + "api_formats": { + "openai:responses": { + "suffixes": ["Future", "VendorFuture"], + "mappings": { + "VendorFuture": { + "reasoning": { "context": "all_turns" } + } + } + } + } + } + })), + ); + let resolution = + snapshot.resolve_reasoning("openai:responses", Some("deployment-alias-VendorFuture")); + + assert_eq!(resolution.base_model(), Some("deployment-alias")); + assert_eq!( + resolution.mapping_patch_for_mapped_model("gpt-5.6-sol"), + Ok(Some(json!({ + "reasoning": { "context": "all_turns" } + }))) + ); + assert!(!snapshot + .resolve_reasoning("openai:responses", Some("deployment-alias-Future"),) + .enabled()); + } + + #[test] + fn configured_mapping_keys_define_suffixes_and_overrides() { + let value = json!({ + "reasoning_effort": { + "api_formats": { + "openai:chat": { + "mappings": { + "low": { "reasoning_effort": "low" }, + "medium": { "reasoning_effort": "medium" }, + "high": { "reasoning_effort": "high" }, + "xhigh": { "reasoning_effort": "xhigh" }, + "MAX": { "reasoning_effort": "xhigh" }, + "fast": { "service_tier": "priority" }, + "VendorFuture": { "keep": true } + } + } + } + } + }); + let settings = + parse_reasoning_model_directive_settings(Some(&value)).expect("settings should parse"); + let suffixes = settings + .api_format_suffixes("openai:chat") + .expect("chat suffixes"); + assert_eq!( + suffixes.iter().map(String::as_str).collect::>(), + vec![ + "VendorFuture", + "fast", + "high", + "low", + "max", + "medium", + "xhigh" + ] + ); + let mappings = settings + .api_format_mappings("openai:chat") + .expect("chat mappings"); + assert_eq!(mappings.len(), 7); + assert_eq!(mappings["max"], json!({ "reasoning_effort": "xhigh" })); + assert_eq!(mappings["VendorFuture"], json!({ "keep": true })); + } + + #[test] + fn explicit_openai_suffix_allowlist_never_reenables_disabled_efforts() { + let cases = [ + vec!["minimal", "low", "medium", "high", "xhigh", "max", "fast"], + vec!["none", "low", "medium", "high", "xhigh", "max", "fast"], + vec!["low", "medium", "high", "xhigh", "max", "fast"], + ]; + + for expected in cases { + let expected_suffixes = expected + .iter() + .copied() + .map(str::to_string) + .collect::>(); + let value = json!({ + "reasoning_effort": { + "api_formats": { + "openai:responses": { + "suffixes": expected + } + } + } + }); + let settings = parse_reasoning_model_directive_settings(Some(&value)) + .expect("settings should parse"); + let suffixes = settings + .api_format_suffixes("openai:responses") + .expect("Responses suffixes"); + + assert_eq!(suffixes, &expected_suffixes); + } + } } diff --git a/apps/aether-gateway/src/task_runtime/mod.rs b/apps/aether-gateway/src/task_runtime/mod.rs index 78a2b76a6..61e8cdde8 100644 --- a/apps/aether-gateway/src/task_runtime/mod.rs +++ b/apps/aether-gateway/src/task_runtime/mod.rs @@ -32,6 +32,8 @@ pub(crate) const TASK_KEY_DB_MAINTENANCE: &str = "maintenance.database"; pub(crate) const TASK_KEY_PENDING_CLEANUP: &str = "maintenance.pending.cleanup"; pub(crate) const TASK_KEY_REQUEST_CANDIDATE_CLEANUP: &str = "maintenance.request.candidate.cleanup"; pub(crate) const TASK_KEY_GEMINI_FILES_CLEANUP: &str = "maintenance.gemini.files.cleanup"; +pub(crate) const TASK_KEY_FIXED_PROVIDER_RECONCILIATION: &str = + "maintenance.provider.fixed_template.reconcile"; pub(crate) const TASK_KEY_OAUTH_TOKEN_REFRESH: &str = "maintenance.oauth.token.refresh"; pub(crate) const TASK_KEY_PROXY_NODE_STALE_CLEANUP: &str = "maintenance.proxy.node.stale.cleanup"; pub(crate) const TASK_KEY_PROXY_NODE_METRICS_CLEANUP: &str = @@ -49,6 +51,7 @@ pub(crate) const TASK_KEY_PROVIDER_BALANCE_REFRESH: &str = "provider.ops.balance const PROVIDER_DELETE_LOCK_TTL_SECS: u64 = 60 * 60 * 6; const RETRY_ONCE: RetryPolicy = RetryPolicy { max_attempts: 1 }; +const RETRY_THREE: RetryPolicy = RetryPolicy { max_attempts: 3 }; const TASK_DEFINITIONS: &[TaskDefinition] = &[ TaskDefinition::new( @@ -187,6 +190,14 @@ const TASK_DEFINITIONS: &[TaskDefinition] = &[ true, RETRY_ONCE, ), + TaskDefinition::new( + TASK_KEY_FIXED_PROVIDER_RECONCILIATION, + TaskKind::FireAndForget, + "startup", + true, + false, + RETRY_THREE, + ), TaskDefinition::new( TASK_KEY_OAUTH_TOKEN_REFRESH, TaskKind::Scheduled, diff --git a/apps/aether-gateway/src/tests/ai_execute/stream/decision.rs b/apps/aether-gateway/src/tests/ai_execute/stream/decision.rs index b93d776fc..cdb8346cb 100644 --- a/apps/aether-gateway/src/tests/ai_execute/stream/decision.rs +++ b/apps/aether-gateway/src/tests/ai_execute/stream/decision.rs @@ -463,8 +463,8 @@ async fn gateway_executes_openai_chat_stream_via_local_openai_responses_cross_fo accept: String, authorization: String, x_client_request_id: String, - session_id: String, - conversation_id: String, + codex_session_id: String, + codex_thread_id: String, instructions: String, user_text: String, prompt_cache_key: String, @@ -742,15 +742,15 @@ async fn gateway_executes_openai_chat_stream_via_local_openai_responses_cross_fo .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), - session_id: payload + codex_session_id: payload .get("headers") - .and_then(|value| value.get("session_id")) + .and_then(|value| value.get("session-id")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), - conversation_id: payload + codex_thread_id: payload .get("headers") - .and_then(|value| value.get("conversation_id")) + .and_then(|value| value.get("thread-id")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), @@ -889,19 +889,16 @@ async fn gateway_executes_openai_chat_stream_via_local_openai_responses_cross_fo ); assert_eq!( seen_execution_runtime_request.x_client_request_id, - "trace-openai-chat-cli-local-123" + seen_execution_runtime_request.codex_thread_id ); assert_eq!( - seen_execution_runtime_request.prompt_cache_key, - "bc749eb7-a9e2-5793-8d14-abd659c700b0" + seen_execution_runtime_request.codex_session_id, + seen_execution_runtime_request.codex_thread_id ); - assert_eq!( - seen_execution_runtime_request.session_id, - "d1e9b802644e1f52" - ); - assert_eq!( - seen_execution_runtime_request.conversation_id, - "d1e9b802644e1f52" + assert!(seen_execution_runtime_request.prompt_cache_key.is_empty()); + assert_ne!( + seen_execution_runtime_request.codex_thread_id, + seen_execution_runtime_request.trace_id ); assert_eq!( seen_execution_runtime_request.instructions, diff --git a/apps/aether-gateway/src/tests/ai_execute/stream/image.rs b/apps/aether-gateway/src/tests/ai_execute/stream/image.rs index a57a2063f..3ab8b0422 100644 --- a/apps/aether-gateway/src/tests/ai_execute/stream/image.rs +++ b/apps/aether-gateway/src/tests/ai_execute/stream/image.rs @@ -2,7 +2,6 @@ use super::{ any, build_router_with_state, build_state_with_execution_runtime_override, json, start_server, to_bytes, Arc, Body, Json, Mutex, Request, Router, StatusCode, TRACE_ID_HEADER, }; -use crate::ai_serving::CODEX_OPENAI_IMAGE_INTERNAL_MODEL; use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY}; use aether_data::repository::auth::{ InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot, @@ -55,13 +54,9 @@ async fn gateway_executes_codex_image_stream_via_local_decision_gate_after_oauth struct SeenExecutionRuntimeStreamRequest { trace_id: String, url: String, - model: String, authorization: String, - x_client_request_id: String, - tool_type: String, - tool_action: String, - tool_partial_images: Option, - request_stream: bool, + headers: serde_json::Value, + body: serde_json::Value, plan_stream: bool, } @@ -180,7 +175,7 @@ async fn gateway_executes_codex_image_stream_via_local_decision_gate_after_oauth None, Some(2), None, - Some(serde_json::json!({"upstream_stream_policy":"force_stream"})), + None, None, None, ) @@ -272,65 +267,26 @@ async fn gateway_executes_codex_image_stream_via_local_decision_gate_after_oauth .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), - model: payload - .get("body") - .and_then(|value| value.get("json_body")) - .and_then(|value| value.get("model")) - .and_then(|value| value.as_str()) - .unwrap_or_default() - .to_string(), authorization: payload .get("headers") .and_then(|value| value.get("authorization")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), - x_client_request_id: payload - .get("headers") - .and_then(|value| value.get("x-client-request-id")) - .and_then(|value| value.as_str()) - .unwrap_or_default() - .to_string(), - tool_type: payload + headers: payload.get("headers").cloned().unwrap_or_default(), + body: payload .get("body") .and_then(|value| value.get("json_body")) - .and_then(|value| value.get("tools")) - .and_then(|value| value.get(0)) - .and_then(|value| value.get("type")) - .and_then(|value| value.as_str()) - .unwrap_or_default() - .to_string(), - tool_action: payload - .get("body") - .and_then(|value| value.get("json_body")) - .and_then(|value| value.get("tools")) - .and_then(|value| value.get(0)) - .and_then(|value| value.get("action")) - .and_then(|value| value.as_str()) - .unwrap_or_default() - .to_string(), - tool_partial_images: payload - .get("body") - .and_then(|value| value.get("json_body")) - .and_then(|value| value.get("tools")) - .and_then(|value| value.get(0)) - .and_then(|value| value.get("partial_images")) - .and_then(|value| value.as_u64()), - request_stream: payload - .get("body") - .and_then(|value| value.get("json_body")) - .and_then(|value| value.get("stream")) - .and_then(|value| value.as_bool()) - .unwrap_or(false), + .cloned() + .unwrap_or_default(), plan_stream: payload .get("stream") .and_then(|value| value.as_bool()) .unwrap_or(false), }); 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: response.output_item.done\\ndata: {\\\"type\\\":\\\"response.output_item.done\\\",\\\"output_index\\\":0,\\\"item\\\":{\\\"id\\\":\\\"ig_123\\\",\\\"type\\\":\\\"image_generation_call\\\",\\\"result\\\":\\\"aGVsbG8=\\\"}}\\n\\n\"}}\n", - "{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: response.completed\\ndata: {\\\"type\\\":\\\"response.completed\\\",\\\"response\\\":{\\\"tool_usage\\\":{\\\"image_gen\\\":{\\\"input_tokens\\\":11,\\\"output_tokens\\\":22,\\\"total_tokens\\\":33}}}}\\n\\n\"}}\n", + "{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"application/json\"}}}\n", + "{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"{\\\"created\\\":1776991097,\\\"data\\\":[{\\\"b64_json\\\":\\\"aGVsbG8=\\\",\\\"revised_prompt\\\":\\\"水墨视觉海报\\\"}],\\\"usage\\\":{\\\"input_tokens\\\":11,\\\"output_tokens\\\":22,\\\"total_tokens\\\":33}}\"}}\n", "{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":41}}}\n", "{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n" ); @@ -397,26 +353,28 @@ async fn gateway_executes_codex_image_stream_via_local_decision_gate_after_oauth ) .header(TRACE_ID_HEADER, "trace-codex-image-stream-local-123") .body( - "{\"model\":\"gpt-image-2\",\"prompt\":\"生成一张中国历史视觉海报\",\"stream\":true,\"partial_images\":1}", + "{\"model\":\"gpt-image-2\",\"prompt\":\"生成一张水墨视觉海报\",\"background\":\"auto\",\"quality\":\"auto\",\"size\":\"auto\",\"stream\":true,\"response_format\":\"b64_json\"}", ) .send() .await .expect("request should succeed"); - assert_eq!(response.status(), StatusCode::OK); - assert_eq!( - response - .headers() - .get(http::header::CONTENT_TYPE) - .and_then(|value| value.to_str().ok()), - Some("text/event-stream") - ); + let response_status = response.status(); + let response_content_type = response + .headers() + .get(http::header::CONTENT_TYPE) + .and_then(|value| value.to_str().ok()) + .map(str::to_string); let response_text = response.text().await.expect("body should read"); - assert!(response_text.contains("event: image_generation.partial_image")); - assert!(response_text.contains("\"type\":\"image_generation.partial_image\"")); - assert!(response_text.contains("\"b64_json\":\"aGVsbG8=\"")); + assert_eq!(response_status, StatusCode::OK, "{response_text}"); + assert_eq!( + response_content_type.as_deref(), + Some("text/event-stream"), + "{response_text}" + ); assert!(response_text.contains("event: image_generation.completed")); assert!(response_text.contains("\"type\":\"image_generation.completed\"")); + assert!(response_text.contains("\"b64_json\":\"aGVsbG8=\"")); assert!(response_text.contains("\"total_tokens\":33")); assert!(!response_text.contains("response.completed")); @@ -444,25 +402,34 @@ async fn gateway_executes_codex_image_stream_via_local_decision_gate_after_oauth ); assert_eq!( seen_execution_runtime_request.url, - "https://chatgpt.com/backend-api/codex/responses" - ); - assert_eq!( - seen_execution_runtime_request.model, - CODEX_OPENAI_IMAGE_INTERNAL_MODEL + "https://chatgpt.com/backend-api/codex/images/generations" ); assert_eq!( seen_execution_runtime_request.authorization, "Bearer refreshed-codex-image-stream-access-token" ); assert_eq!( - seen_execution_runtime_request.x_client_request_id, - "trace-codex-image-stream-local-123" + seen_execution_runtime_request.body, + json!({ + "prompt": "生成一张水墨视觉海报", + "background": "auto", + "model": "gpt-image-2", + "quality": "auto", + "size": "auto" + }) ); - assert_eq!(seen_execution_runtime_request.tool_type, "image_generation"); - assert_eq!(seen_execution_runtime_request.tool_action, "generate"); - assert_eq!(seen_execution_runtime_request.tool_partial_images, Some(1)); - assert!(seen_execution_runtime_request.request_stream); - assert!(seen_execution_runtime_request.plan_stream); + assert_eq!( + seen_execution_runtime_request.headers["user-agent"], + "codex_cli_rs/0.144.1" + ); + assert_eq!( + seen_execution_runtime_request.headers["originator"], + "codex_cli_rs" + ); + for header in ["x-client-request-id", "session-id", "thread-id"] { + assert!(seen_execution_runtime_request.headers.get(header).is_none()); + } + assert!(!seen_execution_runtime_request.plan_stream); gateway_handle.abort(); execution_runtime_handle.abort(); @@ -594,7 +561,7 @@ async fn gateway_bridges_codex_image_sync_json_to_streaming_image_sse_impl() { None, Some(2), None, - Some(serde_json::json!({"upstream_stream_policy":"force_stream"})), + None, None, None, ) @@ -772,8 +739,8 @@ async fn gateway_bridges_codex_image_sync_json_to_streaming_image_sse_impl() { seen_execution_runtime_request.trace_id, "trace-codex-image-stream-json-123" ); - assert!(seen_execution_runtime_request.request_stream); - assert!(seen_execution_runtime_request.plan_stream); + assert!(!seen_execution_runtime_request.request_stream); + assert!(!seen_execution_runtime_request.plan_stream); gateway_handle.abort(); execution_runtime_handle.abort(); @@ -1152,12 +1119,14 @@ async fn gateway_routes_openai_responses_stream_image_intent_to_openai_image_pla seen_plan.url, "https://images.example.com/v1/images/generations" ); - assert!(seen_plan.plan_stream); + assert!(!seen_plan.plan_stream); assert_eq!(seen_plan.auth_header, "Bearer sk-upstream-image-bridge"); - assert_eq!(seen_plan.body_json["stream"], true); - assert_eq!(seen_plan.body_json["input"], "Draw a mountain observatory"); - assert_eq!(seen_plan.body_json["tools"][0]["type"], "image_generation"); - assert_eq!(seen_plan.body_json["tools"][0]["size"], "1024x1024"); + assert_eq!(seen_plan.body_json["model"], "gpt-image-2"); + assert_eq!(seen_plan.body_json["prompt"], "Draw a mountain observatory"); + assert_eq!(seen_plan.body_json["size"], "1024x1024"); + assert!(seen_plan.body_json.get("stream").is_none()); + assert!(seen_plan.body_json.get("input").is_none()); + assert!(seen_plan.body_json.get("tools").is_none()); gateway_handle.abort(); execution_runtime_handle.abort(); diff --git a/apps/aether-gateway/src/tests/ai_execute/stream_cli/compact.rs b/apps/aether-gateway/src/tests/ai_execute/stream_cli/compact.rs index cc0ecb511..c7a0a8841 100644 --- a/apps/aether-gateway/src/tests/ai_execute/stream_cli/compact.rs +++ b/apps/aether-gateway/src/tests/ai_execute/stream_cli/compact.rs @@ -2,7 +2,7 @@ use super::{ any, build_router_with_state, build_state_with_execution_runtime_override, json, run_stream_cli_test, start_server, to_bytes, Arc, Body, Bytes, HeaderName, HeaderValue, Infallible, Json, Mutex, Request, Response, Router, StatusCode, - EXECUTION_PATH_EXECUTION_RUNTIME_STREAM, EXECUTION_PATH_HEADER, TRACE_ID_HEADER, + EXECUTION_PATH_EXECUTION_RUNTIME_SYNC, EXECUTION_PATH_HEADER, TRACE_ID_HEADER, }; use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY}; use aether_data::repository::auth::{ @@ -23,33 +23,37 @@ use aether_data_contracts::repository::provider_catalog::{ use sha2::{Digest, Sha256}; #[test] -fn gateway_executes_openai_responses_compact_stream_via_local_decision_gate_with_local_stream_decision( -) { +fn gateway_executes_openai_responses_compact_as_unary_request() { run_stream_cli_test( - "gateway_executes_openai_responses_compact_stream_via_local_decision_gate_with_local_stream_decision", - gateway_executes_openai_responses_compact_stream_via_local_decision_gate_with_local_stream_decision_impl, + "gateway_executes_openai_responses_compact_as_unary_request", + gateway_executes_openai_responses_compact_as_unary_request_impl, ); } -async fn gateway_executes_openai_responses_compact_stream_via_local_decision_gate_with_local_stream_decision_impl( -) { +async fn gateway_executes_openai_responses_compact_as_unary_request_impl() { #[derive(Debug, Clone)] struct SeenExecutionRuntimeStreamRequest { trace_id: String, url: String, model: String, + content_encoding: String, stream: bool, accept: String, + turn_state: String, authorization: String, + chatgpt_account_id: String, + fedramp: String, + responses_lite: String, + session_id: String, + thread_id: String, + x_client_request_id_present: bool, endpoint_tag: String, conditional_header: String, renamed_header: String, dropped_header_present: bool, - metadata_mode: String, - metadata_source: String, - metadata_origin: String, instructions: String, store_present: bool, + body: serde_json::Value, proxy_node_id: String, transport_profile_id: String, } @@ -71,7 +75,7 @@ async fn gateway_executes_openai_responses_compact_stream_via_local_decision_gat false, Some(serde_json::json!(["openai"])), Some(serde_json::json!(["openai:responses:compact"])), - Some(serde_json::json!(["gpt-5"])), + Some(serde_json::json!(["gpt-5.6-sol"])), api_key_id.to_string(), Some("default".to_string()), true, @@ -82,7 +86,7 @@ async fn gateway_executes_openai_responses_compact_stream_via_local_decision_gat Some(4_102_444_800_i64), Some(serde_json::json!(["openai"])), Some(serde_json::json!(["openai:responses:compact"])), - Some(serde_json::json!(["gpt-5"])), + Some(serde_json::json!(["gpt-5.6-sol"])), ) .expect("auth snapshot should build") } @@ -91,7 +95,7 @@ async fn gateway_executes_openai_responses_compact_stream_via_local_decision_gat StoredMinimalCandidateSelectionRow { provider_id: "provider-openai-compact-local-1".to_string(), provider_name: "openai".to_string(), - provider_type: "custom".to_string(), + provider_type: "codex".to_string(), provider_priority: 10, provider_is_active: true, endpoint_id: "endpoint-openai-compact-local-1".to_string(), @@ -101,7 +105,7 @@ async fn gateway_executes_openai_responses_compact_stream_via_local_decision_gat endpoint_is_active: true, key_id: "key-openai-compact-local-1".to_string(), key_name: "prod".to_string(), - key_auth_type: "bearer".to_string(), + key_auth_type: "oauth".to_string(), key_is_active: true, key_api_formats: Some(vec!["openai:responses:compact".to_string()]), key_allowed_models: None, @@ -110,12 +114,12 @@ async fn gateway_executes_openai_responses_compact_stream_via_local_decision_gat key_global_priority_by_format: Some(serde_json::json!({"openai:responses:compact": 1})), model_id: "model-openai-compact-local-1".to_string(), global_model_id: "global-model-openai-compact-local-1".to_string(), - global_model_name: "gpt-5".to_string(), + global_model_name: "gpt-5.6-sol".to_string(), global_model_mappings: None, global_model_supports_streaming: Some(true), - model_provider_model_name: "gpt-5-upstream".to_string(), + model_provider_model_name: "deployment-production".to_string(), model_provider_model_mappings: Some(vec![StoredProviderModelMapping { - name: "gpt-5-upstream".to_string(), + name: "deployment-production".to_string(), priority: 1, api_formats: Some(vec!["openai:responses:compact".to_string()]), endpoint_ids: None, @@ -131,7 +135,7 @@ async fn gateway_executes_openai_responses_compact_stream_via_local_decision_gat "provider-openai-compact-local-1".to_string(), "openai".to_string(), Some("https://example.com".to_string()), - "custom".to_string(), + "codex".to_string(), ) .expect("provider should build") .with_transport_fields( @@ -161,15 +165,12 @@ async fn gateway_executes_openai_responses_compact_stream_via_local_decision_gat "https://api.openai.example".to_string(), Some(serde_json::json!([ {"action":"set","key":"x-endpoint-tag","value":"openai-compact-local"}, - {"action":"set","key":"x-conditional-tag","value":"header-condition-hit","condition":{"path":"instructions","op":"exists","source":"current"}}, + {"action":"set","key":"x-conditional-tag","value":"header-condition-hit","condition":{"path":"reasoning","op":"exists","source":"current"}}, {"action":"rename","from":"x-client-rename","to":"x-upstream-rename"}, {"action":"drop","key":"x-drop-me"} ])), Some(serde_json::json!([ - {"action":"set","path":"instructions","value":"You are GPT-5.","condition":{"path":"instructions","op":"not_exists","source":"current"}}, - {"action":"set","path":"metadata.mode","value":"safe","condition":{"path":"metadata.mode","op":"not_exists","source":"current"}}, - {"action":"rename","from":"metadata.client","to":"metadata.source"}, - {"action":"set","path":"metadata.origin","value":"from-original","condition":{"path":"metadata.client","op":"exists","source":"original"}}, + {"action":"set","path":"instructions","value":"Use the configured tools.","condition":{"path":"instructions","op":"not_exists","source":"current"}}, {"action":"drop","path":"store"} ])), Some(2), @@ -186,7 +187,7 @@ async fn gateway_executes_openai_responses_compact_stream_via_local_decision_gat "key-openai-compact-local-1".to_string(), "provider-openai-compact-local-1".to_string(), "prod".to_string(), - "bearer".to_string(), + "oauth".to_string(), None, true, ) @@ -198,7 +199,13 @@ async fn gateway_executes_openai_responses_compact_stream_via_local_decision_gat "sk-upstream-openai-compact", ) .expect("api key should encrypt"), - None, + Some( + encrypt_python_fernet_plaintext( + DEVELOPMENT_ENCRYPTION_KEY, + r#"{"account_id":"acc-compact-local-123","is_fedramp":true}"#, + ) + .expect("auth config should encrypt"), + ), None, Some(serde_json::json!({"openai:responses:compact": 1})), None, @@ -242,7 +249,7 @@ async fn gateway_executes_openai_responses_compact_stream_via_local_decision_gat }), ) .route( - "/api/internal/gateway/decision-stream", + "/api/internal/gateway/decision-sync", any(move |_request: Request| { let decision_hits_inner = Arc::clone(&decision_hits_clone); async move { @@ -252,7 +259,7 @@ async fn gateway_executes_openai_responses_compact_stream_via_local_decision_gat }), ) .route( - "/api/internal/gateway/plan-stream", + "/api/internal/gateway/plan-sync", any(move |_request: Request| { let plan_hits_inner = Arc::clone(&plan_hits_clone); async move { @@ -262,7 +269,7 @@ async fn gateway_executes_openai_responses_compact_stream_via_local_decision_gat }), ) .route( - "/api/internal/gateway/report-stream", + "/api/internal/gateway/report-sync", any(move |request: Request| { let seen_report_inner = Arc::clone(&seen_report_clone); async move { @@ -299,135 +306,166 @@ async fn gateway_executes_openai_responses_compact_stream_via_local_decision_gat ); let execution_runtime = Router::new().route( - "/v1/execute/stream", + "/v1/execute/sync", any(move |request: Request| { let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone); async move { let (parts, body) = request.into_parts(); 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"); - *seen_execution_runtime_inner.lock().expect("mutex should lock") = - Some(SeenExecutionRuntimeStreamRequest { - trace_id: parts - .headers - .get(TRACE_ID_HEADER) - .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(), - model: payload - .get("body") - .and_then(|value| value.get("json_body")) - .and_then(|value| value.get("model")) - .and_then(|value| value.as_str()) - .unwrap_or_default() - .to_string(), - stream: payload - .get("stream") - .and_then(|value| value.as_bool()) - .unwrap_or(false), - accept: payload - .get("headers") - .and_then(|value| value.get("accept")) - .and_then(|value| value.as_str()) - .unwrap_or_default() - .to_string(), - authorization: payload - .get("headers") - .and_then(|value| value.get("authorization")) - .and_then(|value| value.as_str()) - .unwrap_or_default() - .to_string(), - endpoint_tag: payload - .get("headers") - .and_then(|value| value.get("x-endpoint-tag")) - .and_then(|value| value.as_str()) - .unwrap_or_default() - .to_string(), - conditional_header: payload - .get("headers") - .and_then(|value| value.get("x-conditional-tag")) - .and_then(|value| value.as_str()) - .unwrap_or_default() - .to_string(), - renamed_header: payload - .get("headers") - .and_then(|value| value.get("x-upstream-rename")) - .and_then(|value| value.as_str()) - .unwrap_or_default() - .to_string(), - dropped_header_present: payload - .get("headers") - .and_then(|value| value.get("x-drop-me")) - .is_some(), - metadata_mode: payload - .get("body") - .and_then(|value| value.get("json_body")) - .and_then(|value| value.get("metadata")) - .and_then(|value| value.get("mode")) - .and_then(|value| value.as_str()) - .unwrap_or_default() - .to_string(), - metadata_source: payload - .get("body") - .and_then(|value| value.get("json_body")) - .and_then(|value| value.get("metadata")) - .and_then(|value| value.get("source")) - .and_then(|value| value.as_str()) - .unwrap_or_default() - .to_string(), - metadata_origin: payload - .get("body") - .and_then(|value| value.get("json_body")) - .and_then(|value| value.get("metadata")) - .and_then(|value| value.get("origin")) - .and_then(|value| value.as_str()) - .unwrap_or_default() - .to_string(), - instructions: payload - .get("body") - .and_then(|value| value.get("json_body")) - .and_then(|value| value.get("instructions")) - .and_then(|value| value.as_str()) - .unwrap_or_default() - .to_string(), - store_present: payload - .get("body") - .and_then(|value| value.get("json_body")) - .and_then(|value| value.get("store")) - .is_some(), - proxy_node_id: payload - .get("proxy") - .and_then(|value| value.get("node_id")) - .and_then(|value| value.as_str()) - .unwrap_or_default() - .to_string(), - transport_profile_id: payload - .get("transport_profile").and_then(|value| value.get("profile_id")) - .and_then(|value| value.as_str()) - .unwrap_or_default() - .to_string(), - }); - let stream = concat!( - "{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n", - "{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: response.completed\\ndata: {\\\"type\\\":\\\"response.completed\\\",\\\"response\\\":{\\\"id\\\":\\\"resp-compact-local-123\\\",\\\"object\\\":\\\"response\\\",\\\"model\\\":\\\"gpt-5-upstream\\\",\\\"output\\\":[]}}\\n\\n\"}}\n", - "{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":41,\"ttfb_ms\":11}}}\n", - "{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n" - ); - let mut response = Response::builder() - .status(StatusCode::OK) - .body(Body::from(stream)) - .expect("response should build"); - response.headers_mut().insert( - http::header::CONTENT_TYPE, - HeaderValue::from_static("application/x-ndjson"), - ); - response + let payload: serde_json::Value = serde_json::from_slice(&raw_body) + .expect("execution runtime payload should parse"); + *seen_execution_runtime_inner + .lock() + .expect("mutex should lock") = Some(SeenExecutionRuntimeStreamRequest { + trace_id: parts + .headers + .get(TRACE_ID_HEADER) + .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(), + model: payload + .get("body") + .and_then(|value| value.get("json_body")) + .and_then(|value| value.get("model")) + .and_then(|value| value.as_str()) + .unwrap_or_default() + .to_string(), + content_encoding: payload + .get("content_encoding") + .and_then(|value| value.as_str()) + .unwrap_or_default() + .to_string(), + stream: payload + .get("stream") + .and_then(|value| value.as_bool()) + .unwrap_or(false), + accept: payload + .get("headers") + .and_then(|value| value.get("accept")) + .and_then(|value| value.as_str()) + .unwrap_or_default() + .to_string(), + turn_state: payload + .get("headers") + .and_then(|value| value.get("x-codex-turn-state")) + .and_then(|value| value.as_str()) + .unwrap_or_default() + .to_string(), + authorization: payload + .get("headers") + .and_then(|value| value.get("authorization")) + .and_then(|value| value.as_str()) + .unwrap_or_default() + .to_string(), + chatgpt_account_id: payload + .get("headers") + .and_then(|value| value.get("chatgpt-account-id")) + .and_then(|value| value.as_str()) + .unwrap_or_default() + .to_string(), + fedramp: payload + .get("headers") + .and_then(|value| value.get("x-openai-fedramp")) + .and_then(|value| value.as_str()) + .unwrap_or_default() + .to_string(), + responses_lite: payload + .get("headers") + .and_then(|value| value.get("x-openai-internal-codex-responses-lite")) + .and_then(|value| value.as_str()) + .unwrap_or_default() + .to_string(), + session_id: payload + .get("headers") + .and_then(|value| value.get("session-id")) + .and_then(|value| value.as_str()) + .unwrap_or_default() + .to_string(), + thread_id: payload + .get("headers") + .and_then(|value| value.get("thread-id")) + .and_then(|value| value.as_str()) + .unwrap_or_default() + .to_string(), + x_client_request_id_present: payload + .get("headers") + .and_then(|value| value.get("x-client-request-id")) + .is_some(), + endpoint_tag: payload + .get("headers") + .and_then(|value| value.get("x-endpoint-tag")) + .and_then(|value| value.as_str()) + .unwrap_or_default() + .to_string(), + conditional_header: payload + .get("headers") + .and_then(|value| value.get("x-conditional-tag")) + .and_then(|value| value.as_str()) + .unwrap_or_default() + .to_string(), + renamed_header: payload + .get("headers") + .and_then(|value| value.get("x-upstream-rename")) + .and_then(|value| value.as_str()) + .unwrap_or_default() + .to_string(), + dropped_header_present: payload + .get("headers") + .and_then(|value| value.get("x-drop-me")) + .is_some(), + instructions: payload + .get("body") + .and_then(|value| value.get("json_body")) + .and_then(|value| value.get("instructions")) + .and_then(|value| value.as_str()) + .unwrap_or_default() + .to_string(), + store_present: payload + .get("body") + .and_then(|value| value.get("json_body")) + .and_then(|value| value.get("store")) + .is_some(), + body: payload + .get("body") + .and_then(|value| value.get("json_body")) + .cloned() + .unwrap_or(serde_json::Value::Null), + proxy_node_id: payload + .get("proxy") + .and_then(|value| value.get("node_id")) + .and_then(|value| value.as_str()) + .unwrap_or_default() + .to_string(), + transport_profile_id: payload + .get("transport_profile") + .and_then(|value| value.get("profile_id")) + .and_then(|value| value.as_str()) + .unwrap_or_default() + .to_string(), + }); + Json(json!({ + "request_id": "trace-openai-compact-local-123", + "status_code": 200, + "headers": { + "content-type": "application/json", + "x-codex-turn-state": "turn-state-compact-123" + }, + "body": { + "json_body": { + "output": [{ + "type": "compaction", + "id": "cmp-compact-local-123", + "encrypted_content": "encrypted-compact-history" + }] + } + }, + "telemetry": {"elapsed_ms": 41, "ttfb_ms": 11} + })) } }), ); @@ -474,23 +512,40 @@ async fn gateway_executes_openai_responses_compact_stream_via_local_decision_gat ) .header("x-client-rename", "rename-openai-compact") .header("x-drop-me", "drop-openai-compact") + .header("x-codex-turn-state", "turn-state-inbound-123") + .header("session-id", "session-compact-local-123") + .header("thread-id", "thread-compact-local-123") .header(TRACE_ID_HEADER, "trace-openai-compact-local-123") - .body("{\"model\":\"gpt-5\",\"input\":\"hello\",\"stream\":true,\"metadata\":{\"client\":\"desktop-openai-compact\"},\"store\":false}") + .body(r#"{"model":"gpt-5.6-sol","input":"hello","client_metadata":{"origin":"codex"},"include":["reasoning.encrypted_content"],"store":false,"stream":true,"stream_options":{"reasoning_summary_delivery":"sequential_cutoff"},"tool_choice":"auto","parallel_tool_calls":true,"reasoning":{"effort":"high"},"text":{"verbosity":"medium"},"tools":[{"type":"function","name":"lookup","parameters":{"type":"object"}}],"prompt_cache_key":"session:compact-e2e"}"#) .send() .await .expect("request should succeed"); - assert_eq!(response.status(), StatusCode::OK); + if response.status() != StatusCode::OK { + let status = response.status(); + let headers = response.headers().clone(); + let body = response.text().await.expect("error body should read"); + panic!("Compact request failed: status={status}, headers={headers:?}, body={body}"); + } + assert_eq!( + response + .headers() + .get("x-codex-turn-state") + .and_then(|value| value.to_str().ok()), + Some("turn-state-compact-123") + ); assert_eq!( response .headers() .get(EXECUTION_PATH_HEADER) .and_then(|value| value.to_str().ok()), - Some(EXECUTION_PATH_EXECUTION_RUNTIME_STREAM) + Some(EXECUTION_PATH_EXECUTION_RUNTIME_SYNC) + ); + let body: serde_json::Value = response.json().await.expect("body should parse"); + assert_eq!( + body["output"][0]["encrypted_content"], + "encrypted-compact-history" ); - let body = response.text().await.expect("body should read"); - assert!(body.contains("event: response.completed")); - assert!(body.contains("\"model\":\"gpt-5-upstream\"")); let seen_execution_runtime_request = seen_execution_runtime .lock() @@ -505,13 +560,36 @@ async fn gateway_executes_openai_responses_compact_stream_via_local_decision_gat seen_execution_runtime_request.url, "https://api.openai.example/custom/v1/responses/compact" ); - assert_eq!(seen_execution_runtime_request.model, "gpt-5-upstream"); - assert!(seen_execution_runtime_request.stream); - assert_eq!(seen_execution_runtime_request.accept, "text/event-stream"); + assert_eq!( + seen_execution_runtime_request.model, + "deployment-production" + ); + assert!(seen_execution_runtime_request.content_encoding.is_empty()); + assert!(!seen_execution_runtime_request.stream); + assert_ne!(seen_execution_runtime_request.accept, "text/event-stream"); + assert_eq!( + seen_execution_runtime_request.turn_state, + "turn-state-inbound-123" + ); assert_eq!( seen_execution_runtime_request.authorization, "Bearer sk-upstream-openai-compact" ); + assert_eq!( + seen_execution_runtime_request.chatgpt_account_id, + "acc-compact-local-123" + ); + assert_eq!(seen_execution_runtime_request.fedramp, "true"); + assert_eq!(seen_execution_runtime_request.responses_lite, "true"); + assert_eq!( + seen_execution_runtime_request.session_id, + "session-compact-local-123" + ); + assert_eq!( + seen_execution_runtime_request.thread_id, + "thread-compact-local-123" + ); + assert!(!seen_execution_runtime_request.x_client_request_id_present); assert_eq!( seen_execution_runtime_request.endpoint_tag, "openai-compact-local" @@ -525,20 +603,61 @@ async fn gateway_executes_openai_responses_compact_stream_via_local_decision_gat "rename-openai-compact" ); assert!(!seen_execution_runtime_request.dropped_header_present); - assert_eq!( - seen_execution_runtime_request.instructions, - "You are GPT-5." - ); - assert_eq!(seen_execution_runtime_request.metadata_mode, "safe"); - assert_eq!( - seen_execution_runtime_request.metadata_source, - "desktop-openai-compact" - ); - assert_eq!( - seen_execution_runtime_request.metadata_origin, - "from-original" - ); + assert!(seen_execution_runtime_request.instructions.is_empty()); assert!(!seen_execution_runtime_request.store_present); + for field in [ + "client_metadata", + "include", + "store", + "stream", + "stream_options", + "tool_choice", + ] { + assert!( + seen_execution_runtime_request.body.get(field).is_none(), + "Compact request must omit {field}" + ); + } + assert_eq!( + seen_execution_runtime_request.body["parallel_tool_calls"], + json!(false) + ); + assert_eq!( + seen_execution_runtime_request.body["reasoning"]["effort"], + json!("high") + ); + assert_eq!( + seen_execution_runtime_request.body["text"]["verbosity"], + json!("medium") + ); + assert_eq!( + seen_execution_runtime_request.body["reasoning"]["context"], + json!("all_turns") + ); + assert_eq!( + seen_execution_runtime_request.body["input"][0]["type"], + json!("additional_tools") + ); + assert_eq!( + seen_execution_runtime_request.body["input"][0]["tools"][0]["name"], + json!("lookup") + ); + assert_eq!( + seen_execution_runtime_request.body["input"][1]["role"], + json!("developer") + ); + assert_eq!( + seen_execution_runtime_request.body["input"][1]["content"][0]["text"], + json!("Use the configured tools.") + ); + assert_eq!( + seen_execution_runtime_request.body["input"][2]["content"][0]["text"], + json!("hello") + ); + assert_eq!( + seen_execution_runtime_request.body["prompt_cache_key"], + json!("session:compact-e2e") + ); assert_eq!( seen_execution_runtime_request.proxy_node_id, "proxy-node-openai-compact-local" @@ -558,7 +677,7 @@ async fn gateway_executes_openai_responses_compact_stream_via_local_decision_gat tokio::time::sleep(std::time::Duration::from_millis(100)).await; assert!( !*seen_report.lock().expect("mutex should lock"), - "report-stream should stay local when request candidate persistence is available" + "report-sync should stay local when request candidate persistence is available" ); assert_eq!(*decision_hits.lock().expect("mutex should lock"), 0); diff --git a/apps/aether-gateway/src/tests/ai_execute/stream_cli/direct.rs b/apps/aether-gateway/src/tests/ai_execute/stream_cli/direct.rs index 2cd551a5a..d885cc479 100644 --- a/apps/aether-gateway/src/tests/ai_execute/stream_cli/direct.rs +++ b/apps/aether-gateway/src/tests/ai_execute/stream_cli/direct.rs @@ -38,10 +38,24 @@ async fn gateway_executes_codex_cli_stream_via_local_decision_gate_after_oauth_r trace_id: String, url: String, model: String, + content_encoding: String, stream: bool, accept: String, authorization: String, + chatgpt_account_id: String, + fedramp: String, x_client_request_id: String, + session_id: String, + thread_id: String, + prompt_cache_key: String, + responses_lite: String, + has_top_level_tools: bool, + has_top_level_instructions: bool, + has_additional_tools: bool, + parallel_tool_calls: bool, + reasoning_effort: String, + reasoning_context: String, + has_compaction_trigger: bool, } #[derive(Debug, Clone)] @@ -74,7 +88,7 @@ async fn gateway_executes_codex_cli_stream_via_local_decision_gate_after_oauth_r false, Some(serde_json::json!(["openai", "codex"])), Some(serde_json::json!(["openai:responses"])), - Some(serde_json::json!(["gpt-5.4"])), + Some(serde_json::json!(["gpt-5.6-sol"])), api_key_id.to_string(), Some("default".to_string()), true, @@ -85,7 +99,7 @@ async fn gateway_executes_codex_cli_stream_via_local_decision_gate_after_oauth_r Some(4_102_444_800_i64), Some(serde_json::json!(["openai", "codex"])), Some(serde_json::json!(["openai:responses"])), - Some(serde_json::json!(["gpt-5.4"])), + Some(serde_json::json!(["gpt-5.6-sol"])), ) .expect("auth snapshot should build") } @@ -113,12 +127,12 @@ async fn gateway_executes_codex_cli_stream_via_local_decision_gate_after_oauth_r key_global_priority_by_format: Some(serde_json::json!({"openai:responses": 1})), model_id: "model-codex-cli-stream-local-1".to_string(), global_model_id: "global-model-codex-cli-stream-local-1".to_string(), - global_model_name: "gpt-5.4".to_string(), + global_model_name: "gpt-5.6-sol".to_string(), global_model_mappings: None, global_model_supports_streaming: Some(true), - model_provider_model_name: "gpt-5.4".to_string(), + model_provider_model_name: "gpt-5.6-sol".to_string(), model_provider_model_mappings: Some(vec![StoredProviderModelMapping { - name: "gpt-5.4".to_string(), + name: "gpt-5.6-sol".to_string(), priority: 1, api_formats: Some(vec!["openai:responses".to_string()]), endpoint_ids: None, @@ -176,7 +190,7 @@ async fn gateway_executes_codex_cli_stream_via_local_decision_gate_after_oauth_r fn sample_provider_catalog_key() -> StoredProviderCatalogKey { let encrypted_auth_config = encrypt_python_fernet_plaintext( DEVELOPMENT_ENCRYPTION_KEY, - r#"{"provider_type":"codex","refresh_token":"rt-codex-stream-local-123"}"#, + r#"{"provider_type":"codex","refresh_token":"rt-codex-stream-local-123","account_id":"acc-codex-stream-local-123","is_fedramp":true}"#, ) .expect("auth config should encrypt"); StoredProviderCatalogKey::new( @@ -367,6 +381,11 @@ async fn gateway_executes_codex_cli_stream_via_local_decision_gate_after_oauth_r .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), + content_encoding: payload + .get("content_encoding") + .and_then(|value| value.as_str()) + .unwrap_or_default() + .to_string(), stream: payload .get("stream") .and_then(|value| value.as_bool()) @@ -383,16 +402,106 @@ async fn gateway_executes_codex_cli_stream_via_local_decision_gate_after_oauth_r .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), + chatgpt_account_id: payload + .get("headers") + .and_then(|value| value.get("chatgpt-account-id")) + .and_then(|value| value.as_str()) + .unwrap_or_default() + .to_string(), + fedramp: payload + .get("headers") + .and_then(|value| value.get("x-openai-fedramp")) + .and_then(|value| value.as_str()) + .unwrap_or_default() + .to_string(), x_client_request_id: payload .get("headers") .and_then(|value| value.get("x-client-request-id")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), + session_id: payload + .get("headers") + .and_then(|value| value.get("session-id")) + .and_then(|value| value.as_str()) + .unwrap_or_default() + .to_string(), + thread_id: payload + .get("headers") + .and_then(|value| value.get("thread-id")) + .and_then(|value| value.as_str()) + .unwrap_or_default() + .to_string(), + prompt_cache_key: payload + .get("body") + .and_then(|value| value.get("json_body")) + .and_then(|value| value.get("prompt_cache_key")) + .and_then(|value| value.as_str()) + .unwrap_or_default() + .to_string(), + responses_lite: payload + .get("headers") + .and_then(|value| { + value.get("x-openai-internal-codex-responses-lite") + }) + .and_then(|value| value.as_str()) + .unwrap_or_default() + .to_string(), + has_top_level_tools: payload + .get("body") + .and_then(|value| value.get("json_body")) + .is_some_and(|body| body.get("tools").is_some()), + has_top_level_instructions: payload + .get("body") + .and_then(|value| value.get("json_body")) + .is_some_and(|body| body.get("instructions").is_some()), + has_additional_tools: payload + .get("body") + .and_then(|value| value.get("json_body")) + .and_then(|value| value.get("input")) + .and_then(|value| value.as_array()) + .and_then(|input| input.first()) + .and_then(|item| item.get("type")) + .and_then(|value| value.as_str()) + == Some("additional_tools"), + parallel_tool_calls: payload + .get("body") + .and_then(|value| value.get("json_body")) + .and_then(|value| value.get("parallel_tool_calls")) + .and_then(|value| value.as_bool()) + .unwrap_or(true), + reasoning_effort: payload + .get("body") + .and_then(|value| value.get("json_body")) + .and_then(|value| value.get("reasoning")) + .and_then(|value| value.get("effort")) + .and_then(|value| value.as_str()) + .unwrap_or_default() + .to_string(), + reasoning_context: payload + .get("body") + .and_then(|value| value.get("json_body")) + .and_then(|value| value.get("reasoning")) + .and_then(|value| value.get("context")) + .and_then(|value| value.as_str()) + .unwrap_or_default() + .to_string(), + has_compaction_trigger: payload + .get("body") + .and_then(|value| value.get("json_body")) + .and_then(|value| value.get("input")) + .and_then(|value| value.as_array()) + .is_some_and(|input| { + input.iter().any(|item| { + item.get("type").and_then(|value| value.as_str()) + == Some("compaction_trigger") + }) + }), }); 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: response.completed\\ndata: {\\\"type\\\":\\\"response.completed\\\",\\\"response\\\":{\\\"id\\\":\\\"resp_codex_cli_stream_local_123\\\",\\\"object\\\":\\\"response\\\",\\\"model\\\":\\\"gpt-5.4\\\",\\\"status\\\":\\\"completed\\\",\\\"usage\\\":{\\\"input_tokens\\\":1,\\\"output_tokens\\\":2,\\\"total_tokens\\\":3}}}\\n\\n\"}}\n", + "{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: response.output_item.done\\ndata: {\\\"type\\\":\\\"response.output_item.done\\\",\\\"item\\\":{\\\"type\\\":\\\"compaction\\\",\\\"encrypted_content\\\":\\\"ENCRYPTED_CONTEXT_COMPACTION_SUMMARY\\\"}}\\n\\n\"}}\n", + "{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: response.completed\\ndata: {\\\"type\\\":\\\"response.completed\\\",\\\"response\\\":{\\\"id\\\":\\\"resp_codex_cli_stream_local_123\\\",\\\"object\\\":\\\"response\\\",\\\"model\\\":\\\"gpt-5.6-sol\\\",\\\"status\\\":\\\"completed\\\",\\\"usage\\\":{\\\"input_tokens\\\":1,\\\"output_tokens\\\":2,\\\"total_tokens\\\":3}}}\\n\\n\"}}\n", "{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":41}}}\n", "{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n" ); @@ -469,8 +578,16 @@ async fn gateway_executes_codex_cli_stream_via_local_decision_gate_after_oauth_r http::header::AUTHORIZATION, format!("Bearer {client_api_key}"), ) + .header("session-id", "session-codex-stream-local-123") + .header("thread-id", "thread-codex-stream-local-123") + .header( + "x-client-request-id", + "thread-codex-stream-local-123", + ) .header(TRACE_ID_HEADER, "trace-codex-cli-stream-local-123") - .body("{\"model\":\"gpt-5.4\",\"input\":\"hello\",\"stream\":true}") + .body( + r#"{"model":"gpt-5.6-sol","instructions":"Use the configured tools.","input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"compact"}]},{"type":"compaction_trigger"}],"tools":[{"type":"function","name":"lookup","parameters":{"type":"object"}}],"parallel_tool_calls":true,"prompt_cache_key":"thread-codex-stream-local-123","client_metadata":{"session_id":"session-codex-stream-local-123","thread_id":"thread-codex-stream-local-123"},"stream":true}"#, + ) .send() .await .expect("request should succeed"); @@ -478,9 +595,20 @@ async fn gateway_executes_codex_cli_stream_via_local_decision_gate_after_oauth_r assert_eq!(response.status(), StatusCode::OK); let response_body = strip_sse_keepalive_comments(&response.text().await.expect("body should read")); + assert!(response_body.contains("event: response.output_item.done\n")); + assert!(response_body.contains("\"type\":\"compaction\"")); + assert!(response_body.contains("ENCRYPTED_CONTEXT_COMPACTION_SUMMARY")); let data_line = response_body .lines() - .find_map(|line| line.strip_prefix("data: ")) + .filter_map(|line| line.strip_prefix("data: ")) + .find(|line| { + serde_json::from_str::(line) + .ok() + .and_then(|event| event.get("type").cloned()) + .and_then(|value| value.as_str().map(ToOwned::to_owned)) + .as_deref() + == Some("response.completed") + }) .expect("completed event data should exist"); let completed_event: serde_json::Value = serde_json::from_str(data_line).expect("completed event should parse"); @@ -495,7 +623,7 @@ async fn gateway_executes_codex_cli_stream_via_local_decision_gate_after_oauth_r "response": { "id": "resp_codex_cli_stream_local_123", "object": "response", - "model": "gpt-5.4", + "model": "gpt-5.6-sol", "status": "completed", "usage": { "input_tokens": 1, @@ -542,7 +670,8 @@ async fn gateway_executes_codex_cli_stream_via_local_decision_gate_after_oauth_r seen_execution_runtime_request.url, "https://chatgpt.com/backend-api/codex/responses" ); - assert_eq!(seen_execution_runtime_request.model, "gpt-5.4"); + assert_eq!(seen_execution_runtime_request.model, "gpt-5.6-sol"); + assert_eq!(seen_execution_runtime_request.content_encoding, "zstd"); assert!(seen_execution_runtime_request.stream); assert_eq!(seen_execution_runtime_request.accept, "text/event-stream"); assert_eq!( @@ -550,9 +679,37 @@ async fn gateway_executes_codex_cli_stream_via_local_decision_gate_after_oauth_r "Bearer refreshed-codex-stream-access-token" ); assert_eq!( - seen_execution_runtime_request.x_client_request_id, - "trace-codex-cli-stream-local-123" + seen_execution_runtime_request.chatgpt_account_id, + "acc-codex-stream-local-123" ); + assert_eq!(seen_execution_runtime_request.fedramp, "true"); + assert_eq!( + seen_execution_runtime_request.x_client_request_id, + "thread-codex-stream-local-123" + ); + assert_eq!( + seen_execution_runtime_request.session_id, + "session-codex-stream-local-123" + ); + assert_eq!( + seen_execution_runtime_request.thread_id, + "thread-codex-stream-local-123" + ); + assert_eq!( + seen_execution_runtime_request.thread_id, + seen_execution_runtime_request.prompt_cache_key + ); + assert_eq!(seen_execution_runtime_request.responses_lite, "true"); + assert!(!seen_execution_runtime_request.has_top_level_tools); + assert!(!seen_execution_runtime_request.has_top_level_instructions); + assert!(seen_execution_runtime_request.has_additional_tools); + assert!(!seen_execution_runtime_request.parallel_tool_calls); + assert_eq!(seen_execution_runtime_request.reasoning_effort, "low"); + assert_eq!( + seen_execution_runtime_request.reasoning_context, + "all_turns" + ); + assert!(seen_execution_runtime_request.has_compaction_trigger); let stored_candidates = request_candidate_repository .list_by_request_id("trace-codex-cli-stream-local-123") diff --git a/apps/aether-gateway/src/tests/ai_execute/sync/cli.rs b/apps/aether-gateway/src/tests/ai_execute/sync/cli.rs index 803bb6dbb..7c8f16999 100644 --- a/apps/aether-gateway/src/tests/ai_execute/sync/cli.rs +++ b/apps/aether-gateway/src/tests/ai_execute/sync/cli.rs @@ -3223,6 +3223,9 @@ async fn gateway_executes_codex_cli_sync_via_local_decision_gate_after_oauth_ref model: String, authorization: String, x_client_request_id: String, + session_id: String, + thread_id: String, + prompt_cache_key: String, stream_present: bool, plan_stream: bool, } @@ -3509,6 +3512,25 @@ async fn gateway_executes_codex_cli_sync_via_local_decision_gate_after_oauth_ref .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), + session_id: payload + .get("headers") + .and_then(|value| value.get("session-id")) + .and_then(|value| value.as_str()) + .unwrap_or_default() + .to_string(), + thread_id: payload + .get("headers") + .and_then(|value| value.get("thread-id")) + .and_then(|value| value.as_str()) + .unwrap_or_default() + .to_string(), + prompt_cache_key: payload + .get("body") + .and_then(|value| value.get("json_body")) + .and_then(|value| value.get("prompt_cache_key")) + .and_then(|value| value.as_str()) + .unwrap_or_default() + .to_string(), stream_present: payload .get("body") .and_then(|value| value.get("json_body")) @@ -3646,7 +3668,16 @@ async fn gateway_executes_codex_cli_sync_via_local_decision_gate_after_oauth_ref ); assert_eq!( seen_execution_runtime_request.x_client_request_id, - "trace-codex-cli-local-123" + seen_execution_runtime_request.thread_id + ); + assert_eq!( + seen_execution_runtime_request.session_id, + seen_execution_runtime_request.thread_id + ); + assert!(seen_execution_runtime_request.prompt_cache_key.is_empty()); + assert_ne!( + seen_execution_runtime_request.thread_id, + seen_execution_runtime_request.trace_id ); assert!(seen_execution_runtime_request.stream_present); assert!(seen_execution_runtime_request.plan_stream); diff --git a/apps/aether-gateway/src/tests/ai_execute/sync/image.rs b/apps/aether-gateway/src/tests/ai_execute/sync/image.rs index 0277939c6..915de228e 100644 --- a/apps/aether-gateway/src/tests/ai_execute/sync/image.rs +++ b/apps/aether-gateway/src/tests/ai_execute/sync/image.rs @@ -2,7 +2,6 @@ use super::{ any, build_router_with_state, build_state_with_execution_runtime_override, json, start_server, to_bytes, Arc, Body, Json, Mutex, Request, Router, StatusCode, TRACE_ID_HEADER, }; -use crate::ai_serving::CODEX_OPENAI_IMAGE_INTERNAL_MODEL; use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY}; use aether_data::repository::auth::{ InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot, @@ -411,10 +410,10 @@ async fn gateway_converts_gemini_image_sync_to_openai_image_provider_impl() { url: String, authorization: String, model: String, - action: String, prompt: String, image_url: String, request_stream: bool, + body_stream: Option, } fn hash_api_key(value: &str) -> String { @@ -574,12 +573,6 @@ async fn gateway_converts_gemini_image_sync_to_openai_image_provider_impl() { .and_then(|value| value.get("json_body")) .cloned() .unwrap_or_else(|| json!({})); - let content = body_json - .get("input") - .and_then(|value| value.get(0)) - .and_then(|value| value.get("content")) - .cloned() - .unwrap_or_else(|| json!([])); *seen_execution_runtime_inner .lock() .expect("mutex should lock") = Some(SeenExecutionRuntimeSyncRequest { @@ -605,39 +598,22 @@ async fn gateway_converts_gemini_image_sync_to_openai_image_provider_impl() { .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), - action: body_json - .get("tools") - .and_then(|value| value.get(0)) - .and_then(|value| value.get("action")) + prompt: body_json + .get("prompt") .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), - prompt: content - .as_array() - .into_iter() - .flatten() - .find(|item| { - item.get("type").and_then(|value| value.as_str()) == Some("input_text") - }) - .and_then(|item| item.get("text")) + image_url: body_json + .get("image") + .and_then(|value| value.get("image_url")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), - image_url: content - .as_array() - .into_iter() - .flatten() - .find(|item| { - item.get("type").and_then(|value| value.as_str()) == Some("input_image") - }) - .and_then(|item| item.get("image_url")) - .and_then(|value| value.as_str()) - .unwrap_or_default() - .to_string(), - request_stream: body_json + request_stream: payload .get("stream") .and_then(|value| value.as_bool()) .unwrap_or(true), + body_stream: body_json.get("stream").and_then(serde_json::Value::as_bool), }); Json(json!({ "request_id": "trace-gemini-image-to-openai-123", @@ -647,21 +623,16 @@ async fn gateway_converts_gemini_image_sync_to_openai_image_provider_impl() { }, "body": { "json_body": { - "id": "resp_img_bridge_123", - "object": "response", + "created": 1776839946, "model": "gpt-image-2-upstream", - "status": "completed", "usage": { "input_tokens": 3, "output_tokens": 4, "total_tokens": 7 }, - "output": [{ - "type": "image_generation_call", - "status": "completed", - "output_format": "png", + "data": [{ "revised_prompt": "converted gemini prompt", - "result": "aGVsbG8=" + "b64_json": "aGVsbG8=" }] } }, @@ -750,14 +721,13 @@ async fn gateway_converts_gemini_image_sync_to_openai_image_provider_impl() { ); assert_eq!( seen_execution_runtime_request.url, - "https://api.openai.com/v1/images/generations" + "https://api.openai.com/v1/images/edits" ); assert_eq!( seen_execution_runtime_request.authorization, "Bearer sk-upstream-openai-image" ); assert_eq!(seen_execution_runtime_request.model, "gpt-image-2-upstream"); - assert_eq!(seen_execution_runtime_request.action, "edit"); assert_eq!( seen_execution_runtime_request.prompt, "Change the background" @@ -767,6 +737,7 @@ async fn gateway_converts_gemini_image_sync_to_openai_image_provider_impl() { "data:image/png;base64,aGVsbG8=" ); assert!(!seen_execution_runtime_request.request_stream); + assert_eq!(seen_execution_runtime_request.body_stream, None); gateway_handle.abort(); execution_runtime_handle.abort(); @@ -785,18 +756,9 @@ async fn gateway_executes_codex_image_sync_via_local_decision_gate_after_oauth_r struct SeenExecutionRuntimeSyncRequest { trace_id: String, url: String, - model: String, authorization: String, - x_client_request_id: String, - prompt: String, - content_is_string: bool, - tool_type: String, - tool_size: String, - tool_quality: String, - tool_background: String, - tool_choice_type: String, - tool_has_n: bool, - request_stream: bool, + headers: serde_json::Value, + body: serde_json::Value, plan_stream: bool, } @@ -1011,98 +973,18 @@ async fn gateway_executes_codex_image_sync_via_local_decision_gate_after_oauth_r .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), - model: payload - .get("body") - .and_then(|value| value.get("json_body")) - .and_then(|value| value.get("model")) - .and_then(|value| value.as_str()) - .unwrap_or_default() - .to_string(), authorization: payload .get("headers") .and_then(|value| value.get("authorization")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), - x_client_request_id: payload - .get("headers") - .and_then(|value| value.get("x-client-request-id")) - .and_then(|value| value.as_str()) - .unwrap_or_default() - .to_string(), - prompt: payload + headers: payload.get("headers").cloned().unwrap_or_default(), + body: payload .get("body") .and_then(|value| value.get("json_body")) - .and_then(|value| value.get("input")) - .and_then(|value| value.get(0)) - .and_then(|value| value.get("content")) - .and_then(|value| value.as_str()) - .unwrap_or_default() - .to_string(), - content_is_string: payload - .get("body") - .and_then(|value| value.get("json_body")) - .and_then(|value| value.get("input")) - .and_then(|value| value.get(0)) - .and_then(|value| value.get("content")) - .is_some_and(|value| value.is_string()), - tool_type: payload - .get("body") - .and_then(|value| value.get("json_body")) - .and_then(|value| value.get("tools")) - .and_then(|value| value.get(0)) - .and_then(|value| value.get("type")) - .and_then(|value| value.as_str()) - .unwrap_or_default() - .to_string(), - tool_size: payload - .get("body") - .and_then(|value| value.get("json_body")) - .and_then(|value| value.get("tools")) - .and_then(|value| value.get(0)) - .and_then(|value| value.get("size")) - .and_then(|value| value.as_str()) - .unwrap_or_default() - .to_string(), - tool_quality: payload - .get("body") - .and_then(|value| value.get("json_body")) - .and_then(|value| value.get("tools")) - .and_then(|value| value.get(0)) - .and_then(|value| value.get("quality")) - .and_then(|value| value.as_str()) - .unwrap_or_default() - .to_string(), - tool_background: payload - .get("body") - .and_then(|value| value.get("json_body")) - .and_then(|value| value.get("tools")) - .and_then(|value| value.get(0)) - .and_then(|value| value.get("background")) - .and_then(|value| value.as_str()) - .unwrap_or_default() - .to_string(), - tool_choice_type: payload - .get("body") - .and_then(|value| value.get("json_body")) - .and_then(|value| value.get("tool_choice")) - .and_then(|value| value.get("type")) - .and_then(|value| value.as_str()) - .unwrap_or_default() - .to_string(), - tool_has_n: payload - .get("body") - .and_then(|value| value.get("json_body")) - .and_then(|value| value.get("tools")) - .and_then(|value| value.get(0)) - .and_then(|value| value.as_object()) - .is_some_and(|object| object.contains_key("n")), - request_stream: payload - .get("body") - .and_then(|value| value.get("json_body")) - .and_then(|value| value.get("stream")) - .and_then(|value| value.as_bool()) - .unwrap_or(false), + .cloned() + .unwrap_or_default(), plan_stream: payload .get("stream") .and_then(|value| value.as_bool()) @@ -1112,20 +994,11 @@ async fn gateway_executes_codex_image_sync_via_local_decision_gate_after_oauth_r "request_id": "trace-codex-image-local-123", "status_code": 200, "headers": { - "content-type": "text/event-stream" + "content-type": "application/json" }, "body": { "body_bytes_b64": base64::engine::general_purpose::STANDARD.encode( - concat!( - "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_img_123\",\"created_at\":1776839946}}\n\n", - "data: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"id\":\"ig_123\",\"type\":\"image_generation_call\",\"status\":\"generating\",\"output_format\":\"png\",\"quality\":\"medium\",\"size\":\"1024x1024\",\"revised_prompt\":\"中国历史视觉海报\",\"result\":\"aGVsbG8=\"}}\n\n", - "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_img_123\",\"object\":\"response\",\"model\":\"__CODEX_IMAGE_MODEL__\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":2440,\"output_tokens\":184,\"total_tokens\":2624},\"tool_usage\":{\"image_gen\":{\"input_tokens\":171,\"input_tokens_details\":{\"image_tokens\":0,\"text_tokens\":171},\"output_tokens\":1372,\"output_tokens_details\":{\"image_tokens\":1372,\"text_tokens\":0},\"total_tokens\":1543}}}}\n\n", - "data: [DONE]\n\n" - ) - .replace( - "__CODEX_IMAGE_MODEL__", - CODEX_OPENAI_IMAGE_INTERNAL_MODEL, - ) + r#"{"created":1776839946,"data":[{"b64_json":"aGVsbG8=","revised_prompt":"水墨视觉海报"}],"usage":{"input_tokens":171,"output_tokens":1372,"total_tokens":1543}}"# ) }, "telemetry": { @@ -1182,7 +1055,7 @@ async fn gateway_executes_codex_image_sync_via_local_decision_gate_after_oauth_r format!("Bearer {client_api_key}"), ) .header(TRACE_ID_HEADER, "trace-codex-image-local-123") - .body("{\"model\":\"gpt-image-2\",\"prompt\":\"生成一张中国历史视觉海报\",\"size\":\"1024x1024\",\"n\":1,\"response_format\":\"b64_json\"}") + .body("{\"model\":\"gpt-image-2\",\"prompt\":\"生成一张水墨视觉海报\",\"background\":\"auto\",\"quality\":\"auto\",\"size\":\"auto\",\"n\":1,\"response_format\":\"b64_json\"}") .send() .await .expect("request should succeed"); @@ -1191,10 +1064,7 @@ async fn gateway_executes_codex_image_sync_via_local_decision_gate_after_oauth_r let response_json: serde_json::Value = response.json().await.expect("body should parse"); assert_eq!(response_json["created"], 1776839946); assert_eq!(response_json["data"][0]["b64_json"], "aGVsbG8="); - assert_eq!( - response_json["data"][0]["revised_prompt"], - "中国历史视觉海报" - ); + assert_eq!(response_json["data"][0]["revised_prompt"], "水墨视觉海报"); assert_eq!(response_json["usage"]["input_tokens"], 171); assert_eq!(response_json["usage"]["output_tokens"], 1372); @@ -1229,35 +1099,34 @@ async fn gateway_executes_codex_image_sync_via_local_decision_gate_after_oauth_r ); assert_eq!( seen_execution_runtime_request.url, - "https://chatgpt.com/backend-api/codex/responses" - ); - assert_eq!( - seen_execution_runtime_request.model, - CODEX_OPENAI_IMAGE_INTERNAL_MODEL + "https://chatgpt.com/backend-api/codex/images/generations" ); assert_eq!( seen_execution_runtime_request.authorization, "Bearer refreshed-codex-image-access-token" ); assert_eq!( - seen_execution_runtime_request.x_client_request_id, - "trace-codex-image-local-123" + seen_execution_runtime_request.body, + json!({ + "prompt": "生成一张水墨视觉海报", + "background": "auto", + "model": "gpt-image-2", + "n": 1, + "quality": "auto", + "size": "auto" + }) ); assert_eq!( - seen_execution_runtime_request.prompt, - "生成一张中国历史视觉海报" + seen_execution_runtime_request.headers["user-agent"], + "codex_cli_rs/0.144.1" ); - assert!(seen_execution_runtime_request.content_is_string); - assert_eq!(seen_execution_runtime_request.tool_type, "image_generation"); - assert_eq!(seen_execution_runtime_request.tool_size, "1024x1024"); - assert_eq!(seen_execution_runtime_request.tool_quality, "high"); - assert_eq!(seen_execution_runtime_request.tool_background, "auto"); assert_eq!( - seen_execution_runtime_request.tool_choice_type, - "image_generation" + seen_execution_runtime_request.headers["originator"], + "codex_cli_rs" ); - assert!(!seen_execution_runtime_request.tool_has_n); - assert!(seen_execution_runtime_request.request_stream); + for header in ["x-client-request-id", "session-id", "thread-id"] { + assert!(seen_execution_runtime_request.headers.get(header).is_none()); + } assert!(!seen_execution_runtime_request.plan_stream); let persisted_transport_state = diff --git a/apps/aether-gateway/src/tests/ai_execute/sync/mod.rs b/apps/aether-gateway/src/tests/ai_execute/sync/mod.rs index f3142119b..d93e4e742 100644 --- a/apps/aether-gateway/src/tests/ai_execute/sync/mod.rs +++ b/apps/aether-gateway/src/tests/ai_execute/sync/mod.rs @@ -29,3 +29,4 @@ mod cli; mod gemini; mod image; mod pii_redaction_formats; +mod search; diff --git a/apps/aether-gateway/src/tests/ai_execute/sync/search.rs b/apps/aether-gateway/src/tests/ai_execute/sync/search.rs new file mode 100644 index 000000000..eb1de4c43 --- /dev/null +++ b/apps/aether-gateway/src/tests/ai_execute/sync/search.rs @@ -0,0 +1,597 @@ +use super::{ + any, build_router_with_state, build_state_with_execution_runtime_override, json, start_server, + to_bytes, Arc, Body, Json, Mutex, Request, Router, StatusCode, + EXECUTION_PATH_EXECUTION_RUNTIME_SYNC, EXECUTION_PATH_HEADER, TRACE_ID_HEADER, +}; +use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY}; +use aether_data::repository::auth::{ + InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot, +}; +use aether_data::repository::candidate_selection::InMemoryMinimalCandidateSelectionReadRepository; +use aether_data::repository::candidates::InMemoryRequestCandidateRepository; +use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository; +use aether_data_contracts::repository::candidate_selection::{ + StoredMinimalCandidateSelectionRow, StoredProviderModelMapping, +}; +use aether_data_contracts::repository::candidates::{ + RequestCandidateReadRepository, RequestCandidateStatus, +}; +use aether_data_contracts::repository::provider_catalog::{ + StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider, +}; +use sha2::{Digest, Sha256}; + +const SEARCH_SYNC_TEST_STACK_BYTES: usize = 16 * 1024 * 1024; + +fn run_search_sync_test(test_name: &'static str, make_future: F) +where + F: FnOnce() -> Fut + Send + 'static, + Fut: std::future::Future + 'static, +{ + let handle = std::thread::Builder::new() + .name(test_name.to_string()) + .stack_size(SEARCH_SYNC_TEST_STACK_BYTES) + .spawn(move || { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("test runtime should build"); + runtime.block_on(make_future()); + }) + .expect("search sync test thread should spawn"); + + if let Err(payload) = handle.join() { + std::panic::resume_unwind(payload); + } +} + +#[test] +fn gateway_executes_codex_search_with_responses_permission_and_search_contract() { + run_search_sync_test( + "gateway_executes_codex_search_with_responses_permission_and_search_contract", + gateway_executes_codex_search_with_responses_permission_and_search_contract_impl, + ); +} + +async fn gateway_executes_codex_search_with_responses_permission_and_search_contract_impl() { + fn hash_api_key(value: &str) -> String { + let mut hasher = Sha256::new(); + hasher.update(value.as_bytes()); + format!("{:x}", hasher.finalize()) + } + + fn auth_snapshot() -> StoredAuthApiKeySnapshot { + StoredAuthApiKeySnapshot::new( + "user-search-1".to_string(), + "alice".to_string(), + Some("alice@example.com".to_string()), + "user".to_string(), + "local".to_string(), + true, + false, + Some(json!(["openai", "codex"])), + Some(json!(["openai:responses"])), + None, + "api-key-search-1".to_string(), + Some("search-client".to_string()), + true, + false, + false, + Some(60), + Some(5), + Some(4_102_444_800_i64), + Some(json!(["openai", "codex"])), + Some(json!(["openai:responses"])), + None, + ) + .expect("auth snapshot should build") + } + + fn candidate_row() -> StoredMinimalCandidateSelectionRow { + StoredMinimalCandidateSelectionRow { + provider_id: "provider-codex-search-1".to_string(), + provider_name: "codex".to_string(), + provider_type: "codex".to_string(), + provider_priority: 10, + provider_is_active: true, + endpoint_id: "endpoint-codex-search-1".to_string(), + endpoint_api_format: "openai:search".to_string(), + endpoint_api_family: Some("openai".to_string()), + endpoint_kind: Some("search".to_string()), + endpoint_is_active: true, + key_id: "key-codex-search-1".to_string(), + key_name: "oauth".to_string(), + key_auth_type: "oauth".to_string(), + key_is_active: true, + key_api_formats: Some(vec!["openai:responses".to_string()]), + key_allowed_models: None, + key_capabilities: None, + key_internal_priority: 5, + key_global_priority_by_format: Some(json!({"openai:search": 1})), + model_id: "model-codex-search-1".to_string(), + global_model_id: "global-model-codex-search-1".to_string(), + global_model_name: "gpt-5.6-sol".to_string(), + global_model_mappings: None, + global_model_supports_streaming: Some(false), + model_provider_model_name: "gpt-5.6-sol".to_string(), + model_provider_model_mappings: Some(vec![StoredProviderModelMapping { + name: "gpt-5.6-sol".to_string(), + priority: 1, + api_formats: Some(vec!["openai:responses".to_string()]), + endpoint_ids: None, + }]), + model_supports_streaming: Some(false), + model_is_active: true, + model_is_available: true, + } + } + + fn provider() -> StoredProviderCatalogProvider { + StoredProviderCatalogProvider::new( + "provider-codex-search-1".to_string(), + "codex".to_string(), + Some("https://chatgpt.com".to_string()), + "codex".to_string(), + ) + .expect("provider should build") + .with_transport_fields( + true, + false, + false, + None, + Some(2), + None, + Some(900.0), + None, + None, + ) + } + + fn endpoint() -> StoredProviderCatalogEndpoint { + StoredProviderCatalogEndpoint::new( + "endpoint-codex-search-1".to_string(), + "provider-codex-search-1".to_string(), + "openai:search".to_string(), + Some("openai".to_string()), + Some("search".to_string()), + true, + ) + .expect("endpoint should build") + .with_transport_fields( + "https://chatgpt.com/backend-api/codex".to_string(), + None, + None, + Some(2), + None, + None, + None, + None, + ) + .expect("endpoint transport should build") + } + + fn key() -> StoredProviderCatalogKey { + let auth_config = encrypt_python_fernet_plaintext( + DEVELOPMENT_ENCRYPTION_KEY, + r#"{"provider_type":"codex","account_id":"account-search-1","is_fedramp":true}"#, + ) + .expect("auth config should encrypt"); + StoredProviderCatalogKey::new( + "key-codex-search-1".to_string(), + "provider-codex-search-1".to_string(), + "oauth".to_string(), + "oauth".to_string(), + None, + true, + ) + .expect("key should build") + .with_transport_fields( + Some(json!(["openai:responses"])), + encrypt_python_fernet_plaintext( + DEVELOPMENT_ENCRYPTION_KEY, + "codex-search-access-token", + ) + .expect("access token should encrypt"), + Some(auth_config), + None, + Some(json!({"openai:search": 1})), + None, + Some(4_102_444_800), + None, + None, + ) + .expect("key transport should build") + } + + let seen_plans = Arc::new(Mutex::new(Vec::::new())); + let seen_plans_clone = Arc::clone(&seen_plans); + let execution_runtime = Router::new().route( + "/v1/execute/sync", + any(move |request: Request| { + let seen_plans_inner = Arc::clone(&seen_plans_clone); + async move { + let (_, body) = request.into_parts(); + let bytes = to_bytes(body, usize::MAX).await.expect("body should read"); + let payload: serde_json::Value = + serde_json::from_slice(&bytes).expect("execution payload should parse"); + let request_id = payload["request_id"] + .as_str() + .unwrap_or_default() + .to_string(); + let provider_id = payload["provider_id"] + .as_str() + .unwrap_or_default() + .to_string(); + seen_plans_inner + .lock() + .expect("mutex should lock") + .push(payload); + let execution_result = if request_id == "trace-search-error-1" { + json!({ + "request_id": request_id, + "status_code": 400, + "headers": { + "content-type": "application/json", + "x-search-upstream": "rate-limited" + }, + "body": { + "json_body": { + "error": { + "type": "rate_limit_error", + "message": "Search capacity reached", + "param": null, + "code": "rate_limit_exceeded" + }, + "future_error_field": {"retryable": true} + } + }, + "telemetry": {"elapsed_ms": 17} + }) + } else if request_id == "trace-search-failover-1" + && provider_id == "provider-codex-search-1" + { + json!({ + "request_id": request_id, + "status_code": 500, + "headers": { + "content-type": "application/json", + "x-search-upstream": "primary" + }, + "body": { + "json_body": { + "error": { + "type": "server_error", + "message": "Search backend unavailable" + } + } + }, + "telemetry": {"elapsed_ms": 11} + }) + } else if request_id == "trace-search-failover-1" { + json!({ + "request_id": request_id, + "status_code": 200, + "headers": { + "content-type": "application/json", + "x-search-upstream": "backup" + }, + "body": { + "json_body": { + "output": "search fallback result" + } + }, + "telemetry": {"elapsed_ms": 23} + }) + } else { + json!({ + "request_id": request_id, + "status_code": 201, + "headers": { + "content-type": "application/json", + "x-search-upstream": "alpha" + }, + "body": { + "json_body": { + "output": "search result", + "encrypted_output": "encrypted-search-result", + "future_response_field": {"enabled": true} + } + }, + "telemetry": {"elapsed_ms": 42} + }) + }; + ( + StatusCode::OK, + [("x-search-source", "codex-alpha")], + Json(execution_result), + ) + } + }), + ); + + let client_api_key = "sk-client-search"; + let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( + Some(hash_api_key(client_api_key)), + auth_snapshot(), + )])); + let candidate_repository = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed({ + let primary = candidate_row(); + let mut backup = primary.clone(); + backup.provider_id = "provider-codex-search-2".to_string(); + backup.provider_name = "codex-backup".to_string(); + backup.provider_priority = 20; + backup.endpoint_id = "endpoint-codex-search-2".to_string(); + backup.key_id = "key-codex-search-2".to_string(); + backup.key_name = "oauth-backup".to_string(); + backup.key_internal_priority = 6; + backup.key_global_priority_by_format = Some(json!({"openai:search": 2})); + backup.model_id = "model-codex-search-2".to_string(); + vec![primary, backup] + })); + let catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( + { + let primary = provider(); + let mut backup = primary.clone(); + backup.id = "provider-codex-search-2".to_string(); + backup.name = "codex-backup".to_string(); + vec![primary, backup] + }, + { + let primary = endpoint(); + let mut backup = primary.clone(); + backup.id = "endpoint-codex-search-2".to_string(); + backup.provider_id = "provider-codex-search-2".to_string(); + vec![primary, backup] + }, + { + let primary = key(); + let mut backup = primary.clone(); + backup.id = "key-codex-search-2".to_string(); + backup.provider_id = "provider-codex-search-2".to_string(); + backup.name = "oauth-backup".to_string(); + backup.global_priority_by_format = Some(json!({"openai:search": 2})); + vec![primary, backup] + }, + )); + let request_candidates = Arc::new(InMemoryRequestCandidateRepository::default()); + let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await; + let data_state = + crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests( + auth_repository, + candidate_repository, + catalog_repository, + Arc::clone(&request_candidates), + DEVELOPMENT_ENCRYPTION_KEY, + ) + .with_system_config_values_for_tests([( + crate::system_features::ENABLE_MODEL_DIRECTIVES_CONFIG_KEY.to_string(), + json!(true), + )]); + let state = build_state_with_execution_runtime_override(execution_runtime_url) + .with_data_state_for_tests(data_state); + let gateway = build_router_with_state(state); + let (gateway_url, gateway_handle) = start_server(gateway).await; + + let response = reqwest::Client::new() + .post(format!("{gateway_url}/v1/alpha/search")) + .header(http::header::CONTENT_TYPE, "application/json") + .header( + http::header::AUTHORIZATION, + format!("Bearer {client_api_key}"), + ) + .header(TRACE_ID_HEADER, "trace-search-1") + .json(&json!({ + "id": "session-search-1", + "model": "gpt-5.6-sol-ultra-fast", + "reasoning": {"effort": "low", "summary": "auto"}, + "input": "find current OpenAI documentation", + "commands": { + "search_query": [{"q": "OpenAI Codex search"}], + "open": [{"ref_id": "turn0search0"}] + }, + "settings": { + "search_context_size": "high", + "allowed_callers": ["direct"] + }, + "max_output_tokens": 4096, + "store": false, + "stream": true, + "future_request_field": {"enabled": true} + })) + .send() + .await + .expect("search request should succeed"); + + if response.status() != StatusCode::CREATED { + let status = response.status(); + let body = response.text().await.expect("error response should read"); + panic!("Search request returned {status}: {body}"); + } + assert_eq!( + response + .headers() + .get("x-search-upstream") + .and_then(|value| value.to_str().ok()), + Some("alpha") + ); + assert_eq!( + response + .headers() + .get(EXECUTION_PATH_HEADER) + .and_then(|value| value.to_str().ok()), + Some(EXECUTION_PATH_EXECUTION_RUNTIME_SYNC) + ); + let response_json: serde_json::Value = response.json().await.expect("response should parse"); + assert_eq!(response_json["output"], "search result"); + assert_eq!(response_json["encrypted_output"], "encrypted-search-result"); + assert_eq!(response_json["future_response_field"]["enabled"], true); + + let plan = seen_plans + .lock() + .expect("mutex should lock") + .first() + .cloned() + .expect("execution plan should be captured"); + assert_eq!( + plan["url"], + "https://chatgpt.com/backend-api/codex/alpha/search" + ); + assert_eq!(plan["client_api_format"], "openai:search"); + assert_eq!(plan["provider_api_format"], "openai:search"); + assert_eq!(plan["stream"], false); + assert_eq!(plan["timeouts"]["total_ms"], 900_000); + assert_eq!( + plan["headers"]["authorization"], + "Bearer codex-search-access-token" + ); + assert_eq!(plan["headers"]["chatgpt-account-id"], "account-search-1"); + assert_eq!(plan["headers"]["x-openai-fedramp"], "true"); + assert_eq!(plan["headers"]["originator"], "codex_cli_rs"); + assert!(plan["headers"]["user-agent"] + .as_str() + .is_some_and(|value| value.starts_with("codex_cli_rs/"))); + assert!(plan["headers"].get("openai-beta").is_none()); + assert!(plan["headers"] + .get("x-openai-internal-codex-responses-lite") + .is_none()); + assert_ne!(plan["headers"]["accept"], "text/event-stream"); + + let body = &plan["body"]["json_body"]; + assert_eq!(body["id"], "session-search-1"); + assert_eq!(body["model"], "gpt-5.6-sol"); + assert_eq!(body["reasoning"]["effort"], "max"); + assert_eq!(body["reasoning"]["summary"], "auto"); + assert_eq!( + body["commands"]["search_query"][0]["q"], + "OpenAI Codex search" + ); + assert_eq!(body["commands"]["open"][0]["ref_id"], "turn0search0"); + assert_eq!(body["settings"]["search_context_size"], "high"); + assert_eq!(body["max_output_tokens"], 4096); + assert!(body.get("store").is_none()); + assert!(body.get("future_request_field").is_none()); + assert!(body.get("stream").is_none()); + assert!(body.get("service_tier").is_none()); + + let candidates = request_candidates + .list_by_request_id("trace-search-1") + .await + .expect("request candidates should read"); + assert_eq!(candidates.len(), 1); + assert_eq!(candidates[0].status, RequestCandidateStatus::Success); + + let expected_error_body = json!({ + "error": { + "type": "rate_limit_error", + "message": "Search capacity reached", + "param": null, + "code": "rate_limit_exceeded" + }, + "future_error_field": {"retryable": true} + }); + let error_response = reqwest::Client::new() + .post(format!("{gateway_url}/v1/alpha/search")) + .header(http::header::CONTENT_TYPE, "application/json") + .header( + http::header::AUTHORIZATION, + format!("Bearer {client_api_key}"), + ) + .header(TRACE_ID_HEADER, "trace-search-error-1") + .json(&json!({ + "id": "session-search-error-1", + "model": "gpt-5.6-sol", + "input": "find current OpenAI documentation", + "commands": {"search_query": [{"q": "OpenAI documentation"}]} + })) + .send() + .await + .expect("search error response should return"); + + assert_eq!(error_response.status(), StatusCode::BAD_REQUEST); + assert_eq!( + error_response + .headers() + .get("x-search-upstream") + .and_then(|value| value.to_str().ok()), + Some("rate-limited") + ); + assert_eq!( + error_response + .json::() + .await + .expect("error response should parse"), + expected_error_body + ); + let error_candidates = request_candidates + .list_by_request_id("trace-search-error-1") + .await + .expect("error request candidates should read"); + assert_eq!(error_candidates.len(), 1); + assert_eq!(error_candidates[0].status, RequestCandidateStatus::Failed); + + let failover_response = reqwest::Client::new() + .post(format!("{gateway_url}/v1/alpha/search")) + .header(http::header::CONTENT_TYPE, "application/json") + .header( + http::header::AUTHORIZATION, + format!("Bearer {client_api_key}"), + ) + .header(TRACE_ID_HEADER, "trace-search-failover-1") + .json(&json!({ + "id": "session-search-failover-1", + "model": "gpt-5.6-sol", + "input": "find current OpenAI documentation", + "commands": {"search_query": [{"q": "OpenAI documentation"}]} + })) + .send() + .await + .expect("search failover response should return"); + + assert_eq!(failover_response.status(), StatusCode::OK); + assert_eq!( + failover_response + .headers() + .get("x-search-upstream") + .and_then(|value| value.to_str().ok()), + Some("backup") + ); + assert_eq!( + failover_response + .json::() + .await + .expect("failover response should parse")["output"], + "search fallback result" + ); + let failover_plans = seen_plans + .lock() + .expect("mutex should lock") + .iter() + .filter(|plan| plan["request_id"] == "trace-search-failover-1") + .map(|plan| plan["provider_id"].clone()) + .collect::>(); + assert_eq!( + failover_plans, + vec![ + json!("provider-codex-search-1"), + json!("provider-codex-search-2") + ] + ); + let failover_candidates = request_candidates + .list_by_request_id("trace-search-failover-1") + .await + .expect("failover request candidates should read"); + assert_eq!(failover_candidates.len(), 2); + assert_eq!( + failover_candidates[0].status, + RequestCandidateStatus::Failed + ); + assert_eq!(failover_candidates[0].status_code, Some(500)); + assert_eq!( + failover_candidates[1].status, + RequestCandidateStatus::Success + ); + assert_eq!(failover_candidates[1].status_code, Some(200)); + + gateway_handle.abort(); + execution_runtime_handle.abort(); +} diff --git a/apps/aether-gateway/src/tests/concurrency.rs b/apps/aether-gateway/src/tests/concurrency.rs index 276f09729..f55a574d5 100644 --- a/apps/aether-gateway/src/tests/concurrency.rs +++ b/apps/aether-gateway/src/tests/concurrency.rs @@ -66,6 +66,7 @@ fn sample_decision() -> crate::control::GatewayControlDecision { auth_context: None, admin_principal: None, local_auth_rejection: None, + model_directive_policy: Default::default(), } } diff --git a/apps/aether-gateway/src/tests/control/admin/endpoints/routes.rs b/apps/aether-gateway/src/tests/control/admin/endpoints/routes.rs index 39e6e9efb..83e4c9517 100644 --- a/apps/aether-gateway/src/tests/control/admin/endpoints/routes.rs +++ b/apps/aether-gateway/src/tests/control/admin/endpoints/routes.rs @@ -559,6 +559,103 @@ async fn gateway_creates_admin_provider_endpoint_locally_with_trusted_admin_prin upstream_handle.abort(); } +#[tokio::test] +async fn gateway_rejects_streaming_policy_for_search_endpoint_before_catalog_write() { + let mut create_provider = sample_provider("provider-search-create", "search-create", 10); + create_provider.provider_type = "custom".to_string(); + let mut update_provider = sample_provider("provider-search-update", "search-update", 20); + update_provider.provider_type = "custom".to_string(); + let mut existing_endpoint = sample_endpoint( + "endpoint-search-update", + "provider-search-update", + "openai:search", + "https://search.example/v1", + ); + existing_endpoint.config = Some(json!({"marker": "kept"})); + let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( + vec![create_provider, update_provider], + vec![existing_endpoint], + vec![], + )); + + let gateway = build_router_with_state( + AppState::new() + .expect("gateway should build") + .with_data_state_for_tests( + GatewayDataState::with_provider_catalog_repository_for_tests( + provider_catalog_repository.clone(), + ), + ), + ); + let (gateway_url, gateway_handle) = start_server(gateway).await; + let client = reqwest::Client::new(); + + let create_response = client + .post(format!( + "{gateway_url}/api/admin/endpoints/providers/provider-search-create/endpoints" + )) + .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(&json!({ + "provider_id": "provider-search-create", + "api_format": "openai:search", + "base_url": "https://search.example/v1", + "config": {"upstream_stream_policy": "force_stream"} + })) + .send() + .await + .expect("request should succeed"); + assert_eq!(create_response.status(), StatusCode::BAD_REQUEST); + let create_payload: serde_json::Value = create_response + .json() + .await + .expect("json body should parse"); + assert_eq!( + create_payload["detail"], + "OpenAI Search 端点仅支持非流式上游请求" + ); + + let update_response = client + .put(format!( + "{gateway_url}/api/admin/endpoints/endpoint-search-update" + )) + .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(&json!({ + "config": {"upstreamStreamPolicy": true} + })) + .send() + .await + .expect("request should succeed"); + assert_eq!(update_response.status(), StatusCode::BAD_REQUEST); + let update_payload: serde_json::Value = update_response + .json() + .await + .expect("json body should parse"); + assert_eq!( + update_payload["detail"], + "OpenAI Search 端点仅支持非流式上游请求" + ); + + let created = provider_catalog_repository + .list_endpoints_by_provider_ids(&["provider-search-create".to_string()]) + .await + .expect("endpoints should read"); + assert!(created.is_empty()); + let unchanged = provider_catalog_repository + .list_endpoints_by_ids(&["endpoint-search-update".to_string()]) + .await + .expect("endpoint should read"); + assert_eq!(unchanged.len(), 1); + assert_eq!(unchanged[0].config, Some(json!({"marker": "kept"}))); + + gateway_handle.abort(); +} + #[tokio::test] async fn gateway_updates_admin_provider_endpoint_locally_with_trusted_admin_principal() { let upstream_hits = Arc::new(Mutex::new(0usize)); diff --git a/apps/aether-gateway/src/tests/control/admin/oauth.rs b/apps/aether-gateway/src/tests/control/admin/oauth.rs index cee2f189a..27afe8625 100644 --- a/apps/aether-gateway/src/tests/control/admin/oauth.rs +++ b/apps/aether-gateway/src/tests/control/admin/oauth.rs @@ -6086,6 +6086,9 @@ async fn gateway_manual_codex_oauth_refresh_reconciles_missing_fixed_endpoint_im .iter() .find(|endpoint| endpoint.api_format == "openai:responses") .expect("openai responses endpoint should be reconciled"); + assert!(endpoints + .iter() + .any(|endpoint| endpoint.api_format == "openai:search")); assert_eq!( responses_endpoint.base_url, "https://chatgpt.com/backend-api/codex" diff --git a/apps/aether-gateway/src/tests/control/admin/provider_query.rs b/apps/aether-gateway/src/tests/control/admin/provider_query.rs index 0085f12a4..51f435991 100644 --- a/apps/aether-gateway/src/tests/control/admin/provider_query.rs +++ b/apps/aether-gateway/src/tests/control/admin/provider_query.rs @@ -559,7 +559,7 @@ async fn gateway_handles_admin_provider_query_models_falls_back_to_codex_preset_ .expect("mutex should lock") += 1; assert_eq!( plan.url, - "https://chatgpt.com/backend-api/codex/models?client_version=0.128.0-alpha.1" + "https://chatgpt.com/backend-api/codex/models?client_version=0.144.1" ); Json(json!({ "request_id": "req-provider-query-codex-invalidated", @@ -625,6 +625,11 @@ async fn gateway_handles_admin_provider_query_models_falls_back_to_codex_preset_ let payload: serde_json::Value = response.json().await.expect("json body should parse"); assert_eq!(payload["success"], json!(true)); assert_eq!(payload["data"]["error"], serde_json::Value::Null); + let warning = payload["data"]["warning"] + .as_str() + .expect("Codex fallback warning should be present"); + assert!(warning.contains("Codex 动态模型目录不可用")); + assert!(warning.contains("invalidated")); let model_ids = payload["data"]["models"] .as_array() .expect("models should be an array") @@ -634,11 +639,14 @@ async fn gateway_handles_admin_provider_query_models_falls_back_to_codex_preset_ assert_eq!( model_ids, vec![ - "gpt-5.3-codex", - "gpt-5.3-codex-spark", + "codex-auto-review", + "gpt-5.2", "gpt-5.4", "gpt-5.4-mini", "gpt-5.5", + "gpt-5.6-luna", + "gpt-5.6-sol", + "gpt-5.6-terra", ] ); assert_eq!( @@ -2391,6 +2399,188 @@ async fn gateway_streams_codex_openai_responses_upstream_for_admin_pool_model_te execution_runtime_handle.abort(); } +#[test] +fn gateway_executes_codex_search_admin_pool_model_test_with_search_contract() { + run_provider_query_test( + "gateway_executes_codex_search_admin_pool_model_test_with_search_contract", + gateway_executes_codex_search_admin_pool_model_test_with_search_contract_impl, + ); +} + +async fn gateway_executes_codex_search_admin_pool_model_test_with_search_contract_impl() { + let execution_runtime = Router::new().route( + "/v1/execute/sync", + any(move |Json(plan): Json| async move { + assert_eq!(plan.provider_id, "provider-codex-search"); + assert_eq!(plan.endpoint_id, "endpoint-codex-search"); + assert_eq!(plan.key_id, "key-codex-search"); + assert_eq!(plan.client_api_format, "openai:search"); + assert_eq!(plan.provider_api_format, "openai:search"); + assert_eq!( + plan.url, + "https://chatgpt.com/backend-api/codex/alpha/search" + ); + assert_eq!(plan.model_name.as_deref(), Some("gpt-5.6-sol")); + assert!(!plan.stream, "Codex Search is a synchronous JSON protocol"); + assert_eq!( + plan.timeouts + .as_ref() + .and_then(|timeouts| timeouts.total_ms), + Some(900_000) + ); + assert_eq!( + plan.headers.get("authorization").map(String::as_str), + Some("Bearer codex-search-access-token") + ); + assert_eq!( + plan.headers.get("chatgpt-account-id").map(String::as_str), + Some("account-search-admin") + ); + assert_eq!( + plan.headers.get("x-openai-fedramp").map(String::as_str), + Some("true") + ); + assert_eq!( + plan.headers.get("originator").map(String::as_str), + Some("codex_cli_rs") + ); + assert!(plan + .headers + .get("user-agent") + .is_some_and(|value| value.starts_with("codex_cli_rs/"))); + assert!(!plan.headers.contains_key("openai-beta")); + assert!(!plan + .headers + .contains_key("x-openai-internal-codex-responses-lite")); + assert_ne!( + plan.headers.get("accept").map(String::as_str), + Some("text/event-stream") + ); + + let body = plan.body.json_body.as_ref().expect("search json body"); + assert_eq!( + body["id"], + json!("aether-model-test-provider-query-search-trace") + ); + assert_eq!(body["model"], json!("gpt-5.6-sol")); + assert_eq!(body["input"], json!("find current OpenAI documentation")); + assert_eq!( + body["commands"]["search_query"][0]["q"], + json!("OpenAI Codex Search") + ); + assert!(body.get("stream").is_none()); + assert!(body.get("store").is_none()); + assert!(body.get("service_tier").is_none()); + assert!(body.get("unknown_field").is_none()); + + Json(json!({ + "request_id": plan.request_id, + "candidate_id": plan.candidate_id, + "status_code": 200, + "headers": { + "content-type": "application/json" + }, + "body": { + "json_body": { + "output": "search result" + } + }, + "telemetry": { + "elapsed_ms": 21 + } + })) + }), + ); + + let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await; + let mut provider = sample_provider("provider-codex-search", "Codex Search", 10); + provider.provider_type = "codex".to_string(); + provider.request_timeout_secs = Some(900.0); + let mut endpoint = sample_endpoint( + "endpoint-codex-search", + "provider-codex-search", + "openai:search", + "https://chatgpt.com/backend-api/codex", + ); + endpoint.config = Some(json!({"upstream_stream_policy": "force_stream"})); + let mut key = sample_key( + "key-codex-search", + "provider-codex-search", + "openai:search", + "codex-search-access-token", + ); + key.auth_type = "oauth".to_string(); + key.encrypted_auth_config = Some( + aether_crypto::encrypt_python_fernet_plaintext( + DEVELOPMENT_ENCRYPTION_KEY, + r#"{"provider_type":"codex","account_id":"account-search-admin","is_fedramp":true}"#, + ) + .expect("auth config should encrypt"), + ); + let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( + vec![provider], + vec![endpoint], + vec![key], + )); + + let gateway = build_router_with_state( + build_state_with_execution_runtime_override(execution_runtime_url) + .with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests( + provider_catalog_repository, + DEVELOPMENT_ENCRYPTION_KEY.to_string(), + )), + ); + let (gateway_url, gateway_handle) = start_server(gateway).await; + + let response = reqwest::Client::new() + .post(format!( + "{gateway_url}/api/admin/provider-query/test-model-failover" + )) + .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(&json!({ + "provider_id": "provider-codex-search", + "mode": "pool", + "model": "gpt-5.6-sol", + "failover_models": ["gpt-5.6-sol"], + "api_format": "openai:search", + "endpoint_id": "endpoint-codex-search", + "request_id": "provider-query-search-trace", + "request_body": { + "model": "gpt-5.6-sol", + "input": "find current OpenAI documentation", + "commands": { + "search_query": [{"q": "OpenAI Codex Search"}] + }, + "max_output_tokens": 256, + "stream": true, + "store": false, + "service_tier": "priority", + "unknown_field": true + } + })) + .send() + .await + .expect("request should succeed"); + + assert_eq!(response.status(), StatusCode::OK); + let payload: serde_json::Value = response.json().await.expect("json body should parse"); + assert_eq!(payload["success"], json!(true), "payload={payload}"); + assert_eq!( + payload["attempts"][0]["request_body"]["id"], + json!("aether-model-test-provider-query-search-trace") + ); + assert_eq!( + payload["attempts"][0]["response_body"]["output"], + json!("search result") + ); + + gateway_handle.abort(); + execution_runtime_handle.abort(); +} + #[test] fn gateway_routes_grok_responses_admin_pool_model_test_through_grok_runtime() { run_provider_query_test( @@ -3837,13 +4027,12 @@ async fn gateway_handles_openai_responses_test_model_locally_impl() { .and_then(|value| value.as_str()), Some(prompt) ); - assert_eq!( - plan.body - .json_body - .as_ref() - .and_then(|body| body.get("instructions")), - Some(&json!("")) - ); + assert!(plan + .body + .json_body + .as_ref() + .and_then(|body| body.get("instructions")) + .is_none()); assert_eq!( plan.body .json_body @@ -3856,7 +4045,7 @@ async fn gateway_handles_openai_responses_test_model_locally_impl() { .json_body .as_ref() .and_then(|body| body.get("prompt_cache_key")) - .is_some()); + .is_none()); Json(json!({ "request_id": plan.request_id, "candidate_id": plan.candidate_id, @@ -3960,8 +4149,8 @@ async fn gateway_handles_openai_image_test_model_locally_impl() { assert_eq!(plan.client_api_format, "openai:image"); assert_eq!(plan.provider_api_format, "openai:image"); assert_eq!(plan.model_name.as_deref(), Some("gpt-image-1")); - assert_eq!(plan.url, "https://api.openai.example/v1/responses"); - assert!(plan.stream); + assert_eq!(plan.url, "https://api.openai.example/v1/images/generations"); + assert!(!plan.stream); assert_eq!( plan.headers.get("authorization").map(String::as_str), Some("Bearer sk-test-image") @@ -3971,38 +4160,42 @@ async fn gateway_handles_openai_image_test_model_locally_impl() { .json_body .as_ref() .and_then(|body| body.get("model")), - Some(&json!(crate::ai_serving::CODEX_OPENAI_IMAGE_INTERNAL_MODEL)) + Some(&json!("gpt-image-1")) ); assert_eq!( plan.body .json_body .as_ref() - .and_then(|body| body.get("input")) - .and_then(|input| input.as_array()) - .and_then(|items| items.first()) - .and_then(|item| item.get("content")) + .and_then(|body| body.get("prompt")) .and_then(|value| value.as_str()), Some("Draw a small blue square") ); + assert!(plan + .body + .json_body + .as_ref() + .is_some_and(|body| body.get("stream").is_none())); Json(json!({ "request_id": plan.request_id, "candidate_id": plan.candidate_id, "status_code": 200, "headers": { - "content-type": "text/event-stream" + "content-type": "application/json" }, "body": { - "body_bytes_b64": base64::engine::general_purpose::STANDARD.encode( - concat!( - "event: response.created\n", - "data: {\"type\":\"response.created\",\"response\":{\"created_at\":1776839946}}\n\n", - "event: response.output_item.done\n", - "data: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"type\":\"image_generation_call\",\"output_format\":\"png\",\"revised_prompt\":\"revised prompt\",\"result\":\"aGVsbG8=\"}}\n\n", - "event: response.completed\n", - "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_img_123\",\"model\":\"gpt-image-1\",\"status\":\"completed\",\"tool_usage\":{\"image_gen\":{\"input_tokens\":171,\"output_tokens\":1372,\"total_tokens\":1543}}}}\n\n" - ) - .as_bytes() - ) + "json_body": { + "created": 1776839946, + "model": "gpt-image-1", + "data": [{ + "b64_json": "aGVsbG8=", + "revised_prompt": "revised prompt" + }], + "usage": { + "input_tokens": 171, + "output_tokens": 1372, + "total_tokens": 1543 + } + } }, "telemetry": { "elapsed_ms": 19 diff --git a/apps/aether-gateway/src/tests/control/admin/providers.rs b/apps/aether-gateway/src/tests/control/admin/providers.rs index b3d406287..39dee06c4 100644 --- a/apps/aether-gateway/src/tests/control/admin/providers.rs +++ b/apps/aether-gateway/src/tests/control/admin/providers.rs @@ -832,7 +832,7 @@ async fn gateway_updates_admin_provider_locally_with_trusted_admin_principal() { "is_active": false, "concurrent_limit": 8, "max_retries": 6, - "request_timeout": 55.0, + "request_timeout": aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_SECS, "stream_first_byte_timeout": 11.0, "enable_format_conversion": false, "config": { @@ -860,7 +860,10 @@ async fn gateway_updates_admin_provider_locally_with_trusted_admin_principal() { assert_eq!(payload["enable_format_conversion"], false); assert_eq!(payload["is_active"], false); assert_eq!(payload["max_retries"], 6); - assert_eq!(payload["request_timeout"], 55.0); + assert_eq!( + payload["request_timeout"].as_f64(), + Some(aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_SECS as f64) + ); assert_eq!(payload["stream_first_byte_timeout"], 11.0); assert_eq!(payload["proxy"], json!({"url": "https://proxy.example"})); assert_eq!(payload["claude_code_advanced"], json!({"pool_size": 2})); @@ -870,6 +873,21 @@ async fn gateway_updates_admin_provider_locally_with_trusted_admin_principal() { assert_eq!(payload["ops_configured"], true); assert_eq!(payload["ops_architecture_id"], "cubence"); + let invalid_timeout_response = reqwest::Client::new() + .patch(format!("{gateway_url}/api/admin/providers/provider-openai")) + .header(crate::constants::GATEWAY_HEADER, "rust-phase3b") + .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123") + .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin") + .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123") + .json(&json!({ + "request_timeout": + aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_SECS + 1 + })) + .send() + .await + .expect("request should succeed"); + assert_eq!(invalid_timeout_response.status(), StatusCode::BAD_REQUEST); + let disable_response = reqwest::Client::new() .patch(format!("{gateway_url}/api/admin/providers/provider-openai")) .header(crate::constants::GATEWAY_HEADER, "rust-phase3b") @@ -924,6 +942,10 @@ async fn gateway_updates_admin_provider_locally_with_trusted_admin_principal() { .iter() .find(|provider| provider.id == "provider-openai") .expect("provider should exist"); + assert_eq!( + updated_provider.request_timeout_secs, + Some(aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_SECS as f64) + ); assert_eq!( updated_provider .config @@ -1000,6 +1022,7 @@ async fn gateway_creates_admin_provider_locally_with_trusted_admin_principal() { "website": "codex.example", "keep_priority_on_conversion": true, "max_retries": 7, + "request_timeout": aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_SECS, "config": {"chat_pii_redaction": {"enabled": true}}, "pool_advanced": {}, "failover_rules": {"strategy": "ordered"}, @@ -1034,6 +1057,10 @@ async fn gateway_creates_admin_provider_locally_with_trusted_admin_principal() { assert_eq!(created.website.as_deref(), Some("https://codex.example")); assert!(created.enable_format_conversion); assert_eq!(created.max_retries, Some(7)); + assert_eq!( + created.request_timeout_secs, + Some(aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_SECS as f64) + ); assert_eq!(created.keep_priority_on_conversion, true); assert_eq!( created @@ -1080,7 +1107,7 @@ async fn gateway_creates_admin_provider_locally_with_trusted_admin_principal() { .list_endpoints_by_provider_ids(std::slice::from_ref(&created.id)) .await .expect("endpoints should list"); - assert_eq!(endpoints.len(), 3); + assert_eq!(endpoints.len(), 4); let responses_endpoint = endpoints .iter() .find(|endpoint| endpoint.api_format == "openai:responses") @@ -1089,6 +1116,10 @@ async fn gateway_creates_admin_provider_locally_with_trusted_admin_principal() { .iter() .find(|endpoint| endpoint.api_format == "openai:responses:compact") .expect("compact endpoint should exist"); + let search_endpoint = endpoints + .iter() + .find(|endpoint| endpoint.api_format == "openai:search") + .expect("search endpoint should exist"); let image_endpoint = endpoints .iter() .find(|endpoint| endpoint.api_format == "openai:image") @@ -1101,12 +1132,17 @@ async fn gateway_creates_admin_provider_locally_with_trusted_admin_principal() { compact_endpoint.base_url, "https://chatgpt.com/backend-api/codex" ); + assert_eq!( + search_endpoint.base_url, + "https://chatgpt.com/backend-api/codex" + ); assert_eq!( image_endpoint.base_url, "https://chatgpt.com/backend-api/codex" ); assert_eq!(responses_endpoint.max_retries, Some(7)); assert_eq!(compact_endpoint.max_retries, Some(7)); + assert_eq!(search_endpoint.max_retries, Some(7)); assert_eq!(image_endpoint.max_retries, Some(7)); assert_eq!( responses_endpoint @@ -1116,16 +1152,25 @@ async fn gateway_creates_admin_provider_locally_with_trusted_admin_principal() { .and_then(serde_json::Value::as_str), Some("force_stream") ); + assert_eq!( + search_endpoint + .config + .as_ref() + .and_then(|value| value.get("upstream_stream_policy")) + .and_then(serde_json::Value::as_str), + None + ); assert_eq!( image_endpoint .config .as_ref() .and_then(|value| value.get("upstream_stream_policy")) .and_then(serde_json::Value::as_str), - Some("force_stream") + None ); assert!(responses_endpoint.body_rules.is_none()); assert!(compact_endpoint.body_rules.is_none()); + assert!(search_endpoint.body_rules.is_none()); assert!(image_endpoint.body_rules.is_none()); assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0); @@ -1216,7 +1261,7 @@ async fn gateway_updates_fixed_provider_and_reconciles_template_managed_endpoint .list_endpoints_by_provider_ids(&["provider-codex".to_string()]) .await .expect("endpoints should list"); - assert_eq!(endpoints.len(), 3); + assert_eq!(endpoints.len(), 4); let responses_endpoint = endpoints .iter() .find(|endpoint| endpoint.api_format == "openai:responses") @@ -1225,6 +1270,10 @@ async fn gateway_updates_fixed_provider_and_reconciles_template_managed_endpoint .iter() .find(|endpoint| endpoint.api_format == "openai:responses:compact") .expect("compact endpoint should exist"); + let search_endpoint = endpoints + .iter() + .find(|endpoint| endpoint.api_format == "openai:search") + .expect("search endpoint should exist"); let image_endpoint = endpoints .iter() .find(|endpoint| endpoint.api_format == "openai:image") @@ -1232,6 +1281,7 @@ async fn gateway_updates_fixed_provider_and_reconciles_template_managed_endpoint assert_eq!(responses_endpoint.max_retries, Some(9)); assert_eq!(compact_endpoint.max_retries, Some(9)); + assert_eq!(search_endpoint.max_retries, Some(9)); assert_eq!(image_endpoint.max_retries, Some(9)); assert_eq!( responses_endpoint @@ -1242,20 +1292,37 @@ async fn gateway_updates_fixed_provider_and_reconciles_template_managed_endpoint .and_then(serde_json::Value::as_bool), Some(true) ); + assert_eq!( + search_endpoint + .config + .as_ref() + .and_then(|value| value.get("_aether_fixed_provider_template")) + .and_then(|value| value.get("managed")) + .and_then(serde_json::Value::as_bool), + Some(true) + ); + assert_eq!( + search_endpoint + .config + .as_ref() + .and_then(|value| value.get("upstream_stream_policy")) + .and_then(serde_json::Value::as_str), + None + ); assert_eq!( image_endpoint .config .as_ref() .and_then(|value| value.get("upstream_stream_policy")) .and_then(serde_json::Value::as_str), - Some("force_stream") + None ); let keys = provider_catalog_repository .list_keys_by_provider_ids(&["provider-codex".to_string()]) .await .expect("keys should list"); assert_eq!(keys.len(), 1); - assert!(keys[0].api_formats.is_none()); + assert_eq!(keys[0].api_formats, Some(json!(["openai:responses"]))); assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0); gateway_handle.abort(); diff --git a/apps/aether-gateway/src/tests/control/admin/stats.rs b/apps/aether-gateway/src/tests/control/admin/stats.rs index c9f439020..15b67483b 100644 --- a/apps/aether-gateway/src/tests/control/admin/stats.rs +++ b/apps/aether-gateway/src/tests/control/admin/stats.rs @@ -1379,7 +1379,7 @@ async fn gateway_handles_admin_stats_leaderboard_models_locally_with_trusted_adm assert_eq!(payload["metric"], "tokens"); assert_eq!(payload["items"][0]["rank"], 1); assert_eq!(payload["items"][0]["id"], "gpt-5"); - assert_eq!(payload["items"][0]["value"], 160); + assert_eq!(payload["items"][0]["value"], 150); assert_eq!(payload["items"][1]["id"], "claude-3-5-sonnet"); assert_eq!(payload["items"][1]["value"], 100); assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0); diff --git a/apps/aether-gateway/src/tests/control/admin/usage.rs b/apps/aether-gateway/src/tests/control/admin/usage.rs index 5bbf7cb6b..67bd9802e 100644 --- a/apps/aether-gateway/src/tests/control/admin/usage.rs +++ b/apps/aether-gateway/src/tests/control/admin/usage.rs @@ -641,7 +641,7 @@ async fn gateway_handles_admin_usage_aggregation_stats_locally_with_trusted_admi assert_eq!(items[0]["model"], "gpt-5"); assert_eq!(items[0]["request_count"], 2); assert_eq!(items[0]["output_tokens"], 40); - assert_eq!(items[0]["effective_input_tokens"], 150); + assert_eq!(items[0]["effective_input_tokens"], 120); assert_eq!(items[0]["total_input_context"], 160); assert_eq!(items[0]["cache_creation_tokens"], 30); assert_eq!(items[0]["cache_creation_ephemeral_5m_tokens"], 12); @@ -1026,7 +1026,7 @@ async fn gateway_handles_admin_usage_active_locally_with_trusted_admin_principal let payload: serde_json::Value = response.json().await.expect("json body should parse"); assert_eq!(payload["requests"].as_array().expect("array").len(), 1); assert_eq!(payload["requests"][0]["id"], "usage-pending"); - assert_eq!(payload["requests"][0]["effective_input_tokens"], 5); + assert_eq!(payload["requests"][0]["effective_input_tokens"], 0); assert_eq!(payload["requests"][0]["provider"], "OpenAI"); assert_eq!(payload["requests"][0]["api_key_name"], "fresh-primary"); assert_eq!(payload["requests"][0]["has_fallback"], true); @@ -1326,7 +1326,7 @@ async fn gateway_handles_admin_usage_records_locally_with_trusted_admin_principa payload["records"][0]["provider_key_name"], "upstream-primary" ); - assert_eq!(payload["records"][0]["effective_input_tokens"], 35); + assert_eq!(payload["records"][0]["effective_input_tokens"], 20); assert_eq!(payload["records"][0]["first_byte_time_ms"], 120); assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0); @@ -2053,8 +2053,8 @@ async fn gateway_handles_admin_usage_detail_locally_with_trusted_admin_principal assert_eq!(payload["api_key"]["name"], "primary"); assert_eq!(payload["provider"], "OpenAI"); assert_eq!(payload["model"], "gpt-5"); - assert_eq!(payload["effective_input_tokens"], 115); - assert_eq!(payload["total_tokens"], 165); + assert_eq!(payload["effective_input_tokens"], 100); + assert_eq!(payload["total_tokens"], 150); assert_eq!(payload["cache_creation_cost"], 0.0); assert_eq!(payload["cache_read_cost"], 0.0); assert_eq!( diff --git a/apps/aether-gateway/src/tests/control/internal.rs b/apps/aether-gateway/src/tests/control/internal.rs index 9a74ec153..6b449cd56 100644 --- a/apps/aether-gateway/src/tests/control/internal.rs +++ b/apps/aether-gateway/src/tests/control/internal.rs @@ -1,6 +1,7 @@ use std::io; use std::sync::{Arc, Mutex}; +use aether_contracts::tunnel::RequestMeta; use aether_data::repository::proxy_nodes::ProxyNodeReadRepository; use axum::body::Body; use axum::routing::{any, post}; @@ -10,12 +11,45 @@ use futures_util::stream; use http::header::HeaderValue; use http::StatusCode; use serde_json::json; +use std::collections::HashMap; +use std::time::Duration; use super::{ build_router_with_state, sample_proxy_node, start_server, AppState, GatewayDataState, InMemoryProxyNodeRepository, TRACE_ID_HEADER, }; +fn relay_request_meta( + stream: bool, + request_timeout_ms: Option, + stream_first_byte_timeout_ms: Option, +) -> RequestMeta { + RequestMeta { + provider_id: Some("provider-1".to_string()), + endpoint_id: Some("endpoint-1".to_string()), + key_id: Some("key-1".to_string()), + method: "POST".to_string(), + url: "https://example.com/responses".to_string(), + headers: HashMap::new(), + stream, + request_timeout_ms, + stream_first_byte_timeout_ms, + timeout: 60, + follow_redirects: None, + http1_only: false, + transport_profile: None, + } +} + +fn relay_envelope(meta: &RequestMeta, body: &[u8]) -> Vec { + let encoded_meta = serde_json::to_vec(meta).expect("metadata should encode"); + let mut envelope = Vec::with_capacity(4 + encoded_meta.len() + body.len()); + envelope.extend_from_slice(&(encoded_meta.len() as u32).to_be_bytes()); + envelope.extend_from_slice(&encoded_meta); + envelope.extend_from_slice(body); + envelope +} + #[tokio::test] async fn gateway_handles_internal_tunnel_heartbeat_locally_with_loopback() { let upstream_hits = Arc::new(Mutex::new(0usize)); @@ -292,10 +326,13 @@ async fn gateway_forwards_tunnel_relay_to_attachment_owner() { ); let (gateway_url, gateway_handle) = start_server(gateway).await; + let envelope = relay_envelope(&relay_request_meta(false, None, None), b"relay-envelope"); + let response = reqwest::Client::new() .post(format!("{gateway_url}/api/internal/tunnel/relay/node-123")) .header(TRACE_ID_HEADER, "trace-owner-forward") - .body("relay-envelope") + .header(http::header::CONTENT_TYPE, "application/octet-stream") + .body(envelope.clone()) .send() .await .expect("request should succeed"); @@ -309,8 +346,8 @@ async fn gateway_forwards_tunnel_relay_to_attachment_owner() { Some("trace-owner-forward") ); assert_eq!( - response.text().await.expect("body should read"), - "relay-envelope" + response.bytes().await.expect("body should read"), + Bytes::from(envelope) ); assert_eq!(*owner_hits.lock().expect("mutex should lock"), 1); @@ -318,6 +355,71 @@ async fn gateway_forwards_tunnel_relay_to_attachment_owner() { owner_handle.abort(); } +#[tokio::test] +async fn gateway_owner_relay_uses_non_stream_timeout_from_envelope() { + let owner = Router::new().route( + "/api/internal/tunnel/relay/node-123", + post(|body: Body| async move { + let body = axum::body::to_bytes(body, usize::MAX) + .await + .expect("body should read"); + tokio::time::sleep(Duration::from_millis(40)).await; + (StatusCode::OK, Body::from(body)) + }), + ); + + let (owner_url, owner_handle) = start_server(owner).await; + let data_state = GatewayDataState::disabled().with_system_config_values_for_tests(vec![( + "tunnel.attachments.node-123".to_string(), + json!({ + "gateway_instance_id": "gateway-b", + "relay_base_url": owner_url, + "conn_count": 1, + "observed_at_unix_secs": 4_102_444_800u64, + }), + )]); + let mut state = AppState::new() + .expect("gateway should build") + .with_data_state_for_tests(data_state) + .with_tunnel_identity_for_tests("gateway-a", Some("http://gateway-a.internal")); + let short_timeout_client = reqwest::Client::builder() + .timeout(Duration::from_millis(10)) + .build() + .expect("test client should build"); + state.client = short_timeout_client.clone(); + state.owner_forward_client = short_timeout_client; + let gateway = build_router_with_state(state); + let (gateway_url, gateway_handle) = start_server(gateway).await; + + let meta = relay_request_meta(false, Some(100), None); + let envelope = relay_envelope(&meta, b"relay-body"); + let encoded_meta = serde_json::to_vec(&meta).expect("metadata should encode"); + let split_at = 4 + encoded_meta.len() / 2; + let request_body = reqwest::Body::wrap_stream(stream::iter(vec![ + Ok::(Bytes::copy_from_slice(&envelope[..split_at])), + Ok::(Bytes::copy_from_slice(&envelope[split_at..])), + ])); + + let response = reqwest::Client::builder() + .timeout(Duration::from_secs(1)) + .build() + .expect("request client should build") + .post(format!("{gateway_url}/api/internal/tunnel/relay/node-123")) + .body(request_body) + .send() + .await + .expect("request should succeed"); + + assert_eq!(response.status(), StatusCode::OK); + assert_eq!( + response.bytes().await.expect("response body should read"), + Bytes::from(envelope) + ); + + gateway_handle.abort(); + owner_handle.abort(); +} + #[tokio::test] async fn gateway_streams_tunnel_relay_body_to_attachment_owner() { let owner_hits = Arc::new(Mutex::new(0usize)); @@ -331,8 +433,12 @@ async fn gateway_streams_tunnel_relay_body_to_attachment_owner() { let body = axum::body::to_bytes(body, usize::MAX) .await .expect("body should read"); - assert_eq!(body, Bytes::from_static(b"relay-stream-envelope")); - (StatusCode::OK, Body::from("stream-ok")) + let response_body = Body::from_stream(async_stream::stream! { + yield Ok::<_, io::Error>(Bytes::from_static(b"stream-")); + tokio::time::sleep(Duration::from_millis(40)).await; + yield Ok::<_, io::Error>(Bytes::from_static(b"ok")); + }); + (StatusCode::OK, response_body) } }), ); @@ -347,19 +453,23 @@ async fn gateway_streams_tunnel_relay_body_to_attachment_owner() { "observed_at_unix_secs": 4_102_444_800u64, }), )]); - let gateway = build_router_with_state( - AppState::new() - .expect("gateway should build") - .with_data_state_for_tests(data_state) - .with_tunnel_identity_for_tests("gateway-a", Some("http://gateway-a.internal")), - ); + let mut state = AppState::new() + .expect("gateway should build") + .with_data_state_for_tests(data_state) + .with_tunnel_identity_for_tests("gateway-a", Some("http://gateway-a.internal")); + state.client = reqwest::Client::builder() + .timeout(Duration::from_millis(10)) + .build() + .expect("short shared client should build"); + let gateway = build_router_with_state(state); let (gateway_url, gateway_handle) = start_server(gateway).await; - let request_body = reqwest::Body::wrap_stream(stream::iter(vec![ - Ok::(Bytes::from_static(b"relay-")), - Ok::(Bytes::from_static(b"stream-")), - Ok::(Bytes::from_static(b"envelope")), - ])); + let meta = relay_request_meta(true, Some(900_000), Some(100)); + let envelope = relay_envelope(&meta, b"relay-stream-envelope"); + let expected_envelope = Bytes::copy_from_slice(&envelope); + let request_body = reqwest::Body::wrap_stream(stream::iter(vec![Ok::( + expected_envelope.clone(), + )])); let response = reqwest::Client::new() .post(format!("{gateway_url}/api/internal/tunnel/relay/node-123")) .body(request_body) @@ -369,8 +479,8 @@ async fn gateway_streams_tunnel_relay_body_to_attachment_owner() { assert_eq!(response.status(), StatusCode::OK); assert_eq!( - response.text().await.expect("body should read"), - "stream-ok" + response.bytes().await.expect("body should read"), + Bytes::from_static(b"stream-ok") ); assert_eq!(*owner_hits.lock().expect("mutex should lock"), 1); diff --git a/apps/aether-gateway/src/tests/frontdoor.rs b/apps/aether-gateway/src/tests/frontdoor.rs index 050730c90..11a59093e 100644 --- a/apps/aether-gateway/src/tests/frontdoor.rs +++ b/apps/aether-gateway/src/tests/frontdoor.rs @@ -35,6 +35,27 @@ use aether_data_contracts::repository::video_tasks::{ use base64::Engine as _; use sha2::{Digest, Sha256}; +fn run_frontdoor_async_test(name: &'static str, future: F) +where + F: std::future::Future + Send + 'static, +{ + let handle = std::thread::Builder::new() + .name(name.to_string()) + .stack_size(16 * 1024 * 1024) + .spawn(move || { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("frontdoor test runtime should build") + .block_on(future); + }) + .expect("large-stack frontdoor test thread should spawn"); + + if let Err(payload) = handle.join() { + std::panic::resume_unwind(payload); + } +} + fn hash_api_key(value: &str) -> String { let mut hasher = Sha256::new(); hasher.update(value.as_bytes()); diff --git a/apps/aether-gateway/src/tests/frontdoor/ai.rs b/apps/aether-gateway/src/tests/frontdoor/ai.rs index 542727bf4..7e4951a61 100644 --- a/apps/aether-gateway/src/tests/frontdoor/ai.rs +++ b/apps/aether-gateway/src/tests/frontdoor/ai.rs @@ -1,8 +1,8 @@ use super::{ hash_api_key, sample_models_candidate_row, unrestricted_models_snapshot, InMemoryAuthApiKeySnapshotRepository, InMemoryMinimalCandidateSelectionReadRepository, - InMemoryVideoTaskRepository, UpsertVideoTask, VideoTaskLookupKey, VideoTaskReadRepository, - VideoTaskStatus, VideoTaskWriteRepository, DEVELOPMENT_ENCRYPTION_KEY, + InMemoryVideoTaskRepository, StoredAuthApiKeySnapshot, UpsertVideoTask, VideoTaskLookupKey, + VideoTaskReadRepository, VideoTaskStatus, VideoTaskWriteRepository, DEVELOPMENT_ENCRYPTION_KEY, }; use crate::image_capabilities::openai_image_gateway_max_generation_count; use crate::tests::{ @@ -26,6 +26,94 @@ use std::collections::HashMap; use std::future::pending; use std::sync::atomic::{AtomicBool, Ordering}; +fn codex_models_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot { + StoredAuthApiKeySnapshot::new( + user_id.to_string(), + "alice".to_string(), + Some("alice@example.com".to_string()), + "user".to_string(), + "local".to_string(), + true, + false, + Some(json!(["codex"])), + Some(json!(["openai:responses"])), + Some(json!(["frontier-sol", "broken-luna"])), + api_key_id.to_string(), + Some("codex-models".to_string()), + true, + false, + false, + Some(10), + Some(5), + Some(4_102_444_800), + Some(json!(["codex"])), + Some(json!(["openai:responses"])), + Some(json!(["frontier-sol", "broken-luna"])), + ) + .expect("Codex models auth snapshot should build") +} + +fn sample_codex_models_candidate_row( + provider_id: &str, + global_model_name: &str, + source_model_name: &str, +) -> StoredMinimalCandidateSelectionRow { + let mut row = sample_models_candidate_row( + provider_id, + "codex", + "openai:responses", + global_model_name, + 10, + ); + row.provider_type = "codex".to_string(); + row.key_auth_type = "oauth".to_string(); + row.model_provider_model_name = source_model_name.to_string(); + row.model_provider_model_mappings = Some(vec![ + aether_data_contracts::repository::candidate_selection::StoredProviderModelMapping { + name: source_model_name.to_string(), + priority: 1, + api_formats: Some(vec!["openai:responses".to_string()]), + endpoint_ids: None, + }, + ]); + row +} + +fn complete_codex_model_card(source_model_name: &str) -> serde_json::Value { + json!({ + "id": source_model_name, + "api_formats": ["openai:responses"], + "slug": source_model_name, + "display_name": "GPT-5.6-Sol", + "description": "Frontier coding model", + "default_reasoning_level": "low", + "supported_reasoning_levels": [ + {"effort": "low", "description": "Low"}, + {"effort": "medium", "description": "Medium"}, + {"effort": "high", "description": "High"}, + {"effort": "xhigh", "description": "XHigh"}, + {"effort": "max", "description": "Max"}, + {"effort": "ultra", "description": "Ultra"} + ], + "shell_type": "shell_command", + "visibility": "list", + "supported_in_api": true, + "priority": 1, + "availability_nux": null, + "upgrade": null, + "base_instructions": "Use the current Codex instructions.", + "model_messages": null, + "support_verbosity": true, + "default_verbosity": "low", + "apply_patch_tool_type": "freeform", + "truncation_policy": {"mode": "tokens", "limit": 10000}, + "supports_parallel_tool_calls": true, + "experimental_supported_tools": [], + "minimal_client_version": "0.144.0", + "future_capability": {"enabled": true} + }) +} + fn gemini_operation_status_label(status: VideoTaskStatus) -> &'static str { match status { VideoTaskStatus::Pending => "Pending", @@ -360,6 +448,121 @@ async fn gateway_handles_public_openai_models_without_hitting_fallback_probe() { fallback_probe_handle.abort(); } +#[tokio::test] +async fn gateway_serves_codex_model_cards_for_versioned_models_requests() { + let codex_row = + sample_codex_models_candidate_row("provider-codex-models", "frontier-sol", "gpt-5.6-sol"); + let incomplete_codex_row = sample_codex_models_candidate_row( + "provider-codex-incomplete", + "broken-luna", + "gpt-5.6-luna", + ); + let candidate_repository = + Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![ + codex_row.clone(), + incomplete_codex_row.clone(), + sample_models_candidate_row( + "provider-openai-responses", + "openai", + "openai:responses", + "custom-responses-model", + 20, + ), + ])); + let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![ + ( + Some(hash_api_key("sk-codex-models")), + codex_models_snapshot("key-codex-models", "user-codex-models"), + ), + ( + Some(hash_api_key("sk-standard-models")), + unrestricted_models_snapshot("key-standard-models", "user-standard-models"), + ), + ])); + let state = AppState::new() + .expect("gateway should build") + .with_data_state_for_tests( + crate::data::GatewayDataState::with_minimal_candidate_selection_and_auth_for_tests( + candidate_repository, + auth_repository, + ), + ); + state + .runtime_kv_setex( + &format!( + "upstream_models:{}:{}", + codex_row.provider_id, codex_row.key_id + ), + &serde_json::to_string(&vec![complete_codex_model_card("gpt-5.6-sol")]) + .expect("model cache should serialize"), + 60, + ) + .await + .expect("model cache should seed"); + state + .runtime_kv_setex( + &format!( + "upstream_models:{}:{}", + incomplete_codex_row.provider_id, incomplete_codex_row.key_id + ), + &serde_json::to_string(&vec![json!({ + "id": "gpt-5.6-luna", + "slug": "gpt-5.6-luna", + "display_name": "GPT-5.6-Luna" + })]) + .expect("incomplete model cache should serialize"), + 60, + ) + .await + .expect("incomplete model cache should seed"); + + let gateway = build_router_with_state(state); + let (gateway_url, gateway_handle) = start_server(gateway).await; + let client = reqwest::Client::new(); + + let codex_response = client + .get(format!("{gateway_url}/v1/models?client_version=0.144.1")) + .header("authorization", "Bearer sk-codex-models") + .send() + .await + .expect("Codex models request should succeed"); + assert_eq!(codex_response.status(), StatusCode::OK); + let codex_payload: serde_json::Value = codex_response + .json() + .await + .expect("Codex models body should parse"); + assert_eq!(codex_payload["models"].as_array().map(Vec::len), Some(1)); + assert_eq!(codex_payload["models"][0]["slug"], "frontier-sol"); + assert_eq!( + codex_payload["models"][0]["supported_reasoning_levels"][5]["effort"], + "ultra" + ); + assert_eq!( + codex_payload["models"][0]["future_capability"], + json!({"enabled": true}) + ); + assert!(codex_payload["models"][0].get("id").is_none()); + assert!(codex_payload["models"][0].get("api_formats").is_none()); + assert!(codex_payload.get("object").is_none()); + + let standard_response = client + .get(format!("{gateway_url}/v1/models")) + .header("authorization", "Bearer sk-standard-models") + .send() + .await + .expect("standard models request should succeed"); + assert_eq!(standard_response.status(), StatusCode::OK); + let standard_payload: serde_json::Value = standard_response + .json() + .await + .expect("standard models body should parse"); + assert_eq!(standard_payload["object"], "list"); + assert!(standard_payload["data"].is_array()); + assert!(standard_payload.get("models").is_none()); + + gateway_handle.abort(); +} + #[tokio::test] async fn gateway_openai_models_list_drops_disabled_global_model_after_cache_invalidation() { let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( @@ -1212,7 +1415,7 @@ async fn gateway_does_not_locally_reject_image_model_name_on_chat_completions() } #[tokio::test] -async fn gateway_rejects_image_request_with_n_greater_than_four_without_hitting_fallback_probe() { +async fn gateway_rejects_image_request_above_gateway_limit_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( @@ -1247,7 +1450,7 @@ async fn gateway_rejects_image_request_with_n_greater_than_four_without_hitting_ serde_json::to_vec(&json!({ "model": "grok-imagine-image-lite", "prompt": "draw", - "n": 5, + "n": openai_image_gateway_max_generation_count() + 1, "response_format": "b64_json" })) .expect("request body should encode"), diff --git a/apps/aether-gateway/src/tests/frontdoor/internal.rs b/apps/aether-gateway/src/tests/frontdoor/internal.rs index 86283873a..c2fba96f9 100644 --- a/apps/aether-gateway/src/tests/frontdoor/internal.rs +++ b/apps/aether-gateway/src/tests/frontdoor/internal.rs @@ -1,8 +1,9 @@ use super::{ - hash_api_key, sample_endpoint, sample_key, sample_models_candidate_row, sample_provider, - unrestricted_models_snapshot, InMemoryAuthApiKeySnapshotRepository, - InMemoryMinimalCandidateSelectionReadRepository, InMemoryProviderCatalogReadRepository, - InMemoryRequestCandidateRepository, DEVELOPMENT_ENCRYPTION_KEY, + hash_api_key, run_frontdoor_async_test, sample_endpoint, sample_key, + sample_models_candidate_row, sample_provider, unrestricted_models_snapshot, + InMemoryAuthApiKeySnapshotRepository, InMemoryMinimalCandidateSelectionReadRepository, + InMemoryProviderCatalogReadRepository, InMemoryRequestCandidateRepository, + DEVELOPMENT_ENCRYPTION_KEY, }; use crate::tests::{ any, build_router, build_router_with_state, build_state_with_execution_runtime_override, json, @@ -160,8 +161,15 @@ async fn gateway_returns_internal_gateway_plan_sync_proxy_public_action_without_ upstream_handle.abort(); } -#[tokio::test] -async fn gateway_handles_internal_gateway_execute_sync_locally() { +#[test] +fn gateway_handles_internal_gateway_execute_sync_locally() { + run_frontdoor_async_test( + "gateway_handles_internal_gateway_execute_sync_locally", + gateway_handles_internal_gateway_execute_sync_locally_impl(), + ); +} + +async fn gateway_handles_internal_gateway_execute_sync_locally_impl() { let upstream_hits = Arc::new(Mutex::new(0usize)); let upstream_hits_clone = Arc::clone(&upstream_hits); let fallback_probe = Router::new().route( diff --git a/apps/aether-gateway/src/tests/frontdoor/ops.rs b/apps/aether-gateway/src/tests/frontdoor/ops.rs index cac36fd9b..e072e16f9 100644 --- a/apps/aether-gateway/src/tests/frontdoor/ops.rs +++ b/apps/aether-gateway/src/tests/frontdoor/ops.rs @@ -70,6 +70,7 @@ async fn gateway_exposes_frontdoor_manifest_without_proxying_upstream() { assert!(owned_routes .iter() .any(|value| value == "/v1/responses/compact")); + assert!(owned_routes.iter().any(|value| value == "/v1/alpha/search")); assert!(owned_routes.iter().any(|value| value == "/health")); assert!(owned_routes.iter().any(|value| value == "/v1/health")); assert!(owned_routes.iter().any(|value| value == "/v1/providers")); diff --git a/apps/aether-gateway/src/tests/frontdoor/public_support.rs b/apps/aether-gateway/src/tests/frontdoor/public_support.rs index 0c10c1f4e..7bd7761a1 100644 --- a/apps/aether-gateway/src/tests/frontdoor/public_support.rs +++ b/apps/aether-gateway/src/tests/frontdoor/public_support.rs @@ -5735,7 +5735,7 @@ async fn gateway_handles_users_me_usage_locally_without_proxying_upstream() { payload["records"][0]["cache_creation_ephemeral_5m_input_tokens"], 4 ); - assert_eq!(payload["records"][0]["effective_input_tokens"], 105); + assert_eq!(payload["records"][0]["effective_input_tokens"], 95); assert_eq!( payload["records"][0]["cache_creation_ephemeral_1h_input_tokens"], 6 @@ -5762,10 +5762,7 @@ async fn gateway_handles_users_me_usage_locally_without_proxying_upstream() { payload["summary_by_model"][0]["cache_creation_ephemeral_1h_tokens"], 6 ); - assert_eq!( - payload["summary_by_model"][0]["effective_input_tokens"], - 105 - ); + assert_eq!(payload["summary_by_model"][0]["effective_input_tokens"], 95); assert_eq!(payload["summary_by_model"][0]["total_input_context"], 120); assert!(payload.get("summary_by_provider").is_none()); assert_eq!(payload["billing"]["id"], "wallet-auth-1"); diff --git a/apps/aether-gateway/src/tests/frontdoor/public_support/dashboard.rs b/apps/aether-gateway/src/tests/frontdoor/public_support/dashboard.rs index b3424d167..86154975e 100644 --- a/apps/aether-gateway/src/tests/frontdoor/public_support/dashboard.rs +++ b/apps/aether-gateway/src/tests/frontdoor/public_support/dashboard.rs @@ -191,7 +191,7 @@ async fn gateway_handles_dashboard_stats_locally_without_proxying_upstream() { assert_eq!(response.status(), StatusCode::OK); let payload: serde_json::Value = response.json().await.expect("json body should parse"); assert_eq!(payload["today"]["requests"], 1); - assert_eq!(payload["today"]["tokens"], 160); + assert_eq!(payload["today"]["tokens"], 150); assert_eq!(payload["api_keys"]["total"], 2); assert_eq!(payload["api_keys"]["active"], 1); assert_eq!(payload["stats"][3]["subValue"], json!("输入 240 / 输出 60")); @@ -646,7 +646,7 @@ async fn gateway_handles_admin_dashboard_stats_locally_without_proxying_upstream assert_eq!(response.status(), StatusCode::OK); let payload: serde_json::Value = response.json().await.expect("json body should parse"); assert_eq!(payload["today"]["requests"], 2); - assert_eq!(payload["today"]["tokens"], 17_450); + assert_eq!(payload["today"]["tokens"], 16_250); assert_eq!(payload["today"]["cost"], json!(2.5)); assert_eq!(payload["cost_stats"]["cost_savings"], json!(0.025)); let stats = payload["stats"].as_array().expect("stats should be array"); @@ -664,10 +664,10 @@ async fn gateway_handles_admin_dashboard_stats_locally_without_proxying_upstream .iter() .find(|item| item["name"] == json!("今日 Token")) .expect("today token stats card should exist"); - assert_eq!(today_token_stats["value"], json!("17.4K")); + assert_eq!(today_token_stats["value"], json!("16.2K")); assert_eq!( today_token_stats["subValue"], - json!("输入 12.1K / 输出 3.1K · 写缓存 1.25K / 读缓存 1K") + json!("输入 10.9K / 输出 3.1K · 写缓存 1.25K / 读缓存 1K") ); assert_eq!(payload["users"]["total"], 2); assert_eq!(payload["users"]["active"], 1); diff --git a/apps/aether-gateway/src/tests/proxy.rs b/apps/aether-gateway/src/tests/proxy.rs index 965c242fe..e4c646bb0 100644 --- a/apps/aether-gateway/src/tests/proxy.rs +++ b/apps/aether-gateway/src/tests/proxy.rs @@ -89,6 +89,13 @@ fn sample_cli_auth_snapshot( } fn sample_provider(provider_id: &str) -> StoredProviderCatalogProvider { + sample_provider_with_request_timeout(provider_id, None) +} + +fn sample_provider_with_request_timeout( + provider_id: &str, + request_timeout_secs: Option, +) -> StoredProviderCatalogProvider { StoredProviderCatalogProvider::new( provider_id.to_string(), provider_id.to_string(), @@ -96,7 +103,17 @@ fn sample_provider(provider_id: &str) -> StoredProviderCatalogProvider { "custom".to_string(), ) .expect("provider should build") - .with_transport_fields(true, false, false, None, None, None, None, None, None) + .with_transport_fields( + true, + false, + false, + None, + None, + None, + request_timeout_secs, + None, + None, + ) } fn sample_endpoint(endpoint_id: &str, provider_id: &str) -> StoredProviderCatalogEndpoint { @@ -435,6 +452,7 @@ async fn gateway_forwards_public_request_to_remote_tunnel_owner_before_fallback_ async move { let (parts, body) = request.into_parts(); let raw_body = to_bytes(body, usize::MAX).await.expect("body should read"); + tokio::time::sleep(Duration::from_millis(40)).await; *seen_owner_inner.lock().expect("mutex should lock") = Some(SeenOwnerRequest { path: parts .uri @@ -511,7 +529,10 @@ async fn gateway_forwards_public_request_to_remote_tunnel_owner_before_fallback_ let (owner_url, owner_handle) = start_server(owner).await; let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( - vec![sample_provider("provider-owner")], + vec![sample_provider_with_request_timeout( + "provider-owner", + Some(0.1), + )], vec![sample_endpoint("endpoint-owner", "provider-owner")], vec![sample_key("key-owner", "provider-owner", "node-owner")], )); @@ -540,6 +561,12 @@ async fn gateway_forwards_public_request_to_remote_tunnel_owner_before_fallback_ state = state .with_data_state_for_tests(data_state) .with_tunnel_identity_for_tests("gateway-a", Some("http://gateway-a:8080")); + let short_timeout_client = reqwest::Client::builder() + .timeout(Duration::from_millis(10)) + .build() + .expect("test client should build"); + state.client = short_timeout_client.clone(); + state.owner_forward_client = short_timeout_client; state.remember_scheduler_affinity_target( "scheduler_affinity:api-key-affinity-1:openai:chat:gpt-4.1", crate::cache::SchedulerAffinityTarget { @@ -922,32 +949,39 @@ async fn gateway_streamifies_sync_json_from_remote_tunnel_owner_before_returning .unwrap_or_default() .to_string(), }); + let encoded_response = serde_json::to_vec(&json!({ + "id": "resp-codex-affinity-stream-123", + "object": "response", + "model": "gpt-5.4", + "status": "completed", + "output": [{ + "type": "message", + "id": "msg-codex-affinity-stream-123", + "role": "assistant", + "content": [{ + "type": "output_text", + "text": "Hello from affinity sync json", + "annotations": [] + }] + }], + "usage": { + "input_tokens": 1, + "output_tokens": 2, + "total_tokens": 3 + } + })) + .expect("body should encode"); + let split_at = encoded_response.len() / 2; + let first = axum::body::Bytes::copy_from_slice(&encoded_response[..split_at]); + let second = axum::body::Bytes::copy_from_slice(&encoded_response[split_at..]); + let response_body = Body::from_stream(async_stream::stream! { + yield Ok::<_, std::io::Error>(first); + tokio::time::sleep(Duration::from_millis(40)).await; + yield Ok::<_, std::io::Error>(second); + }); let mut response = Response::builder() .status(StatusCode::OK) - .body(Body::from( - serde_json::to_vec(&json!({ - "id": "resp-codex-affinity-stream-123", - "object": "response", - "model": "gpt-5.4", - "status": "completed", - "output": [{ - "type": "message", - "id": "msg-codex-affinity-stream-123", - "role": "assistant", - "content": [{ - "type": "output_text", - "text": "Hello from affinity sync json", - "annotations": [] - }] - }], - "usage": { - "input_tokens": 1, - "output_tokens": 2, - "total_tokens": 3 - } - })) - .expect("body should encode"), - )) + .body(response_body) .expect("response should build"); response.headers_mut().insert( http::header::CONTENT_TYPE, @@ -1001,6 +1035,10 @@ async fn gateway_streamifies_sync_json_from_remote_tunnel_owner_before_returning state = state .with_data_state_for_tests(data_state) .with_tunnel_identity_for_tests("gateway-a", Some("http://gateway-a:8080")); + state.client = reqwest::Client::builder() + .timeout(Duration::from_millis(10)) + .build() + .expect("short shared client should build"); state.remember_scheduler_affinity_target( "scheduler_affinity:api-key-affinity-cli-1:openai:responses:gpt-5.4", crate::cache::SchedulerAffinityTarget { diff --git a/apps/aether-gateway/src/tests/usage/pricing.rs b/apps/aether-gateway/src/tests/usage/pricing.rs index 1a426d9ed..d277521a3 100644 --- a/apps/aether-gateway/src/tests/usage/pricing.rs +++ b/apps/aether-gateway/src/tests/usage/pricing.rs @@ -804,7 +804,7 @@ fn gateway_records_openai_sync_usage_and_pricing_with_cache_tokens() { async fn gateway_records_openai_sync_usage_and_pricing_with_cache_tokens_impl() { let expected = ExpectedUsagePricing { input_tokens: 120, - billed_input_tokens: 100, + billed_input_tokens: 20, output_tokens: 40, cache_creation_tokens: 80, cache_creation_ephemeral_5m_tokens: 0, @@ -898,7 +898,7 @@ fn gateway_records_openai_stream_usage_and_pricing_with_cache_tokens() { async fn gateway_records_openai_stream_usage_and_pricing_with_cache_tokens_impl() { let expected = ExpectedUsagePricing { input_tokens: 240, - billed_input_tokens: 200, + billed_input_tokens: 120, output_tokens: 60, cache_creation_tokens: 80, cache_creation_ephemeral_5m_tokens: 0, diff --git a/apps/aether-gateway/src/tunnel/embedded/local_relay.rs b/apps/aether-gateway/src/tunnel/embedded/local_relay.rs index a7767fbfc..f599daa34 100644 --- a/apps/aether-gateway/src/tunnel/embedded/local_relay.rs +++ b/apps/aether-gateway/src/tunnel/embedded/local_relay.rs @@ -2,7 +2,10 @@ use std::io; use std::net::SocketAddr; use std::time::Duration; -use aether_contracts::tunnel::TUNNEL_RELAY_FORWARDED_BY_HEADER; +use aether_contracts::tunnel::{ + resolve_tunnel_request_timeouts, try_decode_tunnel_relay_request_meta, + TUNNEL_RELAY_FORWARDED_BY_HEADER, +}; use aether_runtime::{maybe_hold_axum_response_permit, AdmissionPermit}; use async_stream::stream; use axum::body::{Body, Bytes}; @@ -23,9 +26,6 @@ use super::protocol; use super::AppState; pub const TUNNEL_ERROR_HEADER: &str = "x-aether-tunnel-error"; -const MAX_RELAY_META_LEN: usize = 256 * 1024; -const MIN_RELAY_TIMEOUT_MS: u64 = 1; -const MAX_RELAY_TIMEOUT_MS: u64 = 300_000; struct StreamGuard { hub: std::sync::Arc, @@ -156,15 +156,7 @@ fn map_request_admission_error(error: super::RequestAdmissionError) -> String { } fn relay_header_timeout(meta: &protocol::RequestMeta) -> Duration { - let timeout_ms = if meta.stream { - meta.stream_first_byte_timeout_ms - .unwrap_or_else(|| meta.timeout.saturating_mul(1_000)) - } else { - meta.request_timeout_ms - .or(meta.stream_first_byte_timeout_ms) - .unwrap_or_else(|| meta.timeout.saturating_mul(1_000)) - }; - Duration::from_millis(timeout_ms.clamp(MIN_RELAY_TIMEOUT_MS, MAX_RELAY_TIMEOUT_MS)) + Duration::from_millis(resolve_tunnel_request_timeouts(meta).first_byte_ms) } pub async fn relay_request( @@ -252,15 +244,17 @@ pub async fn relay_request( if stream.is_none() { envelope_buf.extend_from_slice(&chunk); - let Some((parsed_meta, body_offset)) = (match try_decode_envelope_meta(&envelope_buf) { - Ok(result) => result, - Err(error) => { - return release_permit_response( - tunnel_error_response(StatusCode::BAD_REQUEST, "bad_request", &error), - request_permit, - ); - } - }) else { + let Some((parsed_meta, body_offset)) = + (match try_decode_tunnel_relay_request_meta(&envelope_buf) { + Ok(result) => result, + Err(error) => { + return release_permit_response( + tunnel_error_response(StatusCode::BAD_REQUEST, "bad_request", &error), + request_permit, + ); + } + }) + else { continue; }; @@ -427,27 +421,6 @@ fn release_permit_response( response } -fn try_decode_envelope_meta( - buffer: &BytesMut, -) -> Result, String> { - if buffer.len() < 4 { - return Ok(None); - } - let meta_len = u32::from_be_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]) as usize; - if meta_len > MAX_RELAY_META_LEN { - return Err("relay metadata too large".to_string()); - } - let meta_end = 4usize - .checked_add(meta_len) - .ok_or_else(|| "relay envelope length overflow".to_string())?; - if buffer.len() < meta_end { - return Ok(None); - } - let meta = serde_json::from_slice::(&buffer[4..meta_end]) - .map_err(|e| format!("invalid relay metadata: {e}"))?; - Ok(Some((meta, meta_end))) -} - fn append_headers(target: &mut HeaderMap, headers: &[(String, String)]) { for (name, value) in headers { if should_skip_local_relay_response_header(name) { @@ -542,6 +515,30 @@ mod tests { assert_eq!(relay_header_timeout(&meta), Duration::from_secs(7)); } + #[test] + fn relay_header_timeout_keeps_the_protocol_maximum_for_non_stream_requests() { + let meta = protocol::RequestMeta { + provider_id: None, + endpoint_id: None, + key_id: None, + method: "POST".to_string(), + url: "https://example.com/responses".to_string(), + headers: HashMap::new(), + stream: false, + request_timeout_ms: Some(aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_MS), + stream_first_byte_timeout_ms: None, + timeout: 60, + follow_redirects: None, + http1_only: false, + transport_profile: None, + }; + + assert_eq!( + relay_header_timeout(&meta), + Duration::from_millis(aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_MS) + ); + } + fn sample_connected_proxy_node(node_id: &str) -> StoredProxyNode { StoredProxyNode::new( node_id.to_string(), diff --git a/apps/aether-gateway/src/tunnel/mod.rs b/apps/aether-gateway/src/tunnel/mod.rs index 8845db8be..0b517e78a 100644 --- a/apps/aether-gateway/src/tunnel/mod.rs +++ b/apps/aether-gateway/src/tunnel/mod.rs @@ -8,6 +8,7 @@ use std::sync::Arc; use std::time::{Duration, SystemTime}; use aether_contracts::tunnel::{ + resolve_tunnel_request_timeouts, try_decode_tunnel_relay_request_meta, RequestMeta, TUNNEL_RELAY_FORWARDED_BY_HEADER, TUNNEL_RELAY_OWNER_INSTANCE_HEADER, }; use aether_data::repository::proxy_nodes::{ @@ -21,6 +22,7 @@ use axum::extract::ws::WebSocketUpgrade; use axum::extract::{ConnectInfo, Path, Request, State}; use axum::http::{HeaderMap, StatusCode}; use axum::response::IntoResponse; +use bytes::BytesMut; use futures_util::StreamExt; use serde::{Deserialize, Serialize}; use serde_json::json; @@ -61,6 +63,22 @@ const TUNNEL_INSTANCE_ID_ENV: &str = "AETHER_GATEWAY_INSTANCE_ID"; const TUNNEL_RELAY_BASE_URL_ENV: &str = "AETHER_TUNNEL_RELAY_BASE_URL"; const TUNNEL_ATTACHMENT_TTL_ENV: &str = "AETHER_TUNNEL_ATTACHMENT_TTL_SECS"; +pub(crate) async fn send_owner_forward_request( + request: reqwest::RequestBuilder, + first_byte_timeout: Option, +) -> Result { + match first_byte_timeout { + Some(timeout) => match tokio::time::timeout(timeout, request.send()).await { + Ok(result) => result.map_err(|error| error.to_string()), + Err(_) => Err(format!( + "owner gateway first byte timeout after {} ms", + timeout.as_millis() + )), + }, + None => request.send().await.map_err(|error| error.to_string()), + } +} + #[derive(Debug, Deserialize)] struct InternalTunnelHeartbeatRequest { node_id: String, @@ -776,8 +794,29 @@ async fn forward_relay_request_to_owner( ); } let limit_exceeded = Arc::new(AtomicBool::new(false)); + let prepared_body = + match prepare_owner_relay_request_body(body, body_limit, Arc::clone(&limit_exceeded)).await + { + Ok(prepared_body) => prepared_body, + Err(_) if limit_exceeded.load(Ordering::SeqCst) => { + return build_local_http_error_response( + trace_id, + None, + StatusCode::PAYLOAD_TOO_LARGE, + &format!("tunnel relay body exceeds {body_limit} bytes"), + ); + } + Err(error) => { + return build_local_http_error_response( + trace_id, + None, + StatusCode::BAD_REQUEST, + &error, + ); + } + }; - let mut upstream_request = state.client.post(owner_url); + let mut upstream_request = state.owner_forward_client.post(owner_url); for (name, value) in &parts.headers { if should_skip_request_header(name.as_str()) || name == http::header::HOST { continue; @@ -793,18 +832,23 @@ async fn forward_relay_request_to_owner( TUNNEL_RELAY_OWNER_INSTANCE_HEADER, owner.gateway_instance_id.as_str(), ); + let resolved_timeouts = resolve_tunnel_request_timeouts(&prepared_body.meta); + if let Some(timeout_ms) = resolved_timeouts.response_body_ms { + upstream_request = upstream_request.timeout(Duration::from_millis(timeout_ms)); + } if !parts.headers.contains_key(TRACE_ID_HEADER) { upstream_request = upstream_request.header(TRACE_ID_HEADER, trace_id); } - let upstream_response = match upstream_request - .body(build_owner_relay_request_body( - body, - body_limit, - Arc::clone(&limit_exceeded), - )) - .send() - .await + let first_byte_timeout = prepared_body + .meta + .stream + .then_some(Duration::from_millis(resolved_timeouts.first_byte_ms)); + let upstream_response = match send_owner_forward_request( + upstream_request.body(prepared_body.body), + first_byte_timeout, + ) + .await { Ok(response) => response, Err(err) if limit_exceeded.load(Ordering::SeqCst) => { @@ -862,14 +906,53 @@ fn request_content_length_exceeds_limit(headers: &HeaderMap, body_limit: usize) .is_some_and(|value| value > body_limit) } -fn build_owner_relay_request_body( +struct PreparedOwnerRelayRequestBody { + body: reqwest::Body, + meta: RequestMeta, +} + +async fn prepare_owner_relay_request_body( body: Body, body_limit: usize, limit_exceeded: Arc, -) -> reqwest::Body { +) -> Result { let mut body_stream = body.into_data_stream(); - reqwest::Body::wrap_stream(stream! { - let mut forwarded = 0usize; + let mut buffered_chunks = Vec::new(); + let mut meta_buffer = BytesMut::new(); + let mut forwarded = 0usize; + let mut meta = None; + + while meta.is_none() { + let Some(next_chunk) = body_stream.next().await else { + return Err("incomplete tunnel relay metadata".to_string()); + }; + match next_chunk { + Ok(chunk) => { + let next_forwarded = forwarded.saturating_add(chunk.len()); + if next_forwarded > body_limit { + limit_exceeded.store(true, Ordering::SeqCst); + return Err(format!("tunnel relay body exceeds {body_limit} bytes")); + } + forwarded = next_forwarded; + meta_buffer.extend_from_slice(&chunk); + buffered_chunks.push(chunk); + match try_decode_tunnel_relay_request_meta(&meta_buffer) { + Ok(Some((parsed, _))) => meta = Some(parsed), + Ok(None) => {} + Err(error) => return Err(error), + } + } + Err(error) => { + return Err(format!("tunnel relay body read failed: {error}")); + } + } + } + let meta = meta.ok_or_else(|| "incomplete tunnel relay metadata".to_string())?; + + let forwarded_body = reqwest::Body::wrap_stream(stream! { + for chunk in buffered_chunks { + yield Ok::(chunk); + } while let Some(next_chunk) = body_stream.next().await { match next_chunk { Ok(chunk) => { @@ -884,12 +967,17 @@ fn build_owner_relay_request_body( } yield Ok::(chunk); } - Err(err) => { - yield Err::(io::Error::other(err)); + Err(error) => { + yield Err::(io::Error::other(error)); break; } } } + }); + + Ok(PreparedOwnerRelayRequestBody { + body: forwarded_body, + meta, }) } @@ -1047,12 +1135,15 @@ fn parse_embedded_tunnel_heartbeat_request( mod tests { use super::{ apply_embedded_tunnel_heartbeat, apply_embedded_tunnel_node_status, current_unix_secs, - tunnel_attachment_key, GatewayDataState, TunnelAttachmentDirectory, TunnelAttachmentRecord, + prepare_owner_relay_request_body, tunnel_attachment_key, GatewayDataState, + TunnelAttachmentDirectory, TunnelAttachmentRecord, }; use aether_data::repository::proxy_nodes::{ InMemoryProxyNodeRepository, ProxyNodeReadRepository, StoredProxyNode, }; + use axum::body::Body; use serde_json::json; + use std::sync::atomic::AtomicBool; use std::sync::Arc; fn sample_proxy_node(node_id: &str) -> StoredProxyNode { @@ -1092,6 +1183,24 @@ mod tests { ) } + #[tokio::test] + async fn owner_relay_body_preparation_rejects_invalid_metadata() { + let mut envelope = Vec::new(); + envelope.extend_from_slice(&1u32.to_be_bytes()); + envelope.push(b'{'); + + let error = prepare_owner_relay_request_body( + Body::from(envelope), + 1024, + Arc::new(AtomicBool::new(false)), + ) + .await + .err() + .expect("invalid metadata should fail"); + + assert!(error.contains("invalid relay metadata")); + } + #[tokio::test] async fn embedded_tunnel_heartbeat_updates_proxy_node_repository() { let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![sample_proxy_node( diff --git a/apps/aether-tunnel/src/tunnel/stream_handler.rs b/apps/aether-tunnel/src/tunnel/stream_handler.rs index c626ef948..466cf6ab3 100644 --- a/apps/aether-tunnel/src/tunnel/stream_handler.rs +++ b/apps/aether-tunnel/src/tunnel/stream_handler.rs @@ -130,10 +130,6 @@ fn try_send_window_update(frame_tx: &FrameSender, stream_id: u32, bytes: usize) } } -/// Minimum allowed upstream request timeout (milliseconds). -const MIN_TIMEOUT_MS: u64 = 1; -/// Maximum allowed upstream request timeout (milliseconds). -const MAX_TIMEOUT_MS: u64 = 300_000; /// Match reqwest's default redirect budget so direct execution and tunnel relay /// fail at the same point instead of diverging after a different number of hops. const MAX_REDIRECTS: usize = 10; @@ -707,43 +703,14 @@ fn remaining_timeout(deadline: Instant) -> Option { } fn resolve_request_timeouts(meta: &RequestMeta) -> RequestTimeouts { - let first_byte_timeout = if meta.stream { - meta.stream_first_byte_timeout_ms - .map(timeout_duration_from_ms) - .unwrap_or_else(|| timeout_duration_from_legacy_secs(meta.timeout)) - } else { - meta.request_timeout_ms - .or(meta.stream_first_byte_timeout_ms) - .map(timeout_duration_from_ms) - .unwrap_or_else(|| timeout_duration_from_legacy_secs(meta.timeout)) - }; - - let response_body_timeout = if meta.stream { - None - } else { - Some( - meta.request_timeout_ms - .or(meta.stream_first_byte_timeout_ms) - .map(timeout_duration_from_ms) - .unwrap_or_else(|| timeout_duration_from_legacy_secs(meta.timeout)), - ) - }; + let resolved = aether_contracts::tunnel::resolve_tunnel_request_timeouts(meta); RequestTimeouts { - first_byte_timeout, - response_body_timeout, + first_byte_timeout: Duration::from_millis(resolved.first_byte_ms), + response_body_timeout: resolved.response_body_ms.map(Duration::from_millis), } } -fn timeout_duration_from_ms(ms: u64) -> Duration { - Duration::from_millis(ms.clamp(MIN_TIMEOUT_MS, MAX_TIMEOUT_MS)) -} - -fn timeout_duration_from_legacy_secs(secs: u64) -> Duration { - let ms = secs.saturating_mul(1_000); - timeout_duration_from_ms(ms) -} - async fn spool_request_body( stream_id: u32, mut body_rx: mpsc::Receiver, @@ -2176,6 +2143,18 @@ mod tests { ); } + #[test] + fn non_stream_request_timeouts_keep_the_protocol_maximum() { + let mut meta = sample_request_meta(); + meta.request_timeout_ms = Some(aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_MS); + + let timeouts = resolve_request_timeouts(&meta); + + let expected = Duration::from_millis(aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_MS); + assert_eq!(timeouts.first_byte_timeout, expected); + assert_eq!(timeouts.response_body_timeout, Some(expected)); + } + #[test] fn resolve_redirect_changes_post_to_get_for_302() { let current_url = url::Url::parse("https://redirect.test/start").expect("url"); diff --git a/crates/aether-admin/src/observability/stats.rs b/crates/aether-admin/src/observability/stats.rs index 719332f6b..bf93114f7 100644 --- a/crates/aether-admin/src/observability/stats.rs +++ b/crates/aether-admin/src/observability/stats.rs @@ -2222,6 +2222,6 @@ mod tests { build_user_leaderboard_items(&[item], &BTreeMap::new(), false, false, false); assert_eq!(leaderboard.len(), 1); - assert_eq!(leaderboard[0].tokens, 140); + assert_eq!(leaderboard[0].tokens, 120); } } diff --git a/crates/aether-admin/src/observability/usage.rs b/crates/aether-admin/src/observability/usage.rs index d2b9d4d9a..5ac2f2cb0 100644 --- a/crates/aether-admin/src/observability/usage.rs +++ b/crates/aether-admin/src/observability/usage.rs @@ -951,14 +951,7 @@ fn admin_usage_api_format_defaults_to_non_stream(item: &StoredRequestUsageAudit) let Some(value) = api_format else { return false; }; - matches!( - aether_ai_formats::normalize_api_format_alias(value).as_str(), - "openai:chat" - | "openai:responses" - | "openai:responses:compact" - | "openai:image" - | "claude:messages" - ) + aether_ai_formats::api_format_defaults_to_non_stream(value) } fn admin_usage_request_body_implies_default_non_stream(item: &StoredRequestUsageAudit) -> bool { @@ -1240,6 +1233,9 @@ fn admin_usage_active_request_json( if let Some(service_tier) = item.provider_service_tier() { value["service_tier"] = json!(service_tier); } + if let Some(actual_service_tier) = item.provider_actual_service_tier() { + value["actual_service_tier"] = json!(actual_service_tier); + } if let Some(image_progress) = image_progress { value["image_progress"] = image_progress.clone(); } @@ -1359,6 +1355,12 @@ pub fn admin_usage_record_json( if let Some(service_tier) = item.provider_service_tier() { object.insert("service_tier".to_string(), json!(service_tier)); } + if let Some(actual_service_tier) = item.provider_actual_service_tier() { + object.insert( + "actual_service_tier".to_string(), + json!(actual_service_tier), + ); + } payload } @@ -2800,6 +2802,33 @@ mod tests { assert_eq!(record["client_is_stream"], false); } + #[test] + fn client_requested_stream_defaults_to_non_stream_for_openai_search() { + let item = StoredRequestUsageAudit { + is_stream: true, + api_format: Some("openai:search".to_string()), + request_body: Some(json!({ + "id": "session-search-1", + "model": "gpt-5.6-sol", + "input": "current documentation" + })), + ..sample_usage("completed", Some(200), None) + }; + + assert!(!admin_usage_client_is_stream(&item)); + + let record = admin_usage_record_json( + &item, + &BTreeMap::new(), + &BTreeMap::new(), + false, + false, + None, + ); + assert_eq!(record["client_requested_stream"], false); + assert_eq!(record["client_is_stream"], false); + } + #[test] fn upstream_stream_prefers_request_metadata_flag() { let item = StoredRequestUsageAudit { @@ -3435,7 +3464,7 @@ mod tests { assert_eq!(payload["cache_creation_input_tokens"], 20); assert_eq!(payload["cache_creation_ephemeral_5m_input_tokens"], 12); assert_eq!(payload["cache_creation_ephemeral_1h_input_tokens"], 8); - assert_eq!(payload["total_tokens"], 50); + assert_eq!(payload["total_tokens"], 40); } #[test] @@ -3451,7 +3480,7 @@ mod tests { ..sample_usage("completed", Some(200), None) }; - assert_eq!(admin_usage_total_tokens(&item), 140); + assert_eq!(admin_usage_total_tokens(&item), 120); } #[test] diff --git a/crates/aether-admin/src/provider/endpoints.rs b/crates/aether-admin/src/provider/endpoints.rs index 856f4f499..539600704 100644 --- a/crates/aether-admin/src/provider/endpoints.rs +++ b/crates/aether-admin/src/provider/endpoints.rs @@ -46,15 +46,21 @@ pub fn key_api_formats_without_entry( ) } -fn active_endpoint_api_formats(endpoints: &[StoredProviderCatalogEndpoint]) -> Vec { - let mut formats = Vec::new(); - for endpoint in endpoints.iter().filter(|endpoint| endpoint.is_active) { +fn endpoint_api_format_sets( + endpoints: &[StoredProviderCatalogEndpoint], +) -> (Vec, Vec) { + let mut all = Vec::new(); + let mut active = Vec::new(); + for endpoint in endpoints { let api_format = normalize_endpoint_api_format(&endpoint.api_format); - if !formats.iter().any(|existing| existing == &api_format) { - formats.push(api_format); + if !all.iter().any(|existing| existing == &api_format) { + all.push(api_format.clone()); + } + if endpoint.is_active && !active.iter().any(|existing| existing == &api_format) { + active.push(api_format); } } - formats + (all, active) } fn configured_key_api_formats(key: &StoredProviderCatalogKey) -> Vec { @@ -82,33 +88,131 @@ pub fn endpoint_key_counts_by_format( ) -> (BTreeMap, BTreeMap) { let mut total = BTreeMap::new(); let mut active = BTreeMap::new(); - let inherited_api_formats = active_endpoint_api_formats(endpoints); + let (endpoint_api_formats, active_endpoint_api_formats) = endpoint_api_format_sets(endpoints); for key in keys { - if fixed_provider_key_inherits_api_formats( + let inherits_api_formats = fixed_provider_key_inherits_api_formats( provider_type, &key.auth_type, key.encrypted_auth_config.as_deref(), - ) { - for api_format in &inherited_api_formats { - *total.entry(api_format.clone()).or_insert(0) += 1; - if key.is_active { - *active.entry(api_format.clone()).or_insert(0) += 1; - } - } - continue; - } + ); + let has_unrestricted_api_format_scope = key + .api_formats + .as_ref() + .is_none_or(serde_json::Value::is_null); + let configured_api_formats = configured_key_api_formats(key); - for api_format in configured_key_api_formats(key) { + let candidate_api_formats = if inherits_api_formats { + &active_endpoint_api_formats + } else { + &endpoint_api_formats + }; + for api_format in candidate_api_formats.iter().filter(|api_format| { + inherits_api_formats + || has_unrestricted_api_format_scope + || configured_api_formats.iter().any(|allowed| { + aether_ai_formats::api_format_permission_covers(allowed, api_format) + }) + }) { *total.entry(api_format.clone()).or_insert(0) += 1; if key.is_active { - *active.entry(api_format).or_insert(0) += 1; + *active.entry(api_format.clone()).or_insert(0) += 1; } } } (total, active) } +#[cfg(test)] +mod endpoint_key_count_tests { + use super::*; + + fn sample_endpoint(id: &str, api_format: &str) -> StoredProviderCatalogEndpoint { + StoredProviderCatalogEndpoint::new( + id.to_string(), + "provider-1".to_string(), + api_format.to_string(), + None, + None, + true, + ) + .expect("endpoint should build") + } + + fn sample_key(id: &str, api_format: Option<&str>) -> StoredProviderCatalogKey { + let mut key = StoredProviderCatalogKey::new( + id.to_string(), + "provider-1".to_string(), + id.to_string(), + "api_key".to_string(), + None, + true, + ) + .expect("key should build"); + key.api_formats = api_format.map(|api_format| json!([api_format])); + key + } + + #[test] + fn endpoint_counts_follow_one_way_format_permissions() { + let endpoints = vec![ + sample_endpoint("responses", "openai:responses"), + sample_endpoint("search", "openai:search"), + ]; + let mut empty_scope_key = sample_key("empty-scope-key", None); + empty_scope_key.api_formats = Some(json!([])); + let keys = vec![ + sample_key("responses-key", Some("openai:responses")), + sample_key("search-key", Some("openai:search")), + sample_key("unrestricted-key", None), + empty_scope_key, + ]; + + let (total, active) = endpoint_key_counts_by_format("custom", &endpoints, &keys); + + assert_eq!(total.get("openai:responses"), Some(&2)); + assert_eq!(total.get("openai:search"), Some(&3)); + assert_eq!(active, total); + } + + #[test] + fn endpoint_counts_keep_scoped_keys_visible_on_inactive_endpoints() { + let mut endpoint = sample_endpoint("chat", "openai:chat"); + endpoint.is_active = false; + let mut empty_scope_key = sample_key("empty-scope-key", None); + empty_scope_key.api_formats = Some(json!([])); + let keys = vec![ + sample_key("chat-key", Some("openai:chat")), + sample_key("unrestricted-key", None), + empty_scope_key, + ]; + + let (total, active) = endpoint_key_counts_by_format("custom", &[endpoint], &keys); + + assert_eq!(total.get("openai:chat"), Some(&2)); + assert_eq!(active, total); + } + + #[test] + fn inherited_endpoint_counts_only_include_active_formats() { + let responses_endpoint = sample_endpoint("responses", "openai:responses"); + let mut search_endpoint = sample_endpoint("search", "openai:search"); + search_endpoint.is_active = false; + let mut inherited_key = sample_key("codex-key", Some("legacy:mismatch")); + inherited_key.auth_type = "oauth".to_string(); + + let (total, active) = endpoint_key_counts_by_format( + "codex", + &[responses_endpoint, search_endpoint], + &[inherited_key], + ); + + assert_eq!(total.get("openai:responses"), Some(&1)); + assert!(!total.contains_key("openai:search")); + assert_eq!(active, total); + } +} + fn masked_proxy_value(proxy: Option<&serde_json::Value>) -> serde_json::Value { let Some(proxy) = proxy.and_then(serde_json::Value::as_object) else { return serde_json::Value::Null; diff --git a/crates/aether-admin/src/provider/state.rs b/crates/aether-admin/src/provider/state.rs index 210cfc4e9..381ccf9cb 100644 --- a/crates/aether-admin/src/provider/state.rs +++ b/crates/aether-admin/src/provider/state.rs @@ -217,6 +217,14 @@ fn extract_openai_chatgpt_auth_fields_from_object( result.insert("user_id".to_string(), json!(user_id)); } + if let Some(is_fedramp) = auth + .and_then(|value| value.get("chatgpt_account_is_fedramp")) + .and_then(Value::as_bool) + .or_else(|| source.get("is_fedramp").and_then(Value::as_bool)) + { + result.insert("is_fedramp".to_string(), json!(is_fedramp)); + } + if let Some(organizations) = auth .and_then(|value| value.get("organizations")) .and_then(Value::as_array) @@ -250,6 +258,7 @@ pub fn enrich_admin_provider_oauth_auth_config( "plan_type", "user_id", "account_name", + "is_fedramp", ], ); @@ -268,6 +277,7 @@ pub fn enrich_admin_provider_oauth_auth_config( "plan_type", "user_id", "organizations", + "is_fedramp", ], ); @@ -288,6 +298,7 @@ pub fn enrich_admin_provider_oauth_auth_config( "plan_type", "user_id", "account_name", + "is_fedramp", ], ); let chatgpt_claim_fields = extract_openai_chatgpt_auth_fields_from_object(&claims); @@ -301,6 +312,7 @@ pub fn enrich_admin_provider_oauth_auth_config( "plan_type", "user_id", "organizations", + "is_fedramp", ], ); } @@ -428,6 +440,7 @@ mod tests { "chatgpt_account_user_id": "user-image__acc-image", "chatgpt_plan_type": "plus", "chatgpt_user_id": "user-image", + "chatgpt_account_is_fedramp": true, }, })); let token_payload = json!({ @@ -445,5 +458,6 @@ mod tests { ); assert_eq!(auth_config.get("plan_type"), Some(&json!("plus"))); assert_eq!(auth_config.get("user_id"), Some(&json!("user-image"))); + assert_eq!(auth_config.get("is_fedramp"), Some(&json!(true))); } } diff --git a/crates/aether-admin/src/system.rs b/crates/aether-admin/src/system.rs index 736089a6b..020366a9d 100644 --- a/crates/aether-admin/src/system.rs +++ b/crates/aether-admin/src/system.rs @@ -769,6 +769,12 @@ const ADMIN_API_FORMAT_DEFINITIONS: &[AdminApiFormatDefinition] = &[ default_path: "/v1/responses/compact", aliases: &["responses_compact"], }, + AdminApiFormatDefinition { + value: "openai:search", + label: "OpenAI Search", + default_path: "/v1/alpha/search", + aliases: &["openai_search", "search"], + }, AdminApiFormatDefinition { value: "openai:embedding", label: "OpenAI Embedding", @@ -1736,66 +1742,7 @@ pub fn admin_system_config_default_value(key: &str) -> Option "email_suffix_list" => Some(json!([])), "enable_format_conversion" => Some(json!(false)), "enable_model_directives" => Some(json!(false)), - "model_directives" => Some(json!({ - "reasoning_effort": { - "enabled": true, - "api_formats": { - "openai:chat": { - "enabled": true, - "mappings": { - "low": { "reasoning_effort": "low" }, - "medium": { "reasoning_effort": "medium" }, - "high": { "reasoning_effort": "high" }, - "xhigh": { "reasoning_effort": "xhigh" }, - "max": { "reasoning_effort": "max" }, - "fast": { "service_tier": "priority" } - } - }, - "openai:responses": { - "enabled": true, - "mappings": { - "low": { "reasoning": { "effort": "low" } }, - "medium": { "reasoning": { "effort": "medium" } }, - "high": { "reasoning": { "effort": "high" } }, - "xhigh": { "reasoning": { "effort": "xhigh" } }, - "max": { "reasoning": { "effort": "max" } }, - "fast": { "service_tier": "priority" } - } - }, - "openai:responses:compact": { - "enabled": true, - "mappings": { - "low": { "reasoning": { "effort": "low" } }, - "medium": { "reasoning": { "effort": "medium" } }, - "high": { "reasoning": { "effort": "high" } }, - "xhigh": { "reasoning": { "effort": "xhigh" } }, - "max": { "reasoning": { "effort": "max" } }, - "fast": { "service_tier": "priority" } - } - }, - "claude:messages": { - "enabled": true, - "mappings": { - "low": { "thinking": { "type": "enabled", "budget_tokens": 1024 } }, - "medium": { "thinking": { "type": "enabled", "budget_tokens": 4096 } }, - "high": { "thinking": { "type": "enabled", "budget_tokens": 8192 } }, - "xhigh": { "thinking": { "type": "enabled", "budget_tokens": 16384 } }, - "max": { "thinking": { "type": "enabled", "budget_tokens": 32768 } } - } - }, - "gemini:generate_content": { - "enabled": true, - "mappings": { - "low": { "generationConfig": { "thinkingConfig": { "thinkingBudget": 1024 } } }, - "medium": { "generationConfig": { "thinkingConfig": { "thinkingBudget": 4096 } } }, - "high": { "generationConfig": { "thinkingConfig": { "thinkingBudget": 8192 } } }, - "xhigh": { "generationConfig": { "thinkingConfig": { "thinkingBudget": 16384 } } }, - "max": { "generationConfig": { "thinkingConfig": { "thinkingBudget": -1 } } } - } - } - } - } - })), + "model_directives" => Some(aether_ai_formats::default_model_directives_config()), "keep_priority_on_conversion" => Some(json!(false)), "audit_log_retention_days" => Some(json!(30)), "enable_db_maintenance" => Some(json!(true)), diff --git a/crates/aether-ai-formats/src/api.rs b/crates/aether-ai-formats/src/api.rs index 041a504ef..889b7ac45 100644 --- a/crates/aether-ai-formats/src/api.rs +++ b/crates/aether-ai-formats/src/api.rs @@ -39,7 +39,8 @@ pub use crate::contracts::{ OPENAI_RESPONSES_COMPACT_SYNC_SUCCESS_REPORT_KIND, OPENAI_RESPONSES_STREAM_PLAN_KIND, OPENAI_RESPONSES_STREAM_SUCCESS_REPORT_KIND, OPENAI_RESPONSES_SYNC_ERROR_REPORT_KIND, OPENAI_RESPONSES_SYNC_FINALIZE_REPORT_KIND, OPENAI_RESPONSES_SYNC_PLAN_KIND, - OPENAI_RESPONSES_SYNC_SUCCESS_REPORT_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, + OPENAI_RESPONSES_SYNC_SUCCESS_REPORT_KIND, OPENAI_SEARCH_SYNC_PLAN_KIND, + OPENAI_SEARCH_SYNC_SUCCESS_REPORT_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_VIDEO_CONTENT_PLAN_KIND, OPENAI_VIDEO_CREATE_SYNC_FINALIZE_REPORT_KIND, OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND, OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND, @@ -56,11 +57,24 @@ pub use crate::formats::openai::image::stream::{ maybe_build_openai_image_sync_finalize_product, OpenAiImageStreamState, OpenAiImageSyncFinalizeProduct, }; +pub use crate::formats::openai::prompt_cache::resolve_openai_prompt_cache_ttl_minutes; pub use crate::formats::openai::shared::{ copy_request_number_field, copy_request_number_field_as, map_openai_reasoning_effort_to_claude_output, map_openai_reasoning_effort_to_gemini_budget, parse_openai_stop_sequences, resolve_openai_chat_max_tokens, value_as_u64, }; +pub use crate::formats::openai::{ + reasoning::{ + validate_openai_reasoning_request, OpenAiReasoningContractViolation, + OpenAiReasoningViolationKind, + }, + request_contract::{ + finalize_openai_provider_request, + finalize_openai_provider_request_with_codex_model_capabilities, + validate_openai_provider_request_contract, OpenAiProviderRequestContractViolation, + OpenAiProviderRequestFinalization, + }, +}; pub use crate::formats::shared::error_body::{ build_core_error_body_for_client_format, is_core_error_finalize_kind, LocalCoreSyncErrorKind, }; @@ -78,9 +92,16 @@ pub use crate::formats::shared::image_bridge::{ pub use crate::formats::shared::model_directives::{ apply_model_directive_mapping_patch, apply_model_directive_overrides_from_model, apply_model_directive_overrides_from_request, claude_model_uses_adaptive_effort, - extract_gemini_model_from_path, gemini_model_uses_thinking_level, model_directive_base_model, - normalize_model_directive_model, parse_model_directive, ModelDirective, ModelOverride, - ReasoningEffort, ServiceTier, + default_model_directive_mapping_patch, default_model_directive_suffixes, + default_model_directives_config, extract_gemini_model_from_path, + gemini_model_uses_thinking_level, model_directive_base_model, + model_directive_builtin_suffix_supported_for_source_model, + model_directive_suffix_has_builtin_mapping, normalize_model_directive_model, + openai_model_supports_prompt_cache_options, parse_model_directive, + parse_model_directive_with_suffixes, reasoning_effort_supported_for_model, ModelDirective, + ModelDirectiveSuffixResolution, ModelOverride, ReasoningEffort, ServiceTier, + CROSS_PROVIDER_MODEL_DIRECTIVE_SUFFIXES, MODEL_DIRECTIVE_API_FORMATS, + OPENAI_MODEL_DIRECTIVE_SUFFIXES, }; pub use crate::formats::shared::passthrough::{ resolve_stream_spec as resolve_local_same_format_stream_spec, @@ -89,7 +110,8 @@ pub use crate::formats::shared::passthrough::{ }; pub use crate::formats::shared::request::{ endpoint_config_forces_upstream_stream_policy, enforce_request_body_stream_field, - force_upstream_streaming_for_provider, parse_direct_request_body, + forbid_upstream_streaming_for_provider, force_upstream_streaming_for_provider, + parse_direct_request_body, resolve_upstream_is_stream_for_provider, resolve_upstream_is_stream_from_endpoint_config, }; pub use crate::formats::shared::request_matrix::{ @@ -149,11 +171,17 @@ pub use crate::formats::{ embedding::spec::resolve_sync_spec as resolve_openai_embedding_sync_spec, responses::{ codex::{ + apply_codex_openai_compact_terminal_headers, apply_codex_openai_responses_chat_body_edits, + apply_codex_openai_responses_lite_header_with_capabilities, apply_codex_openai_responses_special_body_edits, - apply_codex_openai_responses_special_headers, + apply_codex_openai_responses_special_body_edits_with_source_model_and_capabilities, + apply_codex_openai_special_headers, apply_openai_responses_compact_special_body_edits, - CODEX_OPENAI_IMAGE_DEFAULT_MODEL, CODEX_OPENAI_IMAGE_DEFAULT_OUTPUT_FORMAT, + build_codex_model_catalog_metadata, parse_codex_auth_identity, + resolve_codex_responses_model_capabilities, CodexAuthIdentity, + CodexResponsesModelCapabilities, 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, }, @@ -189,10 +217,13 @@ pub use crate::formats::{ }, openai::image::{ request::{ - build_chatgpt_web_image_request_body, build_openai_image_api_provider_request_body, - build_openai_image_provider_request_body, default_model_for_openai_image_operation, - is_openai_image_stream_request, normalize_openai_image_request, + build_chatgpt_web_image_request_body, + build_codex_openai_image_api_provider_request_body, + build_openai_image_api_provider_request_body, build_openai_image_provider_request_body, + default_model_for_openai_image_operation, is_openai_image_stream_request, + normalize_openai_image_quality, normalize_openai_image_request, normalize_openai_image_request_with_options, openai_image_operation_from_path, + project_codex_openai_image_api_request_body, project_openai_image_api_request_body, resolve_requested_openai_image_model_for_request, ChatGptWebImageRequestError, NormalizedOpenAiImageRequest, OpenAiImageNormalizeOptions, OpenAiImageOperation, OpenAiImageResponseFormat, @@ -244,8 +275,10 @@ pub use aether_ai_formats::formats::conversion::response::{ convert_openai_responses_response_to_openai_chat, OpenAiResponsesResponseUsage, }; pub use aether_ai_formats::{ - api_format_alias_matches, api_format_storage_aliases, is_openai_responses_compact_format, - is_openai_responses_family_format, is_openai_responses_format, normalize_api_format_alias, + api_format_alias_matches, api_format_permission_covers, api_format_permission_storage_aliases, + api_format_storage_aliases, intersect_api_format_allowed_lists, + is_openai_responses_compact_format, is_openai_responses_family_format, + is_openai_responses_format, normalize_api_format_alias, }; pub use aether_ai_formats::{ canonical_request_unknown_block_count, canonical_response_unknown_block_count, diff --git a/crates/aether-ai-formats/src/contracts/mod.rs b/crates/aether-ai-formats/src/contracts/mod.rs index 3f7752f66..62f11b066 100644 --- a/crates/aether-ai-formats/src/contracts/mod.rs +++ b/crates/aether-ai-formats/src/contracts/mod.rs @@ -23,9 +23,9 @@ pub use plan_kinds::{ 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_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, + 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, @@ -52,5 +52,6 @@ pub use report_kinds::{ OPENAI_RESPONSES_COMPACT_SYNC_FINALIZE_REPORT_KIND, OPENAI_RESPONSES_COMPACT_SYNC_SUCCESS_REPORT_KIND, OPENAI_RESPONSES_STREAM_SUCCESS_REPORT_KIND, OPENAI_RESPONSES_SYNC_ERROR_REPORT_KIND, OPENAI_RESPONSES_SYNC_FINALIZE_REPORT_KIND, - OPENAI_RESPONSES_SYNC_SUCCESS_REPORT_KIND, OPENAI_VIDEO_CREATE_SYNC_FINALIZE_REPORT_KIND, + OPENAI_RESPONSES_SYNC_SUCCESS_REPORT_KIND, OPENAI_SEARCH_SYNC_SUCCESS_REPORT_KIND, + OPENAI_VIDEO_CREATE_SYNC_FINALIZE_REPORT_KIND, }; diff --git a/crates/aether-ai-formats/src/contracts/plan_kinds.rs b/crates/aether-ai-formats/src/contracts/plan_kinds.rs index d0d4eecb7..2a9e61329 100644 --- a/crates/aether-ai-formats/src/contracts/plan_kinds.rs +++ b/crates/aether-ai-formats/src/contracts/plan_kinds.rs @@ -23,6 +23,7 @@ pub const OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND: &str = "openai_video_create_sync"; pub const OPENAI_CHAT_SYNC_PLAN_KIND: &str = "openai_chat_sync"; pub const OPENAI_EMBEDDING_SYNC_PLAN_KIND: &str = "openai_embedding_sync"; pub const OPENAI_RERANK_SYNC_PLAN_KIND: &str = "openai_rerank_sync"; +pub const OPENAI_SEARCH_SYNC_PLAN_KIND: &str = "openai_search_sync"; pub const GEMINI_EMBEDDING_SYNC_PLAN_KIND: &str = "gemini_embedding_sync"; pub const OPENAI_RESPONSES_SYNC_PLAN_KIND: &str = "openai_responses_sync"; pub const OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND: &str = "openai_responses_compact_sync"; diff --git a/crates/aether-ai-formats/src/contracts/report_kinds.rs b/crates/aether-ai-formats/src/contracts/report_kinds.rs index bb2e769e0..ce19b1b4f 100644 --- a/crates/aether-ai-formats/src/contracts/report_kinds.rs +++ b/crates/aether-ai-formats/src/contracts/report_kinds.rs @@ -29,6 +29,7 @@ pub const OPENAI_RESPONSES_SYNC_SUCCESS_REPORT_KIND: &str = "openai_responses_sy pub const OPENAI_RESPONSES_COMPACT_SYNC_SUCCESS_REPORT_KIND: &str = "openai_responses_compact_sync_success"; pub const OPENAI_EMBEDDING_SYNC_SUCCESS_REPORT_KIND: &str = "openai_embedding_sync_success"; +pub const OPENAI_SEARCH_SYNC_SUCCESS_REPORT_KIND: &str = "openai_search_sync_success"; pub const GEMINI_EMBEDDING_SYNC_SUCCESS_REPORT_KIND: &str = "gemini_embedding_sync_success"; pub const OPENAI_IMAGE_SYNC_SUCCESS_REPORT_KIND: &str = "openai_image_sync_success"; pub const CLAUDE_CLI_SYNC_SUCCESS_REPORT_KIND: &str = "claude_cli_sync_success"; diff --git a/crates/aether-ai-formats/src/formats/claude/messages/stream.rs b/crates/aether-ai-formats/src/formats/claude/messages/stream.rs index 720f6dd63..ea5156b57 100644 --- a/crates/aether-ai-formats/src/formats/claude/messages/stream.rs +++ b/crates/aether-ai-formats/src/formats/claude/messages/stream.rs @@ -704,6 +704,11 @@ impl ClaudeClientEmitter { }; self.emit_content_part(part) } + CanonicalStreamEvent::OpenAiResponsesOutputItem { .. } => Err( + AiSurfaceFinalizeError::new( + "OpenAI Responses output items cannot be converted losslessly to Claude Messages", + ), + ), CanonicalStreamEvent::ToolCallStart { index, call_id, diff --git a/crates/aether-ai-formats/src/formats/conversion/request.rs b/crates/aether-ai-formats/src/formats/conversion/request.rs index 1178f5910..f339c27dc 100644 --- a/crates/aether-ai-formats/src/formats/conversion/request.rs +++ b/crates/aether-ai-formats/src/formats/conversion/request.rs @@ -221,7 +221,7 @@ mod tests { } #[test] - fn claude_request_to_chat_maps_max_reasoning_effort_to_xhigh() { + fn claude_request_to_chat_preserves_max_reasoning_effort() { let body = json!({ "model": "claude-sonnet", "messages": [{"role": "user", "content": "hello"}], @@ -233,7 +233,7 @@ mod tests { let converted = normalize_claude_request_to_openai_chat_request(&body).expect("openai chat request"); - assert_eq!(converted["reasoning_effort"], "xhigh"); + assert_eq!(converted["reasoning_effort"], "max"); } #[test] @@ -289,11 +289,6 @@ mod tests { "name": "mcp__mapsWeather", "arguments": "{\"city\":\"Hangzhou\"}" }, - { - "type": "web_search_call", - "id": "ignored_web_search", - "action": {"query": "should be skipped"} - }, { "type": "function_call", "call_id": call_id_two, @@ -395,7 +390,7 @@ mod tests { assert_eq!(converted["prompt_cache_key"], "cache_123"); assert_eq!(converted["safety_identifier"], "user_123"); assert!(converted.get("include").is_none()); - assert!(converted.get("store").is_none()); + assert_eq!(converted["store"], false); assert!(converted.get("text").is_none()); assert!(converted.get("reasoning").is_none()); } @@ -880,7 +875,7 @@ mod tests { } #[test] - fn openai_responses_request_normalizer_strips_content_cache_control() { + fn openai_responses_same_format_preserves_content_extensions() { let body = json!({ "model": "gpt-5.1", "input": [{ @@ -904,7 +899,10 @@ mod tests { .expect("responses request"); assert_eq!(converted["prompt_cache_key"], "cache_123"); - assert!(!converted["input"].to_string().contains("cache_control")); + assert_eq!( + converted["input"][0]["content"][0]["cache_control"], + json!({"type": "ephemeral"}) + ); } #[test] @@ -925,7 +923,7 @@ mod tests { ) .expect("responses request"); - assert_eq!(converted["reasoning"]["effort"], "xhigh"); + assert_eq!(converted["reasoning"]["effort"], "max"); assert_eq!(converted["reasoning"]["summary"], "auto"); } diff --git a/crates/aether-ai-formats/src/formats/gemini/generate_content/stream.rs b/crates/aether-ai-formats/src/formats/gemini/generate_content/stream.rs index e643a7d8e..01e954666 100644 --- a/crates/aether-ai-formats/src/formats/gemini/generate_content/stream.rs +++ b/crates/aether-ai-formats/src/formats/gemini/generate_content/stream.rs @@ -490,6 +490,11 @@ impl GeminiClientEmitter { None, ) } + CanonicalStreamEvent::OpenAiResponsesOutputItem { .. } => Err( + AiSurfaceFinalizeError::new( + "OpenAI Responses output items cannot be converted losslessly to Gemini GenerateContent", + ), + ), CanonicalStreamEvent::ToolCallStart { index, call_id, diff --git a/crates/aether-ai-formats/src/formats/id.rs b/crates/aether-ai-formats/src/formats/id.rs index 1e064f928..fc83ffdc4 100644 --- a/crates/aether-ai-formats/src/formats/id.rs +++ b/crates/aether-ai-formats/src/formats/id.rs @@ -23,6 +23,7 @@ pub enum FormatId { OpenAiChat, OpenAiResponses, OpenAiResponsesCompact, + OpenAiSearch, OpenAiEmbedding, OpenAiRerank, ClaudeMessages, @@ -49,6 +50,7 @@ impl FormatId { Self::OpenAiChat | Self::OpenAiResponses | Self::OpenAiResponsesCompact + | Self::OpenAiSearch | Self::OpenAiEmbedding | Self::OpenAiRerank => FormatFamily::OpenAi, Self::ClaudeMessages => FormatFamily::Claude, @@ -73,6 +75,7 @@ impl FormatId { Self::OpenAiChat => "openai:chat", Self::OpenAiResponses => "openai:responses", Self::OpenAiResponsesCompact => "openai:responses:compact", + Self::OpenAiSearch => "openai:search", Self::OpenAiEmbedding => "openai:embedding", Self::OpenAiRerank => "openai:rerank", Self::ClaudeMessages => "claude:messages", @@ -103,6 +106,9 @@ impl FromStr for FormatId { "openai:responses:compact" | "/v1/responses/compact" => { Ok(Self::OpenAiResponsesCompact) } + "openai:search" | "openai_search" | "search" | "/v1/alpha/search" => { + Ok(Self::OpenAiSearch) + } "openai:embedding" | "/v1/embeddings" => Ok(Self::OpenAiEmbedding), "openai:rerank" | "/v1/rerank" => Ok(Self::OpenAiRerank), "claude:messages" | "/v1/messages" => Ok(Self::ClaudeMessages), @@ -139,6 +145,56 @@ pub fn api_format_alias_matches(left: &str, right: &str) -> bool { normalize_api_format_alias(left) == normalize_api_format_alias(right) } +pub fn api_format_defaults_to_non_stream(value: &str) -> bool { + matches!( + normalize_api_format_alias(value).as_str(), + "openai:chat" + | "openai:responses" + | "openai:responses:compact" + | "openai:search" + | "openai:image" + | "claude:messages" + ) +} + +pub fn api_format_defaults_to_client_error_failover(value: &str) -> bool { + !matches!( + FormatId::parse(value).map(FormatId::canonical), + Some(FormatId::OpenAiSearch) + ) +} + +pub fn api_format_permission_covers(allowed_value: &str, requested_api_format: &str) -> bool { + let allowed_value = normalize_api_format_alias(allowed_value); + let requested_api_format = normalize_api_format_alias(requested_api_format); + !allowed_value.is_empty() + && !requested_api_format.is_empty() + && (allowed_value == requested_api_format + || allowed_value == "openai:responses" && requested_api_format == "openai:search") +} + +pub fn intersect_api_format_allowed_lists(left: &[String], right: &[String]) -> Vec { + let mut effective = Vec::new(); + for left_value in left { + for right_value in right { + let intersection = if api_format_permission_covers(right_value, left_value) { + Some(left_value) + } else if api_format_permission_covers(left_value, right_value) { + Some(right_value) + } else { + None + }; + if let Some(value) = intersection { + let normalized = normalize_api_format_alias(value); + if !effective.iter().any(|item| item == &normalized) { + effective.push(normalized); + } + } + } + } + effective +} + pub fn api_format_storage_aliases(value: &str) -> Vec { match FormatId::parse(value).map(FormatId::canonical) { Some(FormatId::AliyunMultimodalEmbedding) => vec![ @@ -149,6 +205,22 @@ pub fn api_format_storage_aliases(value: &str) -> Vec { } } +pub fn api_format_permission_storage_aliases(value: &str) -> Vec { + let requested_api_format = normalize_api_format_alias(value); + let mut aliases = api_format_storage_aliases(&requested_api_format); + for allowed_api_format in [FormatId::OpenAiResponses.as_str()] { + if !api_format_permission_covers(allowed_api_format, &requested_api_format) { + continue; + } + for alias in api_format_storage_aliases(allowed_api_format) { + if !aliases.iter().any(|existing| existing == &alias) { + aliases.push(alias); + } + } + } + aliases +} + pub fn is_openai_responses_format(value: &str) -> bool { normalize_api_format_alias(value) == "openai:responses" } @@ -179,7 +251,10 @@ pub fn api_format_uses_body_stream_field(value: &str) -> bool { #[cfg(test)] mod tests { use super::{ - api_format_alias_matches, api_format_storage_aliases, api_format_uses_body_stream_field, + api_format_alias_matches, api_format_defaults_to_client_error_failover, + api_format_defaults_to_non_stream, api_format_permission_covers, + api_format_permission_storage_aliases, api_format_storage_aliases, + api_format_uses_body_stream_field, intersect_api_format_allowed_lists, normalize_api_format_alias, FormatId, }; @@ -193,6 +268,102 @@ mod tests { assert_eq!(FormatId::parse("gemini:cli"), None); } + #[test] + fn responses_permission_covers_only_its_search_companion() { + assert!(api_format_permission_covers( + "OPENAI:RESPONSES", + "openai:search" + )); + assert!(api_format_permission_covers( + "openai:search", + "openai:search" + )); + assert!(!api_format_permission_covers( + "openai:search", + "openai:responses" + )); + assert!(!api_format_permission_covers( + "openai:responses", + "openai:chat" + )); + assert_eq!( + api_format_permission_storage_aliases("openai:search"), + vec!["openai:search".to_string(), "openai:responses".to_string()] + ); + assert_eq!( + api_format_permission_storage_aliases("openai:responses"), + vec!["openai:responses".to_string()] + ); + } + + #[test] + fn normalizes_openai_search_aliases() { + for alias in [ + "openai:search", + "OPENAI_SEARCH", + "search", + "/v1/alpha/search", + ] { + assert_eq!(FormatId::parse(alias), Some(FormatId::OpenAiSearch)); + assert_eq!(normalize_api_format_alias(alias), "openai:search"); + } + assert!(!api_format_uses_body_stream_field("openai:search")); + } + + #[test] + fn identifies_default_non_stream_formats_from_aliases() { + for format in [ + "/v1/chat/completions", + "/v1/responses", + "/v1/responses/compact", + "/v1/alpha/search", + "openai:image", + "/v1/messages", + ] { + assert!(api_format_defaults_to_non_stream(format), "{format}"); + } + assert!(!api_format_defaults_to_non_stream("gemini:interactions")); + } + + #[test] + fn search_defaults_to_passthrough_for_client_errors() { + for format in ["openai:search", "OPENAI_SEARCH", "/v1/alpha/search"] { + assert!( + !api_format_defaults_to_client_error_failover(format), + "{format}" + ); + } + assert!(api_format_defaults_to_client_error_failover( + "openai:responses" + )); + assert!(api_format_defaults_to_client_error_failover( + "custom:unknown" + )); + } + + #[test] + fn api_format_policy_intersection_keeps_the_narrowest_companion_scope() { + assert_eq!( + intersect_api_format_allowed_lists( + &["openai:responses".to_string()], + &["openai:search".to_string()], + ), + vec!["openai:search".to_string()] + ); + assert_eq!( + intersect_api_format_allowed_lists( + &["openai:search".to_string()], + &["OPENAI:RESPONSES".to_string()], + ), + vec!["openai:search".to_string()] + ); + assert!(intersect_api_format_allowed_lists( + &["openai:search".to_string()], + &["openai:chat".to_string()], + ) + .is_empty()); + } + #[test] fn parses_embedding_api_formats() { assert_eq!( diff --git a/crates/aether-ai-formats/src/formats/matrix.rs b/crates/aether-ai-formats/src/formats/matrix.rs index be66d4833..0c2568e54 100644 --- a/crates/aether-ai-formats/src/formats/matrix.rs +++ b/crates/aether-ai-formats/src/formats/matrix.rs @@ -57,6 +57,9 @@ pub fn request_candidate_api_format_preference( if client_api_format == "openai:responses:compact" { return (provider_api_format == "openai:responses:compact").then_some((0, 0)); } + if client_api_format == "openai:search" { + return (provider_api_format == "openai:search").then_some((0, 0)); + } if is_gemini_interactions_api_format(client_api_format.as_str()) { return (provider_api_format == "gemini:interactions").then_some((0, 0)); } @@ -109,6 +112,9 @@ pub fn request_candidate_api_formats( if client_api_format == "openai:responses:compact" { return vec!["openai:responses:compact"]; } + if client_api_format == "openai:search" { + return vec!["openai:search"]; + } if is_gemini_interactions_api_format(client_api_format.as_str()) { return GEMINI_INTERACTIONS_CANDIDATE_API_FORMATS.to_vec(); } @@ -377,6 +383,10 @@ mod tests { request_conversion_kind("openai:compact", "openai:responses:compact"), None ); + assert_eq!( + request_conversion_kind("openai:search", "openai:responses"), + None + ); assert_eq!( request_conversion_kind("gemini:generate_content", "claude:messages"), Some(RequestConversionKind::ToClaudeStandard) @@ -409,6 +419,26 @@ mod tests { } } + #[test] + fn search_candidate_registry_keeps_exact_protocol_identity() { + assert_eq!( + request_candidate_api_formats("openai:search", false), + vec!["openai:search"] + ); + assert_eq!( + request_candidate_api_formats("/v1/alpha/search", true), + vec!["openai:search"] + ); + assert_eq!( + request_candidate_api_format_preference("openai:search", "openai:search"), + Some((0, 0)) + ); + assert_eq!( + request_candidate_api_format_preference("openai:search", "openai:responses"), + None + ); + } + #[test] fn sync_response_conversion_registry_supports_bidirectional_standard_matrix() { assert_eq!( diff --git a/crates/aether-ai-formats/src/formats/mod.rs b/crates/aether-ai-formats/src/formats/mod.rs index 6dd9a9bf1..a50a036cd 100644 --- a/crates/aether-ai-formats/src/formats/mod.rs +++ b/crates/aether-ai-formats/src/formats/mod.rs @@ -13,7 +13,10 @@ pub mod shared; pub use context::{FormatContext, FormatError}; pub use id::{ - api_format_alias_matches, api_format_storage_aliases, is_openai_responses_compact_format, + api_format_alias_matches, api_format_defaults_to_client_error_failover, + api_format_defaults_to_non_stream, api_format_permission_covers, + api_format_permission_storage_aliases, api_format_storage_aliases, + intersect_api_format_allowed_lists, is_openai_responses_compact_format, is_openai_responses_family_format, is_openai_responses_format, normalize_api_format_alias, FormatFamily, FormatId, FormatProfile, }; diff --git a/crates/aether-ai-formats/src/formats/openai/chat/request.rs b/crates/aether-ai-formats/src/formats/openai/chat/request.rs index 2b06232df..71fb4c1ac 100644 --- a/crates/aether-ai-formats/src/formats/openai/chat/request.rs +++ b/crates/aether-ai-formats/src/formats/openai/chat/request.rs @@ -2,7 +2,6 @@ use serde_json::{json, Map, Value}; use crate::{ formats::context::FormatContext, - formats::openai::shared::OpenAiChatReasoningEffort, protocol::canonical::{ canonical_extension_object_mut, canonical_message_to_openai_chat_messages, canonical_response_format_to_openai, canonical_tool_choice_to_openai, @@ -367,11 +366,8 @@ fn non_empty_source_str<'a>(source: &'a Map, key: &str) -> Option .filter(|value| !value.trim().is_empty()) } -fn openai_chat_reasoning_effort(value: &str) -> Option<&'static str> { - if value.trim().eq_ignore_ascii_case("max") { - return Some("xhigh"); - } - OpenAiChatReasoningEffort::parse(value).map(OpenAiChatReasoningEffort::as_str) +fn openai_chat_reasoning_effort(value: &str) -> Option<&str> { + (!value.trim().is_empty()).then_some(value) } fn chat_compatible_openai_responses_extension_object( @@ -384,7 +380,14 @@ fn chat_compatible_openai_responses_extension_object( .filter(|(key, _)| { matches!( key.as_str(), - "verbosity" | "service_tier" | "prompt_cache_key" | "safety_identifier" | "user" + "verbosity" + | "store" + | "service_tier" + | "prompt_cache_key" + | "prompt_cache_options" + | "prompt_cache_retention" + | "safety_identifier" + | "user" ) }) .collect() diff --git a/crates/aether-ai-formats/src/formats/openai/chat/response.rs b/crates/aether-ai-formats/src/formats/openai/chat/response.rs index e137b265e..69d503169 100644 --- a/crates/aether-ai-formats/src/formats/openai/chat/response.rs +++ b/crates/aether-ai-formats/src/formats/openai/chat/response.rs @@ -7,8 +7,8 @@ use crate::{ protocol::canonical::{ canonical_blocks_to_openai_chat_message, canonical_stop_reason_to_openai, canonical_usage_to_openai, openai_extensions, openai_finish_reason_to_canonical, - openai_message_content_blocks, openai_usage_to_canonical, CanonicalContentBlock, - CanonicalResponse, CanonicalResponseOutput, CanonicalRole, + openai_message_content_blocks, openai_service_tier_extension, openai_usage_to_canonical, + CanonicalContentBlock, CanonicalResponse, CanonicalResponseOutput, CanonicalRole, OPENAI_RESPONSES_EXTENSION_NAMESPACE, OPENAI_RESPONSES_LEGACY_EXTENSION_NAMESPACE, }, }; @@ -17,20 +17,8 @@ pub fn from(body: &Value, _ctx: &FormatContext) -> Option { from_raw(body) } -pub fn to(response: &CanonicalResponse, ctx: &FormatContext) -> Option { - let mut body = to_raw(response); - if body.get("service_tier").is_none() { - if let Some(service_tier) = ctx - .report_context_value() - .get("original_request_body") - .and_then(Value::as_object) - .and_then(|request| request.get("service_tier")) - .cloned() - { - body["service_tier"] = service_tier; - } - } - Some(body) +pub fn to(response: &CanonicalResponse, _ctx: &FormatContext) -> Option { + Some(to_raw(response)) } pub fn from_raw(body_json: &Value) -> Option { @@ -186,17 +174,7 @@ pub fn to_raw(canonical: &CanonicalResponse) -> Value { { response["created"] = Value::from(created_at); } - if let Some(service_tier) = canonical - .extensions - .get(OPENAI_RESPONSES_EXTENSION_NAMESPACE) - .or_else(|| { - canonical - .extensions - .get(OPENAI_RESPONSES_LEGACY_EXTENSION_NAMESPACE) - }) - .and_then(|value| value.get("service_tier")) - .cloned() - { + if let Some(service_tier) = openai_service_tier_extension(&canonical.extensions).cloned() { response["service_tier"] = service_tier; } response diff --git a/crates/aether-ai-formats/src/formats/openai/chat/stream.rs b/crates/aether-ai-formats/src/formats/openai/chat/stream.rs index 8207f98aa..c432b8df7 100644 --- a/crates/aether-ai-formats/src/formats/openai/chat/stream.rs +++ b/crates/aether-ai-formats/src/formats/openai/chat/stream.rs @@ -10,6 +10,13 @@ use crate::formats::shared::sse::{encode_done_sse, encode_json_sse}; use crate::formats::shared::stream_core::common::*; use crate::formats::shared::AiSurfaceFinalizeError; +fn normalize_openai_service_tier(value: Option<&str>) -> Option { + value + .map(str::trim) + .filter(|value| !value.is_empty() && value.len() <= 64) + .map(str::to_ascii_lowercase) +} + #[derive(Default)] struct OpenAIChatProviderToolState { id: Option, @@ -21,6 +28,7 @@ struct OpenAIChatProviderToolState { pub struct OpenAIChatProviderState { response_id: Option, model: Option, + actual_service_tier: Option, started: bool, finished: bool, pending_finish_reason: Option, @@ -46,6 +54,7 @@ struct OpenAIResponsesProviderToolResultState { pub struct OpenAIResponsesProviderState { response_id: Option, model: Option, + actual_service_tier: Option, started: bool, finished: bool, text_parts: BTreeMap, @@ -55,10 +64,15 @@ pub struct OpenAIResponsesProviderState { tool_results: BTreeMap, tool_index_by_key: BTreeMap, image_item_keys: BTreeSet, + opaque_completed_item_keys: BTreeSet, last_tool_index: Option, } impl OpenAIChatProviderState { + pub(crate) fn actual_service_tier(&self) -> Option<&str> { + self.actual_service_tier.as_deref() + } + fn finish_usage(value: Option<&Value>) -> Option { let usage_object = value?.as_object()?; let has_token_fields = [ @@ -128,6 +142,11 @@ impl OpenAIChatProviderState { .and_then(Value::as_str) .map(ToOwned::to_owned) .or_else(|| self.model.clone()); + if let Some(service_tier) = + normalize_openai_service_tier(chunk_object.get("service_tier").and_then(Value::as_str)) + { + self.actual_service_tier = Some(service_tier); + } let mut out = Vec::new(); let Some(chunk_choices) = chunk_object.get("choices").and_then(Value::as_array) else { @@ -368,6 +387,10 @@ impl OpenAIChatProviderState { } impl OpenAIResponsesProviderState { + pub(crate) fn actual_service_tier(&self) -> Option<&str> { + self.actual_service_tier.as_deref() + } + fn identity(&self, report_context: &Value) -> (String, String) { resolve_identity( self.response_id.as_deref(), @@ -655,6 +678,14 @@ impl OpenAIResponsesProviderState { if item.get("type").and_then(Value::as_str) != Some("function_call") { return; } + const MAPPED_FIELDS: &[&str] = &["type", "id", "call_id", "status", "name", "arguments"]; + if item + .keys() + .any(|field| !MAPPED_FIELDS.contains(&field.as_str())) + { + out.push(self.unknown_frame(report_context, Value::Object(item.clone()))); + return; + } self.ensure_started(report_context, out); let key = item .get("call_id") @@ -1048,6 +1079,20 @@ impl OpenAIResponsesProviderState { }); } + fn output_item_key(item: &Map) -> String { + let item_type = item.get("type").and_then(Value::as_str).unwrap_or_default(); + if let Some(item_id) = item.get("id").and_then(Value::as_str) { + return format!("{item_type}:id:{item_id}"); + } + if let Some(encrypted_content) = item.get("encrypted_content").and_then(Value::as_str) { + return format!("{item_type}:encrypted_content:{encrypted_content}"); + } + format!( + "{item_type}:{}", + serde_json::to_string(item).unwrap_or_default() + ) + } + fn emit_output_item( &mut self, report_context: &Value, @@ -1055,23 +1100,31 @@ impl OpenAIResponsesProviderState { item: &Map, output_index: Option, final_item: bool, - ) { + ) -> bool { match item.get("type").and_then(Value::as_str).unwrap_or_default() { - "function_call" => self.emit_tool_call_item(report_context, out, item, output_index), + "function_call" => { + self.emit_tool_call_item(report_context, out, item, output_index); + true + } "function_call_output" => { self.emit_tool_result_item(report_context, out, item, output_index); + true } "custom_tool_call" => { self.emit_custom_tool_call_item(report_context, out, item, output_index); + true } "local_shell_call" | "shell_call" => { self.emit_shell_tool_call_item(report_context, out, item, output_index); + true } "apply_patch_call" => { self.emit_apply_patch_tool_call_item(report_context, out, item, output_index); + true } "computer_call" => { self.emit_computer_tool_call_item(report_context, out, item, output_index); + true } "custom_tool_call_output" | "local_shell_call_output" @@ -1079,10 +1132,20 @@ impl OpenAIResponsesProviderState { | "apply_patch_call_output" | "computer_call_output" => { self.emit_generic_tool_result_item(report_context, out, item, output_index); + true + } + "message" => { + self.emit_message_item(report_context, out, item, output_index); + true + } + "reasoning" if final_item => { + self.emit_reasoning_item(report_context, out, item); + true + } + "reasoning" => { + self.ensure_started(report_context, out); + true } - "message" => self.emit_message_item(report_context, out, item, output_index), - "reasoning" if final_item => self.emit_reasoning_item(report_context, out, item), - "reasoning" => self.ensure_started(report_context, out), "image_generation_call" => { self.emit_image_generation_item( report_context, @@ -1091,16 +1154,57 @@ impl OpenAIResponsesProviderState { output_index, final_item, ); + true } "web_search_call" | "file_search_call" | "code_interpreter_call" | "mcp_call" => { if !final_item { self.ensure_started(report_context, out); } + true } - _ => out.push(self.unknown_frame(report_context, Value::Object(item.clone()))), + _ => false, } } + fn emit_output_item_event( + &mut self, + report_context: &Value, + out: &mut Vec, + raw_event: &Value, + item: &Map, + output_index: Option, + final_item: bool, + ) { + if self.emit_output_item(report_context, out, item, output_index, final_item) { + return; + } + + if item + .get("type") + .and_then(Value::as_str) + .is_none_or(str::is_empty) + { + out.push(self.unknown_frame(report_context, raw_event.clone())); + return; + } + + if final_item { + self.opaque_completed_item_keys + .insert(Self::output_item_key(item)); + } + self.ensure_started(report_context, out); + let (id, model) = self.identity(report_context); + out.push(CanonicalStreamFrame { + id, + model, + event: CanonicalStreamEvent::OpenAiResponsesOutputItem { + output_index, + item: Value::Object(item.clone()), + raw_event: raw_event.clone(), + }, + }); + } + fn emit_response_output_items( &mut self, report_context: &Value, @@ -1117,7 +1221,13 @@ impl OpenAIResponsesProviderState { let Some(item) = raw_item.as_object() else { continue; }; - self.emit_output_item(report_context, out, item, Some(output_index), true); + if !self.emit_output_item(report_context, out, item, Some(output_index), true) + && !self + .opaque_completed_item_keys + .contains(&Self::output_item_key(item)) + { + out.push(self.unknown_frame(report_context, Value::Object(item.clone()))); + } } } @@ -1141,6 +1251,11 @@ impl OpenAIResponsesProviderState { .and_then(Value::as_str) .map(ToOwned::to_owned) .or_else(|| self.model.clone()); + if let Some(service_tier) = + normalize_openai_service_tier(response.get("service_tier").and_then(Value::as_str)) + { + self.actual_service_tier = Some(service_tier); + } } match value @@ -1326,7 +1441,14 @@ impl OpenAIResponsesProviderState { .get("output_index") .and_then(Value::as_u64) .map(|value| value as usize); - self.emit_output_item(report_context, &mut out, item, output_index, false); + self.emit_output_item_event( + report_context, + &mut out, + &value, + item, + output_index, + false, + ); } "response.custom_tool_call_input.delta" => { let delta = value @@ -1578,7 +1700,14 @@ impl OpenAIResponsesProviderState { .get("output_index") .and_then(Value::as_u64) .map(|value| value as usize); - self.emit_output_item(report_context, &mut out, item, output_index, true); + self.emit_output_item_event( + report_context, + &mut out, + &value, + item, + output_index, + true, + ); } "response.incomplete" => { let Some(response) = value.get("response").and_then(Value::as_object) else { @@ -1679,6 +1808,7 @@ impl OpenAIResponsesProviderState { pub struct OpenAIChatClientEmitter { response_id: Option, model: Option, + actual_service_tier: Option, started: bool, finished: bool, next_tool_call_index: usize, @@ -1724,6 +1854,7 @@ fn web_search_query_from_arguments(arguments: &str) -> String { pub struct OpenAIResponsesClientEmitter { response_id: Option, model: Option, + actual_service_tier: Option, created_at: Option, message_item_id: Option, reasoning_item_id: Option, @@ -1744,9 +1875,36 @@ pub struct OpenAIResponsesClientEmitter { tool_calls: BTreeMap, tool_results: BTreeMap, image_generation_items: BTreeMap, + opaque_output_items: BTreeMap, + opaque_output_indexes: BTreeMap, } impl OpenAIChatClientEmitter { + pub(crate) fn set_actual_service_tier(&mut self, value: Option<&str>) { + if value.is_some_and(|value| { + self.actual_service_tier + .as_deref() + .is_some_and(|current| current.eq_ignore_ascii_case(value.trim())) + }) { + return; + } + if let Some(value) = normalize_openai_service_tier(value) { + self.actual_service_tier = Some(value); + } + } + + fn encode_chunk(&self, mut chunk: Value) -> Result, AiSurfaceFinalizeError> { + if let (Some(service_tier), Some(object)) = + (self.actual_service_tier.as_ref(), chunk.as_object_mut()) + { + object.insert( + "service_tier".to_string(), + Value::String(service_tier.clone()), + ); + } + encode_json_sse(None, &chunk) + } + fn update_identity(&mut self, frame: &CanonicalStreamFrame) { self.response_id = Some(frame.id.clone()); self.model = Some(frame.model.clone()); @@ -1757,15 +1915,12 @@ impl OpenAIChatClientEmitter { return Ok(Vec::new()); } self.started = true; - encode_json_sse( - None, - &build_openai_chat_role_chunk( - self.response_id - .as_deref() - .unwrap_or("chatcmpl-local-stream"), - self.model.as_deref().unwrap_or("unknown"), - ), - ) + self.encode_chunk(build_openai_chat_role_chunk( + self.response_id + .as_deref() + .unwrap_or("chatcmpl-local-stream"), + self.model.as_deref().unwrap_or("unknown"), + )) } fn chat_tool_call_index(&mut self, canonical_index: usize) -> usize { @@ -1785,9 +1940,8 @@ impl OpenAIChatClientEmitter { CanonicalStreamEvent::Start => self.ensure_started(), CanonicalStreamEvent::TextDelta(text) => { let mut out = self.ensure_started()?; - out.extend(encode_json_sse( - None, - &build_openai_chat_chunk( + out.extend( + self.encode_chunk(build_openai_chat_chunk( self.response_id .as_deref() .unwrap_or("chatcmpl-local-stream"), @@ -1795,61 +1949,54 @@ impl OpenAIChatClientEmitter { text, None, None, - ), - )?); + ))?, + ); Ok(out) } CanonicalStreamEvent::ReasoningDelta(text) => { let mut out = self.ensure_started()?; - out.extend(encode_json_sse( - None, - &json!({ - "id": self.response_id - .as_deref() - .unwrap_or("chatcmpl-local-stream"), - "object": "chat.completion.chunk", - "model": self.model.as_deref().unwrap_or("unknown"), - "choices": [{ - "index": 0, - "delta": { - "reasoning_content": text, - }, - "finish_reason": Value::Null - }] - }), - )?); + out.extend(self.encode_chunk(json!({ + "id": self.response_id + .as_deref() + .unwrap_or("chatcmpl-local-stream"), + "object": "chat.completion.chunk", + "model": self.model.as_deref().unwrap_or("unknown"), + "choices": [{ + "index": 0, + "delta": { + "reasoning_content": text, + }, + "finish_reason": Value::Null + }] + }))?); Ok(out) } CanonicalStreamEvent::ReasoningSummaryDone => { // CPA strategy: emit "\n\n" as paragraph separator between // reasoning sections, matching CPA's Chat downstream behavior. let mut out = self.ensure_started()?; - out.extend(encode_json_sse( - None, - &json!({ - "id": self.response_id - .as_deref() - .unwrap_or("chatcmpl-local-stream"), - "object": "chat.completion.chunk", - "model": self.model.as_deref().unwrap_or("unknown"), - "choices": [{ - "index": 0, - "delta": { - "reasoning_content": "\n\n", - }, - "finish_reason": Value::Null - }] - }), - )?); + out.extend(self.encode_chunk(json!({ + "id": self.response_id + .as_deref() + .unwrap_or("chatcmpl-local-stream"), + "object": "chat.completion.chunk", + "model": self.model.as_deref().unwrap_or("unknown"), + "choices": [{ + "index": 0, + "delta": { + "reasoning_content": "\n\n", + }, + "finish_reason": Value::Null + }] + }))?); Ok(out) } CanonicalStreamEvent::ReasoningSignature(_) => Ok(Vec::new()), CanonicalStreamEvent::ContentPart(part) => { let placeholder = openai_stream_placeholder_for_content_part(&part); let mut out = self.ensure_started()?; - out.extend(encode_json_sse( - None, - &build_openai_chat_chunk( + out.extend( + self.encode_chunk(build_openai_chat_chunk( self.response_id .as_deref() .unwrap_or("chatcmpl-local-stream"), @@ -1857,8 +2004,8 @@ impl OpenAIChatClientEmitter { placeholder, None, None, - ), - )?); + ))?, + ); Ok(out) } CanonicalStreamEvent::ImageGenerationCall { item, .. } => { @@ -1867,9 +2014,8 @@ impl OpenAIChatClientEmitter { }; let placeholder = openai_stream_placeholder_for_content_part(&part); let mut out = self.ensure_started()?; - out.extend(encode_json_sse( - None, - &build_openai_chat_chunk( + out.extend( + self.encode_chunk(build_openai_chat_chunk( self.response_id .as_deref() .unwrap_or("chatcmpl-local-stream"), @@ -1877,10 +2023,15 @@ impl OpenAIChatClientEmitter { placeholder, None, None, - ), - )?); + ))?, + ); Ok(out) } + CanonicalStreamEvent::OpenAiResponsesOutputItem { .. } => { + Err(AiSurfaceFinalizeError::new( + "OpenAI Responses output items cannot be converted losslessly to OpenAI Chat", + )) + } CanonicalStreamEvent::ToolCallStart { index, call_id, @@ -1888,9 +2039,8 @@ impl OpenAIChatClientEmitter { } => { let mut out = self.ensure_started()?; let chat_index = self.chat_tool_call_index(index); - out.extend(encode_json_sse( - None, - &build_openai_chat_chunk( + out.extend( + self.encode_chunk(build_openai_chat_chunk( self.response_id .as_deref() .unwrap_or("chatcmpl-local-stream"), @@ -1906,35 +2056,32 @@ impl OpenAIChatClientEmitter { } })]), None, - ), - )?); + ))?, + ); Ok(out) } CanonicalStreamEvent::ToolCallArgumentsDelta { index, arguments } => { let mut out = self.ensure_started()?; let chat_index = self.chat_tool_call_index(index); - out.extend(encode_json_sse( - None, - &json!({ - "id": self.response_id - .as_deref() - .unwrap_or("chatcmpl-local-stream"), - "object": "chat.completion.chunk", - "model": self.model.as_deref().unwrap_or("unknown"), - "choices": [{ - "index": 0, - "delta": { - "tool_calls": [{ - "index": chat_index, - "function": { - "arguments": arguments, - } - }] - }, - "finish_reason": Value::Null - }] - }), - )?); + out.extend(self.encode_chunk(json!({ + "id": self.response_id + .as_deref() + .unwrap_or("chatcmpl-local-stream"), + "object": "chat.completion.chunk", + "model": self.model.as_deref().unwrap_or("unknown"), + "choices": [{ + "index": 0, + "delta": { + "tool_calls": [{ + "index": chat_index, + "function": { + "arguments": arguments, + } + }] + }, + "finish_reason": Value::Null + }] + }))?); Ok(out) } CanonicalStreamEvent::ToolResultDelta { @@ -1951,21 +2098,18 @@ impl OpenAIChatClientEmitter { delta.insert("name".to_string(), Value::String(name)); } delta.insert("content".to_string(), Value::String(content)); - out.extend(encode_json_sse( - None, - &json!({ - "id": self.response_id - .as_deref() - .unwrap_or("chatcmpl-local-stream"), - "object": "chat.completion.chunk", - "model": self.model.as_deref().unwrap_or("unknown"), - "choices": [{ - "index": 0, - "delta": Value::Object(delta), - "finish_reason": Value::Null - }] - }), - )?); + out.extend(self.encode_chunk(json!({ + "id": self.response_id + .as_deref() + .unwrap_or("chatcmpl-local-stream"), + "object": "chat.completion.chunk", + "model": self.model.as_deref().unwrap_or("unknown"), + "choices": [{ + "index": 0, + "delta": Value::Object(delta), + "finish_reason": Value::Null + }] + }))?); Ok(out) } CanonicalStreamEvent::UnknownEvent(payload) @@ -1984,27 +2128,25 @@ impl OpenAIChatClientEmitter { return Ok(Vec::new()); } let mut out = self.ensure_started()?; - out.extend(encode_json_sse( - None, - &build_openai_chat_finish_chunk( + out.extend( + self.encode_chunk(build_openai_chat_finish_chunk( self.response_id .as_deref() .unwrap_or("chatcmpl-local-stream"), self.model.as_deref().unwrap_or("unknown"), finish_reason.as_deref(), - ), - )?); + ))?, + ); if let Some(usage) = usage { - out.extend(encode_json_sse( - None, - &build_openai_chat_usage_chunk_from_usage( + out.extend( + self.encode_chunk(build_openai_chat_usage_chunk_from_usage( self.response_id .as_deref() .unwrap_or("chatcmpl-local-stream"), self.model.as_deref().unwrap_or("unknown"), &usage, - ), - )?); + ))?, + ); } out.extend(encode_done_sse()); self.finished = true; @@ -2017,16 +2159,13 @@ impl OpenAIChatClientEmitter { if !self.started || self.finished { return Ok(Vec::new()); } - let out = encode_json_sse( + let out = self.encode_chunk(build_openai_chat_finish_chunk( + self.response_id + .as_deref() + .unwrap_or("chatcmpl-local-stream"), + self.model.as_deref().unwrap_or("unknown"), None, - &build_openai_chat_finish_chunk( - self.response_id - .as_deref() - .unwrap_or("chatcmpl-local-stream"), - self.model.as_deref().unwrap_or("unknown"), - None, - ), - )?; + ))?; self.finished = true; let mut bytes = out; bytes.extend(encode_done_sse()); @@ -2035,6 +2174,19 @@ impl OpenAIChatClientEmitter { } impl OpenAIResponsesClientEmitter { + pub(crate) fn set_actual_service_tier(&mut self, value: Option<&str>) { + if value.is_some_and(|value| { + self.actual_service_tier + .as_deref() + .is_some_and(|current| current.eq_ignore_ascii_case(value.trim())) + }) { + return; + } + if let Some(value) = normalize_openai_service_tier(value) { + self.actual_service_tier = Some(value); + } + } + fn response_id(&self) -> &str { self.response_id.as_deref().unwrap_or("resp-local-stream") } @@ -2082,6 +2234,14 @@ impl OpenAIResponsesClientEmitter { { response_object.insert("created_at".to_string(), Value::from(created_at)); } + if let (Some(service_tier), Some(response_object)) = + (self.actual_service_tier.as_ref(), response.as_object_mut()) + { + response_object.insert( + "service_tier".to_string(), + Value::String(service_tier.clone()), + ); + } response } @@ -2682,6 +2842,9 @@ impl OpenAIResponsesClientEmitter { for (output_index, item) in &self.image_generation_items { ordered_output.push((*output_index, item.clone())); } + for (output_index, item) in &self.opaque_output_items { + ordered_output.push((*output_index, item.clone())); + } ordered_output.sort_by_key(|(output_index, _)| *output_index); let mut response = json!({ @@ -2703,6 +2866,12 @@ impl OpenAIResponsesClientEmitter { response_object.insert("created_at".to_string(), Value::from(created_at)); } ensure_modern_openai_responses_response_fields(response_object); + if let Some(service_tier) = self.actual_service_tier.as_ref() { + response_object.insert( + "service_tier".to_string(), + Value::String(service_tier.clone()), + ); + } } response } @@ -2715,6 +2884,62 @@ impl OpenAIResponsesClientEmitter { self.terminal_response(usage, "incomplete", Some(reason)) } + pub(crate) fn finish_with_authoritative_response_event( + &mut self, + response: Value, + event_type: &'static str, + ) -> Result, AiSurfaceFinalizeError> { + if self.finished { + return Ok(Vec::new()); + } + if !matches!( + event_type, + "response.completed" | "response.incomplete" | "response.failed" + ) { + return Err(AiSurfaceFinalizeError::new(format!( + "unsupported authoritative OpenAI Responses terminal event: {event_type}" + ))); + } + let response_object = response.as_object().ok_or_else(|| { + AiSurfaceFinalizeError::new( + "authoritative OpenAI Responses terminal payload must be an object", + ) + })?; + self.set_actual_service_tier(response_object.get("service_tier").and_then(Value::as_str)); + if let Some(response_id) = response_object + .get("id") + .and_then(Value::as_str) + .filter(|value| !value.is_empty()) + { + self.response_id = Some(response_id.to_string()); + } + if let Some(model) = response_object + .get("model") + .and_then(Value::as_str) + .filter(|value| !value.is_empty()) + { + self.model = Some(model.to_string()); + } + if let Some(created_at) = response_object.get("created_at").and_then(Value::as_i64) { + self.created_at = Some(created_at); + } + + let mut out = self.ensure_started()?; + out.extend(self.finish_reasoning_item()?); + out.extend(self.finish_text_item()?); + out.extend(self.finish_tool_items()?); + out.extend(self.finish_tool_result_items()?); + out.extend(self.encode_response_event( + event_type, + json!({ + "type": event_type, + "response": response, + }), + )?); + self.finished = true; + Ok(out) + } + pub fn emit(&mut self, frame: CanonicalStreamFrame) -> Result, AiSurfaceFinalizeError> { self.update_identity(&frame); match frame.event { @@ -2814,6 +3039,50 @@ impl OpenAIResponsesClientEmitter { CanonicalStreamEvent::ImageGenerationCall { index, item } => { self.emit_image_generation_call_item(index, item) } + CanonicalStreamEvent::OpenAiResponsesOutputItem { + output_index, + item, + raw_event, + } => { + let event_type = raw_event + .get("type") + .and_then(Value::as_str) + .filter(|value| { + matches!( + *value, + "response.output_item.added" | "response.output_item.done" + ) + }) + .ok_or_else(|| { + AiSurfaceFinalizeError::new( + "OpenAI Responses output item event must be added or done", + ) + })? + .to_string(); + let item_key = item + .as_object() + .map(OpenAIResponsesProviderState::output_item_key) + .unwrap_or_else(|| item.to_string()); + let materialized_output_index = if let Some(output_index) = output_index { + self.next_output_index = + self.next_output_index.max(output_index.saturating_add(1)); + self.opaque_output_indexes.insert(item_key, output_index); + output_index + } else if let Some(output_index) = self.opaque_output_indexes.get(&item_key) { + *output_index + } else { + let output_index = self.allocate_output_index(); + self.opaque_output_indexes.insert(item_key, output_index); + output_index + }; + let mut out = self.ensure_started()?; + if event_type == "response.output_item.done" { + self.opaque_output_items + .insert(materialized_output_index, item); + } + out.extend(encode_json_sse(Some(event_type.as_str()), &raw_event)?); + Ok(out) + } CanonicalStreamEvent::ToolCallStart { index, call_id, @@ -3362,6 +3631,93 @@ mod tests { ))); } + #[test] + fn openai_responses_provider_state_recognizes_compaction_output_without_index() { + let mut state = OpenAIResponsesProviderState::default(); + let report_context = json!({}); + let raw_event = json!({ + "type": "response.output_item.done", + "item": { + "type": "compaction", + "encrypted_content": "ENCRYPTED_CONTEXT_COMPACTION_SUMMARY" + } + }); + let frames = state + .push_line(&report_context, data_line(raw_event.clone())) + .expect("compaction output item should parse"); + + assert!(frames.iter().any(|frame| matches!( + &frame.event, + CanonicalStreamEvent::OpenAiResponsesOutputItem { + output_index: None, + item, + raw_event: emitted_event, + } if item["type"] == "compaction" && emitted_event == &raw_event + ))); + assert!(!frames + .iter() + .any(|frame| matches!(frame.event, CanonicalStreamEvent::UnknownEvent(_)))); + + let terminal = state + .push_line( + &report_context, + data_line(json!({ + "type": "response.completed", + "response": { + "id": "resp-compact", + "usage": { + "input_tokens": 0, + "output_tokens": 0, + "total_tokens": 0 + } + } + })), + ) + .expect("completed response should parse"); + assert!(terminal + .iter() + .any(|frame| matches!(frame.event, CanonicalStreamEvent::Finish { .. }))); + assert!(!terminal + .iter() + .any(|frame| matches!(frame.event, CanonicalStreamEvent::UnknownEvent(_)))); + } + + #[test] + fn openai_responses_provider_state_rejects_function_call_provenance() { + let mut state = OpenAIResponsesProviderState::default(); + let report_context = json!({}); + let frames = state + .push_line( + &report_context, + data_line(json!({ + "type": "response.output_item.added", + "response_id": "resp_ptc_123", + "output_index": 0, + "item": { + "type": "function_call", + "id": "fc_123", + "call_id": "call_123", + "status": "completed", + "name": "lookup", + "arguments": "{}", + "caller": {"type": "program", "id": "program_123"} + } + })), + ) + .expect("function call event should parse"); + + assert!(frames.iter().any(|frame| matches!( + frame.event, + CanonicalStreamEvent::UnknownEvent(ref payload) + if payload.get("caller").is_some() + ))); + assert!(!frames.iter().any(|frame| matches!( + frame.event, + CanonicalStreamEvent::ToolCallStart { .. } + | CanonicalStreamEvent::ToolCallArgumentsDelta { .. } + ))); + } + #[test] fn openai_responses_provider_state_treats_failed_event_as_terminal() { let mut state = OpenAIResponsesProviderState::default(); @@ -3489,6 +3845,7 @@ mod tests { "input_tokens": 26, "input_tokens_details": { "cached_tokens": 0, + "cache_write_tokens": 6, }, "output_tokens": 137, "output_tokens_details": { @@ -3508,6 +3865,7 @@ mod tests { usage: Some(CanonicalUsage { input_tokens: 26, output_tokens: 137, + cache_creation_tokens: 6, cache_read_tokens: 0, .. }), @@ -4560,6 +4918,7 @@ mod tests { assert!(sse.contains("\"completion_tokens\":2")); assert!(sse.contains("\"completion_tokens_details\":{\"reasoning_tokens\":1}")); assert!(sse.contains("\"cache_write_tokens\":5")); + assert!(!sse.contains("\"cached_creation_tokens\"")); assert!(sse.contains("\"cached_tokens\":4")); assert!(sse.contains("\"total_tokens\":3")); assert!(sse.contains("data: [DONE]\n\n")); @@ -4684,6 +5043,7 @@ mod tests { assert!(sse.contains("\"output_tokens_details\":{\"reasoning_tokens\":1}")); assert!(sse.contains("\"input_tokens_details\"")); assert!(sse.contains("\"cache_write_tokens\":5")); + assert!(!sse.contains("\"cached_creation_tokens\"")); assert!(sse.contains("\"cached_tokens\":4")); } diff --git a/crates/aether-ai-formats/src/formats/openai/image/request.rs b/crates/aether-ai-formats/src/formats/openai/image/request.rs index d68b0ac06..45b6aa827 100644 --- a/crates/aether-ai-formats/src/formats/openai/image/request.rs +++ b/crates/aether-ai-formats/src/formats/openai/image/request.rs @@ -44,7 +44,7 @@ pub struct NormalizedOpenAiImageRequest { images: Vec, tool: Map, image_count: Option, - stream: Option, + max_generation_count: u64, user: Option, } @@ -56,7 +56,7 @@ pub struct OpenAiImageNormalizeOptions { impl Default for OpenAiImageNormalizeOptions { fn default() -> Self { Self { - max_generation_count: 1, + max_generation_count: OPENAI_IMAGE_MAX_GENERATION_COUNT, } } } @@ -70,6 +70,7 @@ impl OpenAiImageNormalizeOptions { } pub const CHATGPT_WEB_IMAGE_MAX_AREA: u64 = 1_500_000; +pub const OPENAI_IMAGE_MAX_GENERATION_COUNT: u64 = 10; #[derive(Debug, Clone, PartialEq, Eq)] pub struct ChatGptWebImageRequestError { @@ -540,7 +541,8 @@ pub fn build_openai_image_provider_request_body(request: &NormalizedOpenAiImageR pub fn build_openai_image_api_provider_request_body( request: &NormalizedOpenAiImageRequest, mapped_model: Option<&str>, -) -> Value { + upstream_is_stream: bool, +) -> Option { let model = mapped_model .map(str::trim) .filter(|value| !value.is_empty()) @@ -557,8 +559,8 @@ pub fn build_openai_image_api_provider_request_body( if let Some(user) = request.user.as_ref() { body.insert("user".to_string(), Value::String(user.clone())); } - if let Some(stream) = request.stream { - body.insert("stream".to_string(), Value::Bool(stream)); + if upstream_is_stream { + body.insert("stream".to_string(), Value::Bool(true)); } for (key, value) in &request.tool { match key.as_str() { @@ -575,14 +577,503 @@ pub fn build_openai_image_api_provider_request_body( body.entry("response_format".to_string()) .or_insert_with(|| response_format.clone()); } - if !request.images.is_empty() { - if request.images.len() == 1 { - body.insert("image".to_string(), request.images[0].clone()); + insert_standard_openai_image_inputs(&mut body, request.images.clone()); + project_openai_image_api_request_body( + &Value::Object(body), + model, + request.operation, + request.max_generation_count, + ) +} + +pub fn build_codex_openai_image_api_provider_request_body( + request: &NormalizedOpenAiImageRequest, + mapped_model: Option<&str>, + upstream_is_stream: bool, +) -> Option { + let body = + build_openai_image_api_provider_request_body(request, mapped_model, upstream_is_stream)?; + project_codex_openai_image_api_request_body(&body, request.operation) +} + +pub fn project_openai_image_api_request_body( + body: &Value, + provider_model: &str, + operation: OpenAiImageOperation, + max_generation_count: u64, +) -> Option { + let mut projected = body.as_object()?.clone(); + let model = provider_model.trim(); + if model.is_empty() { + return None; + } + let model_family = openai_image_model_family(model); + if operation == OpenAiImageOperation::Edit && model_family == OpenAiImageModelFamily::DallE3 { + return None; + } + + let prompt = non_empty_image_string(projected.get("prompt"))?; + let max_prompt_chars = match model_family { + OpenAiImageModelFamily::DallE2 => Some(1_000), + OpenAiImageModelFamily::DallE3 => Some(4_000), + OpenAiImageModelFamily::GptImage + | OpenAiImageModelFamily::GptImageMini + | OpenAiImageModelFamily::GptImage2 => Some(32_000), + OpenAiImageModelFamily::Other => None, + }; + if max_prompt_chars.is_some_and(|limit| prompt.chars().count() > limit) { + return None; + } + + let image_inputs = openai_image_api_inputs(&projected)? + .into_iter() + .cloned() + .collect::>(); + projected.remove("image"); + projected.remove("images"); + let image_input_count = image_inputs.len(); + insert_standard_openai_image_inputs(&mut projected, image_inputs); + match operation { + OpenAiImageOperation::Generate + if image_input_count > 0 || projected.contains_key("mask") => + { + return None; + } + OpenAiImageOperation::Edit if image_input_count == 0 => return None, + OpenAiImageOperation::Edit => { + let max_images = match model_family { + OpenAiImageModelFamily::DallE2 => 1, + OpenAiImageModelFamily::GptImage + | OpenAiImageModelFamily::GptImageMini + | OpenAiImageModelFamily::GptImage2 => 16, + OpenAiImageModelFamily::DallE3 => 0, + OpenAiImageModelFamily::Other => usize::MAX, + }; + if image_input_count > max_images { + return None; + } + } + OpenAiImageOperation::Generate => {} + } + + if let Some(n) = projected.get("n") { + let model_limit = if model_family == OpenAiImageModelFamily::DallE3 { + 1 } else { - body.insert("images".to_string(), Value::Array(request.images.clone())); + OPENAI_IMAGE_MAX_GENERATION_COUNT + }; + let max_generation_count = max_generation_count.max(1).min(model_limit); + let n = n + .as_u64() + .filter(|value| (1..=max_generation_count).contains(value))?; + projected.insert("n".to_string(), Value::Number(Number::from(n))); + } + if let Some(quality) = projected.get("quality") { + let quality = non_empty_image_string(Some(quality))?; + let canonical = normalize_openai_image_quality(quality)?; + let wire = match model_family { + OpenAiImageModelFamily::DallE3 => match canonical { + "medium" => "standard", + "high" => "hd", + _ => return None, + }, + OpenAiImageModelFamily::DallE2 => match canonical { + "medium" => "standard", + _ => return None, + }, + OpenAiImageModelFamily::GptImage + | OpenAiImageModelFamily::GptImageMini + | OpenAiImageModelFamily::GptImage2 + | OpenAiImageModelFamily::Other => canonical, + }; + projected.insert("quality".to_string(), Value::String(wire.to_string())); + } + + let output_format = if let Some(value) = projected.get("output_format") { + let value = non_empty_image_string(Some(value))?.to_ascii_lowercase(); + if !matches!(value.as_str(), "png" | "jpeg" | "webp") + || !matches!( + model_family, + OpenAiImageModelFamily::GptImage + | OpenAiImageModelFamily::GptImageMini + | OpenAiImageModelFamily::GptImage2 + | OpenAiImageModelFamily::Other + ) + { + return None; + } + projected.insert("output_format".to_string(), Value::String(value.clone())); + Some(value) + } else { + None + }; + + if let Some(value) = projected.get("response_format") { + let value = non_empty_image_string(Some(value))?.to_ascii_lowercase(); + if !matches!(value.as_str(), "url" | "b64_json") { + return None; + } + match model_family { + OpenAiImageModelFamily::GptImage + | OpenAiImageModelFamily::GptImageMini + | OpenAiImageModelFamily::GptImage2 => { + if value != "b64_json" { + return None; + } + projected.remove("response_format"); + } + OpenAiImageModelFamily::DallE2 + | OpenAiImageModelFamily::DallE3 + | OpenAiImageModelFamily::Other => { + projected.insert("response_format".to_string(), Value::String(value)); + } } } - Value::Object(body) + + if let Some(value) = projected.get("background") { + let value = non_empty_image_string(Some(value))?.to_ascii_lowercase(); + if !matches!(value.as_str(), "transparent" | "opaque" | "auto") + || matches!( + model_family, + OpenAiImageModelFamily::DallE2 | OpenAiImageModelFamily::DallE3 + ) + || (model_family == OpenAiImageModelFamily::GptImage2 && value == "transparent") + || (value == "transparent" + && !matches!(output_format.as_deref().unwrap_or("png"), "png" | "webp")) + { + return None; + } + projected.insert("background".to_string(), Value::String(value)); + } + + if let Some(value) = projected.get("moderation") { + let value = non_empty_image_string(Some(value))?.to_ascii_lowercase(); + if operation != OpenAiImageOperation::Generate + || !matches!(value.as_str(), "low" | "auto") + || matches!( + model_family, + OpenAiImageModelFamily::DallE2 | OpenAiImageModelFamily::DallE3 + ) + { + return None; + } + projected.insert("moderation".to_string(), Value::String(value)); + } + + if let Some(value) = projected.get("input_fidelity") { + let value = non_empty_image_string(Some(value))?.to_ascii_lowercase(); + if operation != OpenAiImageOperation::Edit + || !matches!(value.as_str(), "low" | "high") + || matches!( + model_family, + OpenAiImageModelFamily::DallE2 + | OpenAiImageModelFamily::DallE3 + | OpenAiImageModelFamily::GptImageMini + ) + { + return None; + } + projected.insert("input_fidelity".to_string(), Value::String(value)); + } + + if let Some(value) = projected.get("output_compression") { + let value = value.as_u64().filter(|value| *value <= 100)?; + if !matches!(output_format.as_deref(), Some("jpeg" | "webp")) + || matches!( + model_family, + OpenAiImageModelFamily::DallE2 | OpenAiImageModelFamily::DallE3 + ) + { + return None; + } + projected.insert( + "output_compression".to_string(), + Value::Number(Number::from(value)), + ); + } + + let stream = match projected.get("stream") { + Some(value) => Some(value.as_bool()?), + None => None, + }; + if let Some(stream) = stream { + match model_family { + OpenAiImageModelFamily::DallE2 + | OpenAiImageModelFamily::DallE3 + | OpenAiImageModelFamily::GptImage2 => { + if stream { + return None; + } + projected.remove("stream"); + } + OpenAiImageModelFamily::GptImage + | OpenAiImageModelFamily::GptImageMini + | OpenAiImageModelFamily::Other => {} + } + } + + if let Some(value) = projected.get("partial_images") { + let value = value.as_u64().filter(|value| *value <= 3)?; + if stream != Some(true) + || matches!( + model_family, + OpenAiImageModelFamily::DallE2 + | OpenAiImageModelFamily::DallE3 + | OpenAiImageModelFamily::GptImage2 + ) + { + return None; + } + projected.insert( + "partial_images".to_string(), + Value::Number(Number::from(value)), + ); + } + + if let Some(value) = projected.get("style") { + let value = non_empty_image_string(Some(value))?.to_ascii_lowercase(); + if operation != OpenAiImageOperation::Generate + || !matches!(value.as_str(), "vivid" | "natural") + || !matches!( + model_family, + OpenAiImageModelFamily::DallE3 | OpenAiImageModelFamily::Other + ) + { + return None; + } + projected.insert("style".to_string(), Value::String(value)); + } + + if let Some(value) = projected.get("size") { + let value = non_empty_image_string(Some(value))?.to_ascii_lowercase(); + if !openai_image_size_supported(model_family, &value) { + return None; + } + projected.insert("size".to_string(), Value::String(value)); + } + + Some(Value::Object(projected)) +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum OpenAiImageModelFamily { + GptImage, + GptImageMini, + GptImage2, + DallE3, + DallE2, + Other, +} + +fn openai_image_model_family(model: &str) -> OpenAiImageModelFamily { + let normalized = model.trim().to_ascii_lowercase(); + let model = normalized.rsplit('/').next().unwrap_or_default(); + if model == "gpt-image-2" || model.starts_with("gpt-image-2-") { + OpenAiImageModelFamily::GptImage2 + } else if model == "gpt-image-1-mini" || model.starts_with("gpt-image-1-mini-") { + OpenAiImageModelFamily::GptImageMini + } else if model.starts_with("gpt-image-") || model == "chatgpt-image-latest" { + OpenAiImageModelFamily::GptImage + } else if model.starts_with("dall-e-3") { + OpenAiImageModelFamily::DallE3 + } else if model.starts_with("dall-e-2") { + OpenAiImageModelFamily::DallE2 + } else { + OpenAiImageModelFamily::Other + } +} + +pub(crate) fn insert_standard_openai_image_inputs( + object: &mut Map, + images: Vec, +) { + if images.is_empty() { + return; + } + let image = if images.len() == 1 { + images.into_iter().next().expect("one image should exist") + } else { + Value::Array(images) + }; + object.insert("image".to_string(), image); +} + +fn openai_image_api_inputs(object: &Map) -> Option> { + match (object.get("image"), object.get("images")) { + (Some(_), Some(_)) => None, + (Some(Value::Array(images)), None) | (None, Some(Value::Array(images))) => { + Some(images.iter().collect()) + } + (Some(image), None) => Some(vec![image]), + (None, Some(_)) => None, + (None, None) => Some(Vec::new()), + } +} + +fn openai_image_size_supported(family: OpenAiImageModelFamily, size: &str) -> bool { + match family { + OpenAiImageModelFamily::DallE2 => { + matches!(size, "256x256" | "512x512" | "1024x1024") + } + OpenAiImageModelFamily::DallE3 => { + matches!(size, "1024x1024" | "1792x1024" | "1024x1792") + } + OpenAiImageModelFamily::GptImage | OpenAiImageModelFamily::GptImageMini => { + matches!(size, "auto" | "1024x1024" | "1536x1024" | "1024x1536") + } + OpenAiImageModelFamily::GptImage2 => { + size == "auto" || gpt_image_2_resolution_supported(size) + } + OpenAiImageModelFamily::Other => !size.is_empty(), + } +} + +fn gpt_image_2_resolution_supported(size: &str) -> bool { + let mut dimensions = size.split('x'); + let Some(width) = dimensions + .next() + .and_then(|value| value.parse::().ok()) + else { + return false; + }; + let Some(height) = dimensions + .next() + .and_then(|value| value.parse::().ok()) + else { + return false; + }; + if dimensions.next().is_some() + || width == 0 + || height == 0 + || width % 16 != 0 + || height % 16 != 0 + { + return false; + } + let shorter = width.min(height); + let longer = width.max(height); + longer <= shorter.saturating_mul(3) + && longer <= 3_840 + && width.saturating_mul(height) <= 3_840 * 2_160 +} + +pub fn project_codex_openai_image_api_request_body( + body: &Value, + operation: OpenAiImageOperation, +) -> Option { + let model = body.get("model").and_then(Value::as_str)?; + let projected_body = project_openai_image_api_request_body( + body, + model, + operation, + OPENAI_IMAGE_MAX_GENERATION_COUNT, + )?; + let object = projected_body.as_object()?; + if object.keys().any(|key| { + !matches!( + key.as_str(), + "model" + | "prompt" + | "background" + | "n" + | "quality" + | "size" + | "image" + | "response_format" + | "stream" + ) + }) { + return None; + } + + let model = non_empty_image_string(object.get("model"))?; + let prompt = non_empty_image_string(object.get("prompt"))?; + if object + .get("stream") + .is_some_and(|value| !value.is_boolean()) + { + return None; + } + if object + .get("response_format") + .is_some_and(|value| value.as_str() != Some("b64_json")) + { + return None; + } + + let mut projected = Map::new(); + if operation == OpenAiImageOperation::Edit { + let images = collect_codex_openai_image_urls(object)?; + if images.is_empty() || images.len() > 5 { + return None; + } + projected.insert("images".to_string(), Value::Array(images)); + } else if object.contains_key("image") || object.contains_key("images") { + return None; + } + projected.insert("prompt".to_string(), Value::String(prompt.to_string())); + if let Some(background) = + optional_codex_image_enum(object.get("background"), &["transparent", "opaque", "auto"])? + { + projected.insert("background".to_string(), Value::String(background)); + } + projected.insert("model".to_string(), Value::String(model.to_string())); + if let Some(n) = object.get("n") { + let n = n + .as_u64() + .filter(|value| (1..=OPENAI_IMAGE_MAX_GENERATION_COUNT).contains(value))?; + projected.insert("n".to_string(), Value::Number(Number::from(n))); + } + if let Some(quality) = + optional_codex_image_enum(object.get("quality"), &["low", "medium", "high", "auto"])? + { + projected.insert("quality".to_string(), Value::String(quality)); + } + if let Some(size) = object.get("size") { + projected.insert( + "size".to_string(), + Value::String(non_empty_image_string(Some(size))?.to_string()), + ); + } + Some(Value::Object(projected)) +} + +fn non_empty_image_string(value: Option<&Value>) -> Option<&str> { + value? + .as_str() + .map(str::trim) + .filter(|value| !value.is_empty()) +} + +fn optional_codex_image_enum(value: Option<&Value>, allowed: &[&str]) -> Option> { + let Some(value) = value else { + return Some(None); + }; + let value = non_empty_image_string(Some(value))?.to_ascii_lowercase(); + allowed + .iter() + .any(|allowed| value == *allowed) + .then_some(Some(value)) +} + +fn collect_codex_openai_image_urls(object: &Map) -> Option> { + openai_image_api_inputs(object)? + .into_iter() + .map(|image| { + let image = image.as_object()?; + if image + .keys() + .any(|key| !matches!(key.as_str(), "type" | "image_url")) + || image + .get("type") + .is_some_and(|value| value.as_str() != Some("input_image")) + { + return None; + } + let image_url = non_empty_image_string(image.get("image_url"))?; + Some(json!({ "image_url": image_url })) + }) + .collect() } fn normalize_openai_image_json_request( @@ -591,18 +1082,8 @@ fn normalize_openai_image_json_request( options: OpenAiImageNormalizeOptions, ) -> Option { let object = body_json.as_object()?; - if object - .get("style") - .and_then(Value::as_str) - .map(str::trim) - .is_some_and(|value| !value.is_empty()) - { - return None; - } let image_count = object.get("n").and_then(image_request_count); - if image_count - .is_some_and(|value| value == 0 || value > max_count_for_operation(operation, options)) - { + if image_count.is_some_and(|value| value == 0 || value > max_generation_count(options)) { return None; } let requested_model = @@ -613,7 +1094,12 @@ fn normalize_openai_image_json_request( let output_format = normalize_output_format(object.get("output_format").and_then(Value::as_str))?; let partial_images = normalize_partial_images(object.get("partial_images"))?; - let stream = object.get("stream").and_then(value_as_bool); + if object + .get("stream") + .is_some_and(|value| value_as_bool(value).is_none()) + { + return None; + } let user = object .get("user") .and_then(Value::as_str) @@ -644,7 +1130,7 @@ fn normalize_openai_image_json_request( images, tool, image_count, - stream, + max_generation_count: max_generation_count(options), user, summary_json: build_image_request_summary_json( operation, @@ -665,14 +1151,9 @@ fn normalize_openai_image_multipart_request( let requested_model = normalize_requested_image_model( find_multipart_text_field(&multipart_fields, "model").as_deref(), ); - if find_multipart_text_field(&multipart_fields, "style").is_some() { - return None; - } let image_count = find_multipart_text_field(&multipart_fields, "n") .and_then(|value| value.trim().parse::().ok()); - if image_count - .is_some_and(|value| value == 0 || value > max_count_for_operation(operation, options)) - { + if image_count.is_some_and(|value| value == 0 || value > max_generation_count(options)) { return None; } let prompt = normalize_prompt( @@ -693,9 +1174,12 @@ fn normalize_openai_image_multipart_request( .map(Value::String) .as_ref(), )?; - let stream = find_multipart_text_field(&multipart_fields, "stream") + if find_multipart_text_field(&multipart_fields, "stream") .as_deref() - .and_then(parse_bool_string); + .is_some_and(|value| parse_bool_string(value).is_none()) + { + return None; + } let user = find_multipart_text_field(&multipart_fields, "user") .map(|value| value.trim().to_string()) .filter(|value| !value.is_empty()); @@ -727,6 +1211,7 @@ fn normalize_openai_image_multipart_request( | "moderation" | "input_fidelity" | "partial_images" + | "style" ) { raw_tool_values.insert( name, @@ -748,7 +1233,7 @@ fn normalize_openai_image_multipart_request( images, tool, image_count, - stream, + max_generation_count: max_generation_count(options), user, summary_json: build_image_request_summary_json( operation, @@ -865,6 +1350,7 @@ fn build_tool_options_from_json( "moderation", "input_fidelity", "partial_images", + "style", ] { if let Some(value) = object.get(key) { raw_values.insert(key.to_string(), value.clone()); @@ -907,7 +1393,7 @@ fn build_tool_options( ); for (key, value) in raw_values { let normalized = match key.as_str() { - "size" | "background" | "moderation" | "input_fidelity" => { + "size" | "background" | "moderation" | "input_fidelity" | "style" => { normalize_non_empty_string_value(&value) } "output_format" => normalize_output_format_value(&value), @@ -934,20 +1420,22 @@ fn normalize_non_empty_string_value(value: &Value) -> Option { .map(|value| Value::String(value.to_string())) } +pub fn normalize_openai_image_quality(value: &str) -> Option<&'static str> { + match value.trim().to_ascii_lowercase().as_str() { + "low" => Some("low"), + "medium" | "standard" => Some("medium"), + "high" | "hd" => Some("high"), + "auto" => Some("auto"), + _ => None, + } +} + fn normalize_quality_value(value: &Value) -> Option { let quality = value .as_str() .map(str::trim) - .filter(|value| !value.is_empty())? - .to_ascii_lowercase(); - let normalized = match quality.as_str() { - "low" => "low", - "medium" => "medium", - "high" => "high", - "standard" => "medium", - "hd" => "high", - _ => return None, - }; + .filter(|value| !value.is_empty())?; + let normalized = normalize_openai_image_quality(quality)?; Some(Value::String(normalized.to_string())) } @@ -988,14 +1476,8 @@ fn image_request_count(value: &Value) -> Option { }) } -fn max_count_for_operation( - operation: OpenAiImageOperation, - options: OpenAiImageNormalizeOptions, -) -> u64 { - match operation { - OpenAiImageOperation::Generate => options.max_generation_count.max(1), - OpenAiImageOperation::Edit => 1, - } +fn max_generation_count(options: OpenAiImageNormalizeOptions) -> u64 { + options.max_generation_count.max(1) } fn normalize_image_value(value: &Value) -> Vec { @@ -1206,15 +1688,14 @@ mod tests { use serde_json::json; use super::{ - build_chatgpt_web_image_request_body, build_openai_image_api_provider_request_body, - build_openai_image_provider_request_body, is_openai_image_stream_request, + build_chatgpt_web_image_request_body, build_codex_openai_image_api_provider_request_body, + build_openai_image_api_provider_request_body, build_openai_image_provider_request_body, + is_openai_image_stream_request, normalize_openai_image_quality, normalize_openai_image_request, normalize_openai_image_request_with_options, - openai_image_operation_from_path, OpenAiImageNormalizeOptions, OpenAiImageOperation, + openai_image_operation_from_path, project_codex_openai_image_api_request_body, + project_openai_image_api_request_body, OpenAiImageNormalizeOptions, OpenAiImageOperation, }; use crate::formats::openai::image::spec::{resolve_stream_spec, resolve_sync_spec}; - use crate::formats::openai::responses::codex::{ - apply_codex_openai_responses_special_body_edits, CODEX_OPENAI_IMAGE_INTERNAL_MODEL, - }; fn request_parts(path: &str, content_type: Option<&str>) -> http::request::Parts { let mut builder = Request::builder().method(Method::POST).uri(path); @@ -1228,6 +1709,21 @@ mod tests { .0 } + #[test] + fn image_quality_uses_one_normalized_contract() { + for (input, expected) in [ + ("auto", Some("auto")), + ("low", Some("low")), + ("medium", Some("medium")), + ("standard", Some("medium")), + ("high", Some("high")), + ("hd", Some("high")), + ("unsupported", None), + ] { + assert_eq!(normalize_openai_image_quality(input), expected); + } + } + #[test] fn resolves_openai_image_sync_spec() { let spec = resolve_sync_spec("openai_image_sync").expect("spec"); @@ -1411,14 +1907,25 @@ mod tests { } #[test] - fn normalize_generate_json_request_rejects_multi_image_count_by_default() { + fn normalize_generate_json_request_uses_the_openai_image_count_range_by_default() { let parts = request_parts("/v1/images/generations", Some("application/json")); + let request = normalize_openai_image_request( + &parts, + &json!({ + "model": "gpt-image-2", + "prompt": "generate image", + "n": 10 + }), + None, + ) + .expect("OpenAI image generation should allow n up to ten"); + assert_eq!(request.image_count, Some(10)); assert!(normalize_openai_image_request( &parts, &json!({ "model": "gpt-image-2", "prompt": "generate image", - "n": 2 + "n": 11 }), None, ) @@ -1426,14 +1933,14 @@ mod tests { } #[test] - fn normalize_edit_request_rejects_multi_image_count_even_with_generation_override() { + fn normalize_edit_request_uses_the_configured_image_count_range() { let parts = request_parts("/v1/images/edits", Some("application/json")); - assert!(normalize_openai_image_request_with_options( + let request = normalize_openai_image_request_with_options( &parts, &json!({ "model": "grok-imagine-image-edit", "prompt": "edit image", - "n": 2, + "n": 4, "image": { "b64_json": "aGVsbG8=", "mime_type": "image/png" @@ -1442,110 +1949,305 @@ mod tests { None, OpenAiImageNormalizeOptions::with_max_generation_count(4), ) - .is_none()); + .expect("image edits should use the configured output count range"); + assert_eq!(request.image_count, Some(4)); } #[test] - fn build_generate_request_defaults_codex_image_tool_and_tool_choice() { + fn builds_codex_image_generation_with_the_typed_images_contract() { let parts = request_parts("/v1/images/generations", Some("application/json")); let request = normalize_openai_image_request( &parts, &json!({ "model": "gpt-image-2", - "prompt": "generate image" + "prompt": "generate image", + "background": "auto", + "quality": "auto", + "size": "auto", + "stream": true, + "response_format": "b64_json" }), None, ) .expect("generation request should normalize"); - assert!(request.tool.get("size").is_none()); - assert!(request.tool.get("quality").is_none()); - assert!(request.tool.get("background").is_none()); - assert!(request.tool.get("output_format").is_none()); assert_eq!( - request.tool.get("action").and_then(|value| value.as_str()), - Some("generate") + build_codex_openai_image_api_provider_request_body( + &request, + Some("gpt-image-2"), + false, + ), + Some(json!({ + "prompt": "generate image", + "background": "auto", + "model": "gpt-image-2", + "quality": "auto", + "size": "auto" + })) ); + } - let mut provider_request_body = build_openai_image_provider_request_body(&request); - assert!(provider_request_body.get("model").is_none()); - assert!(provider_request_body.get("tool_choice").is_none()); - assert!(provider_request_body.get("stream").is_none()); - apply_codex_openai_responses_special_body_edits( - &mut provider_request_body, - "codex", - "openai:image", + #[test] + fn builds_codex_image_edit_with_plural_image_urls() { + let parts = request_parts("/v1/images/edits", Some("application/json")); + let request = normalize_openai_image_request( + &parts, + &json!({ + "model": "gpt-image-2", + "prompt": "add a red hat", + "images": [ + {"image_url": "data:image/png;base64,Zm9v"}, + {"image_url": "https://example.test/reference.png"} + ], + "background": "auto", + "quality": "auto", + "size": "auto" + }), None, - None, - ); + ) + .expect("edit request should normalize"); assert_eq!( - provider_request_body - .get("tools") - .and_then(|value| value.get(0)) - .and_then(|value| value.get("type")) - .and_then(|value| value.as_str()), - Some("image_generation") + build_codex_openai_image_api_provider_request_body(&request, None, false), + Some(json!({ + "images": [ + {"image_url": "data:image/png;base64,Zm9v"}, + {"image_url": "https://example.test/reference.png"} + ], + "prompt": "add a red hat", + "background": "auto", + "model": "gpt-image-2", + "quality": "auto", + "size": "auto" + })) ); + } + + #[test] + fn rejects_fields_outside_the_codex_images_contract() { + let parts = request_parts("/v1/images/edits", Some("application/json")); + for unsupported in [ + json!({"mask": "data:image/png;base64,bWFzaw=="}), + json!({"input_fidelity": "high"}), + json!({"output_format": "png"}), + json!({"partial_images": 1}), + json!({"response_format": "url"}), + json!({"user": "user-123"}), + ] { + let mut body = json!({ + "model": "gpt-image-2", + "prompt": "edit image", + "image": {"image_url": "data:image/png;base64,aW1hZ2U="} + }); + body.as_object_mut() + .expect("body object") + .extend(unsupported.as_object().expect("unsupported object").clone()); + let request = normalize_openai_image_request(&parts, &body, None) + .expect("request should normalize before provider projection"); + assert!( + build_codex_openai_image_api_provider_request_body(&request, None, false).is_none(), + "unsupported body should be rejected: {body}" + ); + } + } + + #[test] + fn codex_image_projection_enforces_the_openai_output_count_range() { + let valid = json!({ + "model": "gpt-image-2", + "prompt": "generate image", + "n": 10 + }); assert_eq!( - provider_request_body - .get("tools") - .and_then(|value| value.get(0)) - .and_then(|value| value.get("action")) - .and_then(|value| value.as_str()), - Some("generate") - ); - assert_eq!( - provider_request_body - .get("tools") - .and_then(|value| value.get(0)) - .and_then(|value| value.get("size")) - .and_then(|value| value.as_str()), - Some("1024x1024") - ); - assert_eq!( - provider_request_body - .get("tools") - .and_then(|value| value.get(0)) - .and_then(|value| value.get("quality")) - .and_then(|value| value.as_str()), - Some("high") - ); - assert_eq!( - provider_request_body - .get("tools") - .and_then(|value| value.get(0)) - .and_then(|value| value.get("background")) - .and_then(|value| value.as_str()), - Some("auto") - ); - assert_eq!( - provider_request_body - .get("tools") - .and_then(|value| value.get(0)) - .and_then(|value| value.get("output_format")) - .and_then(|value| value.as_str()), - Some("png") - ); - assert_eq!( - provider_request_body - .get("model") - .and_then(|value| value.as_str()), - Some(CODEX_OPENAI_IMAGE_INTERNAL_MODEL) - ); - assert_eq!( - provider_request_body - .get("stream") - .and_then(|value| value.as_bool()), - Some(true) - ); - assert_eq!( - provider_request_body - .get("tool_choice") - .and_then(|value| value.get("type")) - .and_then(|value| value.as_str()), - Some("image_generation") + project_codex_openai_image_api_request_body(&valid, OpenAiImageOperation::Generate) + .and_then(|body| body.get("n").cloned()), + Some(json!(10)) ); + + for count in [0, 11] { + let body = json!({ + "model": "gpt-image-2", + "prompt": "generate image", + "n": count + }); + assert!(project_codex_openai_image_api_request_body( + &body, + OpenAiImageOperation::Generate + ) + .is_none()); + } + } + + #[test] + fn image_api_projection_renders_quality_for_the_mapped_model_family() { + for (model, quality, expected) in [ + ("gpt-image-2", "standard", Some("medium")), + ("gpt-image-2", "hd", Some("high")), + ("dall-e-3", "medium", Some("standard")), + ("dall-e-3", "high", Some("hd")), + ("dall-e-3", "low", None), + ("dall-e-2", "standard", Some("standard")), + ("dall-e-2", "high", None), + ] { + let body = json!({ + "model": model, + "prompt": "generate image", + "quality": quality + }); + let projected = project_openai_image_api_request_body( + &body, + model, + OpenAiImageOperation::Generate, + 10, + ); + assert_eq!( + projected + .as_ref() + .and_then(|body| body.get("quality")) + .and_then(|quality| quality.as_str()), + expected, + "model={model}, quality={quality}" + ); + } + + let dall_e_3_multi = json!({ + "model": "dall-e-3", + "prompt": "generate image", + "n": 2 + }); + assert!(project_openai_image_api_request_body( + &dall_e_3_multi, + "dall-e-3", + OpenAiImageOperation::Generate, + 1, + ) + .is_none()); + } + + #[test] + fn image_api_projection_enforces_model_field_profiles() { + let gpt_image_2 = project_openai_image_api_request_body( + &json!({ + "model": "gpt-image-2", + "prompt": "generate image", + "background": "opaque", + "size": "1536x864", + "response_format": "b64_json" + }), + "gpt-image-2", + OpenAiImageOperation::Generate, + 10, + ) + .expect("GPT Image 2 should accept its synchronous generation fields"); + assert!(gpt_image_2.get("response_format").is_none()); + assert_eq!(gpt_image_2["size"], "1536x864"); + + for invalid in [ + json!({ + "model": "gpt-image-2", + "prompt": "generate image", + "background": "transparent" + }), + json!({ + "model": "gpt-image-2", + "prompt": "generate image", + "stream": true + }), + json!({ + "model": "gpt-image-2", + "prompt": "generate image", + "size": "1537x864" + }), + ] { + assert!(project_openai_image_api_request_body( + &invalid, + "gpt-image-2", + OpenAiImageOperation::Generate, + 10, + ) + .is_none()); + } + + let dall_e_3 = project_openai_image_api_request_body( + &json!({ + "model": "dall-e-3", + "prompt": "generate image", + "style": "natural", + "quality": "high" + }), + "dall-e-3", + OpenAiImageOperation::Generate, + 10, + ) + .expect("DALL-E 3 should accept style and render its quality vocabulary"); + assert_eq!(dall_e_3["style"], "natural"); + assert_eq!(dall_e_3["quality"], "hd"); + assert!(project_openai_image_api_request_body( + &json!({ + "model": "dall-e-3", + "prompt": "generate image", + "output_format": "png" + }), + "dall-e-3", + OpenAiImageOperation::Generate, + 10, + ) + .is_none()); + + let streaming = project_openai_image_api_request_body( + &json!({ + "model": "gpt-image-1.5", + "prompt": "generate image", + "stream": true, + "partial_images": 2, + "output_format": "webp", + "output_compression": 80 + }), + "gpt-image-1.5", + OpenAiImageOperation::Generate, + 10, + ) + .expect("stream-capable GPT Image models should accept partial images"); + assert_eq!(streaming["partial_images"], 2); + + let edit = json!({ + "model": "gpt-image-2", + "prompt": "edit image", + "image": {"type": "input_image", "image_url": "data:image/png;base64,aA=="}, + "input_fidelity": "high" + }); + assert!(project_openai_image_api_request_body( + &edit, + "gpt-image-2", + OpenAiImageOperation::Edit, + 10, + ) + .is_some()); + assert!(project_openai_image_api_request_body( + &edit, + "gpt-image-1-mini", + OpenAiImageOperation::Edit, + 10, + ) + .is_none()); + } + + #[test] + fn image_api_builder_preserves_dall_e_3_style() { + let parts = request_parts("/v1/images/generations", Some("application/json")); + let request = normalize_openai_image_request( + &parts, + &json!({ + "model": "dall-e-3", + "prompt": "generate image", + "style": "vivid" + }), + None, + ) + .expect("DALL-E 3 style should normalize"); + let body = build_openai_image_api_provider_request_body(&request, Some("dall-e-3"), false) + .expect("DALL-E 3 style should project"); + assert_eq!(body["style"], "vivid"); } #[test] @@ -1565,8 +2267,12 @@ mod tests { ) .expect("generation request should normalize"); - let provider_request_body = - build_openai_image_api_provider_request_body(&request, Some("mapped-image-model")); + let provider_request_body = build_openai_image_api_provider_request_body( + &request, + Some("mapped-image-model"), + true, + ) + .expect("Images API body should satisfy the mapped model contract"); assert_eq!(provider_request_body["model"], "mapped-image-model"); assert_eq!(provider_request_body["prompt"], "draw a cat"); @@ -1596,8 +2302,12 @@ mod tests { ) .expect("edit request should normalize"); - let provider_request_body = - build_openai_image_api_provider_request_body(&request, Some("mapped-edit-model")); + let provider_request_body = build_openai_image_api_provider_request_body( + &request, + Some("mapped-edit-model"), + false, + ) + .expect("Images API edit body should satisfy the mapped model contract"); assert_eq!(provider_request_body["model"], "mapped-edit-model"); assert_eq!(provider_request_body["prompt"], "replace the background"); @@ -1618,6 +2328,34 @@ mod tests { assert!(provider_request_body.get("action").is_none()); } + #[test] + fn build_image_api_provider_edit_request_uses_one_standard_image_field() { + let parts = request_parts("/v1/images/edits", Some("application/json")); + let request = normalize_openai_image_request( + &parts, + &json!({ + "model": "gpt-image-2", + "prompt": "combine the references", + "images": [ + {"image_url": "data:image/png;base64,Zm9v"}, + {"image_url": "https://example.test/reference.png"} + ] + }), + None, + ) + .expect("multi-image edit request should normalize"); + + let provider_request_body = + build_openai_image_api_provider_request_body(&request, None, false) + .expect("standard Images edit body should project"); + + assert_eq!( + provider_request_body["image"].as_array().map(Vec::len), + Some(2) + ); + assert!(provider_request_body.get("images").is_none()); + } + #[test] fn chatgpt_web_accepts_1k_tier_and_1024_size() { let parts = request_parts("/v1/images/generations", Some("application/json")); diff --git a/crates/aether-ai-formats/src/formats/openai/image/stream.rs b/crates/aether-ai-formats/src/formats/openai/image/stream.rs index b20dc87fe..93eca6b93 100644 --- a/crates/aether-ai-formats/src/formats/openai/image/stream.rs +++ b/crates/aether-ai-formats/src/formats/openai/image/stream.rs @@ -742,6 +742,7 @@ impl OpenAiImageStreamTerminalState { .model .clone() .or_else(|| image_bridge_model(Some(report_context))), + provider_actual_service_tier: None, observed_finish: self.observed_finish, unknown_event_count: 0, parser_error: self.parser_error.clone(), @@ -856,6 +857,7 @@ fn openai_image_usage_to_standardized_usage(value: &Value) -> Option) -> Option<(u64, u64, u6 details .get("cache_write_tokens") .or_else(|| details.get("cached_creation_tokens")) + .or_else(|| details.get("cache_creation_tokens")) }) .and_then(Value::as_u64) }) diff --git a/crates/aether-ai-formats/src/formats/openai/mod.rs b/crates/aether-ai-formats/src/formats/openai/mod.rs index b8cdcc78c..ba6e87780 100644 --- a/crates/aether-ai-formats/src/formats/openai/mod.rs +++ b/crates/aether-ai-formats/src/formats/openai/mod.rs @@ -1,7 +1,11 @@ pub mod chat; pub mod embedding; pub mod image; +pub mod prompt_cache; +pub mod reasoning; +pub mod request_contract; pub mod rerank; pub mod responses; +pub mod search; pub mod shared; pub mod video; diff --git a/crates/aether-ai-formats/src/formats/openai/prompt_cache.rs b/crates/aether-ai-formats/src/formats/openai/prompt_cache.rs new file mode 100644 index 000000000..7c24e3d0f --- /dev/null +++ b/crates/aether-ai-formats/src/formats/openai/prompt_cache.rs @@ -0,0 +1,663 @@ +use serde_json::Value; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum OpenAiPromptCacheViolationKind { + InvalidType, + InvalidEnum, + UnsupportedForModel, + UnsupportedContentBlock, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct OpenAiPromptCacheContractViolation { + pub kind: OpenAiPromptCacheViolationKind, + pub field: String, + pub value: Option, + pub reason: String, +} + +pub fn validate_openai_prompt_cache_request( + source_api_format: &str, + provider_model: &str, + body: &Value, +) -> Result<(), OpenAiPromptCacheContractViolation> { + let source_model = body + .get("model") + .and_then(Value::as_str) + .unwrap_or_default(); + validate_openai_prompt_cache_request_with_source_model( + source_api_format, + provider_model, + source_model, + body, + ) +} + +pub fn resolve_openai_prompt_cache_ttl_minutes( + provider_api_format: &str, + provider_model: &str, + source_model: &str, + body: &Value, +) -> Option { + OpenAiPromptCacheApi::parse(provider_api_format)?; + let request = body.as_object()?; + let explicit_ttl = request + .get("prompt_cache_options") + .and_then(Value::as_object) + .and_then(|options| options.get("ttl")) + .and_then(Value::as_str); + if explicit_ttl == Some("30m") { + return Some(30); + } + + let capability_model = + crate::formats::shared::model_directives::openai_model_capability_identity( + provider_model, + source_model, + ); + crate::openai_model_supports_prompt_cache_options(&capability_model).then_some(30) +} + +pub(crate) fn validate_openai_prompt_cache_request_with_source_model( + source_api_format: &str, + provider_model: &str, + source_model: &str, + body: &Value, +) -> Result<(), OpenAiPromptCacheContractViolation> { + let Some(api) = OpenAiPromptCacheApi::parse(source_api_format) else { + return Ok(()); + }; + let Some(request) = body.as_object() else { + return Ok(()); + }; + let capability_model = + crate::formats::shared::model_directives::openai_model_capability_identity( + provider_model, + source_model, + ); + let supports_prompt_cache_options = + crate::formats::shared::model_directives::openai_model_capability_is_opaque( + provider_model, + source_model, + ) || crate::openai_model_supports_prompt_cache_options(&capability_model); + + if let Some(options) = request + .get("prompt_cache_options") + .filter(|value| !value.is_null()) + { + validate_prompt_cache_options(options, supports_prompt_cache_options)?; + } + if let Some(retention) = request + .get("prompt_cache_retention") + .filter(|value| !value.is_null()) + { + validate_prompt_cache_retention(retention, &capability_model)?; + } + match api { + OpenAiPromptCacheApi::Chat => { + validate_chat_prompt_cache_breakpoints(request, supports_prompt_cache_options) + } + OpenAiPromptCacheApi::Responses => { + validate_responses_prompt_cache_breakpoints(request, supports_prompt_cache_options) + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum OpenAiPromptCacheApi { + Chat, + Responses, +} + +impl OpenAiPromptCacheApi { + fn parse(api_format: &str) -> Option { + match crate::normalize_api_format_alias(api_format).as_str() { + "openai:chat" => Some(Self::Chat), + "openai:responses" | "openai:responses:compact" => Some(Self::Responses), + _ => None, + } + } +} + +fn normalize_provider_model(model: &str) -> &str { + model.trim().rsplit('/').next().unwrap_or_default() +} + +fn validate_prompt_cache_options( + value: &Value, + supported_for_model: bool, +) -> Result<(), OpenAiPromptCacheContractViolation> { + if !supported_for_model { + return Err(unsupported_for_model( + "prompt_cache_options", + "provider model does not support prompt_cache_options", + )); + } + let Some(options) = value.as_object() else { + return Err(invalid_type( + "prompt_cache_options", + value, + "prompt_cache_options must be an object", + )); + }; + if let Some(mode) = options.get("mode") { + let Some(raw) = mode.as_str() else { + return Err(invalid_type( + "prompt_cache_options.mode", + mode, + "prompt_cache_options.mode must be a string", + )); + }; + if !matches!(raw, "implicit" | "explicit") { + return Err(invalid_enum( + "prompt_cache_options.mode", + raw, + "prompt_cache_options.mode supports implicit or explicit", + )); + } + } + if let Some(ttl) = options.get("ttl") { + let Some(raw) = ttl.as_str() else { + return Err(invalid_type( + "prompt_cache_options.ttl", + ttl, + "prompt_cache_options.ttl must be a string", + )); + }; + if raw != "30m" { + return Err(invalid_enum( + "prompt_cache_options.ttl", + raw, + "prompt_cache_options.ttl supports 30m", + )); + } + } + Ok(()) +} + +fn validate_prompt_cache_retention( + value: &Value, + provider_model: &str, +) -> Result<(), OpenAiPromptCacheContractViolation> { + if crate::openai_model_supports_prompt_cache_options(provider_model) { + return Err(unsupported_for_model( + "prompt_cache_retention", + "provider model uses prompt_cache_options.ttl", + )); + } + let Some(raw) = value.as_str() else { + return Err(invalid_type( + "prompt_cache_retention", + value, + "prompt_cache_retention must be a string", + )); + }; + if super::shared::OpenAiPromptCacheRetention::parse(raw).is_none() { + return Err(invalid_enum( + "prompt_cache_retention", + raw, + "prompt_cache_retention supports in_memory or 24h", + )); + } + if gpt_5_5_retention_is_24h_only(provider_model) && raw != "24h" { + return Err(invalid_enum( + "prompt_cache_retention", + raw, + "GPT-5.5 family models support 24h retention", + )); + } + Ok(()) +} + +fn gpt_5_5_retention_is_24h_only(model: &str) -> bool { + let normalized = normalize_provider_model(model) + .to_ascii_lowercase() + .replace('_', "-"); + matches!(normalized.as_str(), "gpt-5.5" | "gpt-5.5-pro") +} + +fn validate_chat_prompt_cache_breakpoints( + request: &serde_json::Map, + supported_for_model: bool, +) -> Result<(), OpenAiPromptCacheContractViolation> { + let Some(messages) = request.get("messages").and_then(Value::as_array) else { + return Ok(()); + }; + for (message_index, message) in messages.iter().enumerate() { + let Some(content) = message + .as_object() + .and_then(|message| message.get("content")) + .and_then(Value::as_array) + else { + continue; + }; + for (content_index, part) in content.iter().enumerate() { + let Some(part) = part.as_object() else { + continue; + }; + let Some(breakpoint) = part.get("prompt_cache_breakpoint") else { + continue; + }; + let block_type = part.get("type").and_then(Value::as_str); + let supported = matches!( + block_type, + Some("text" | "image_url" | "input_audio" | "file" | "refusal") + ); + validate_prompt_cache_breakpoint( + breakpoint, + &format!( + "messages[{message_index}].content[{content_index}].prompt_cache_breakpoint" + ), + supported, + supported_for_model, + )?; + } + } + Ok(()) +} + +fn validate_responses_prompt_cache_breakpoints( + request: &serde_json::Map, + supported_for_model: bool, +) -> Result<(), OpenAiPromptCacheContractViolation> { + let Some(input) = request.get("input").and_then(Value::as_array) else { + return Ok(()); + }; + for (item_index, item) in input.iter().enumerate() { + let Some(item) = item.as_object() else { + continue; + }; + if let Some(breakpoint) = item.get("prompt_cache_breakpoint") { + validate_prompt_cache_breakpoint( + breakpoint, + &format!("input[{item_index}].prompt_cache_breakpoint"), + responses_cache_breakpoint_block_is_supported(item), + supported_for_model, + )?; + } + let Some(content) = item.get("content").and_then(Value::as_array) else { + continue; + }; + for (content_index, part) in content.iter().enumerate() { + let Some(part) = part.as_object() else { + continue; + }; + let Some(breakpoint) = part.get("prompt_cache_breakpoint") else { + continue; + }; + validate_prompt_cache_breakpoint( + breakpoint, + &format!("input[{item_index}].content[{content_index}].prompt_cache_breakpoint"), + responses_cache_breakpoint_block_is_supported(part), + supported_for_model, + )?; + } + } + Ok(()) +} + +fn responses_cache_breakpoint_block_is_supported(block: &serde_json::Map) -> bool { + matches!( + block.get("type").and_then(Value::as_str), + Some("input_text" | "input_image" | "input_file") + ) +} + +fn validate_prompt_cache_breakpoint( + value: &Value, + field: &str, + supported_content_block: bool, + supported_for_model: bool, +) -> Result<(), OpenAiPromptCacheContractViolation> { + if !supported_for_model { + return Err(unsupported_for_model( + field, + "provider model does not support prompt_cache_breakpoint", + )); + } + if !supported_content_block { + return Err(OpenAiPromptCacheContractViolation { + kind: OpenAiPromptCacheViolationKind::UnsupportedContentBlock, + field: field.to_string(), + value: None, + reason: "content block does not support prompt_cache_breakpoint".to_string(), + }); + } + let Some(breakpoint) = value.as_object() else { + return Err(invalid_type( + field, + value, + "prompt_cache_breakpoint must be an object", + )); + }; + let mode_field = format!("{field}.mode"); + let Some(mode) = breakpoint.get("mode") else { + return Err(OpenAiPromptCacheContractViolation { + kind: OpenAiPromptCacheViolationKind::InvalidType, + field: mode_field, + value: None, + reason: "prompt_cache_breakpoint.mode is required".to_string(), + }); + }; + let Some(raw) = mode.as_str() else { + return Err(invalid_type( + &mode_field, + mode, + "prompt_cache_breakpoint.mode must be a string", + )); + }; + if raw != "explicit" { + return Err(invalid_enum( + &mode_field, + raw, + "prompt_cache_breakpoint.mode supports explicit", + )); + } + Ok(()) +} + +fn invalid_type(field: &str, value: &Value, reason: &str) -> OpenAiPromptCacheContractViolation { + OpenAiPromptCacheContractViolation { + kind: OpenAiPromptCacheViolationKind::InvalidType, + field: field.to_string(), + value: Some(value.to_string()), + reason: reason.to_string(), + } +} + +fn invalid_enum(field: &str, value: &str, reason: &str) -> OpenAiPromptCacheContractViolation { + OpenAiPromptCacheContractViolation { + kind: OpenAiPromptCacheViolationKind::InvalidEnum, + field: field.to_string(), + value: Some(value.to_string()), + reason: reason.to_string(), + } +} + +fn unsupported_for_model(field: &str, reason: &str) -> OpenAiPromptCacheContractViolation { + OpenAiPromptCacheContractViolation { + kind: OpenAiPromptCacheViolationKind::UnsupportedForModel, + field: field.to_string(), + value: None, + reason: reason.to_string(), + } +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::{ + resolve_openai_prompt_cache_ttl_minutes, validate_openai_prompt_cache_request, + OpenAiPromptCacheViolationKind, + }; + + #[test] + fn resolves_effective_prompt_cache_ttl_from_current_openai_contract() { + for body in [ + json!({"model": "client-alias"}), + json!({ + "model": "client-alias", + "prompt_cache_options": {"mode": "implicit"} + }), + json!({ + "model": "client-alias", + "prompt_cache_options": {"mode": "explicit", "ttl": "30m"} + }), + ] { + assert_eq!( + resolve_openai_prompt_cache_ttl_minutes( + "openai:responses", + "gpt-5.6-sol", + "client-alias", + &body, + ), + Some(30) + ); + } + + assert_eq!( + resolve_openai_prompt_cache_ttl_minutes( + "openai:chat", + "deployment-alias", + "gpt-5.6-terra", + &json!({"model": "gpt-5.6-terra"}), + ), + Some(30) + ); + assert_eq!( + resolve_openai_prompt_cache_ttl_minutes( + "openai:chat", + "gpt-5.5", + "gpt-5.6-terra", + &json!({"model": "gpt-5.6-terra"}), + ), + None + ); + assert_eq!( + resolve_openai_prompt_cache_ttl_minutes( + "claude:messages", + "gpt-5.6-sol", + "gpt-5.6-sol", + &json!({"model": "gpt-5.6-sol"}), + ), + None + ); + } + + #[test] + fn gpt_5_6_accepts_current_prompt_cache_options_and_breakpoints() { + for (format, body) in [ + ( + "openai:chat", + json!({ + "model": "client-alias", + "prompt_cache_options": {"mode": "implicit", "ttl": "30m"}, + "messages": [{ + "role": "user", + "content": [{ + "type": "input_audio", + "input_audio": {"data": "ZmFrZQ==", "format": "mp3"}, + "prompt_cache_breakpoint": {"mode": "explicit"} + }] + }] + }), + ), + ( + "openai:responses", + json!({ + "model": "client-alias", + "prompt_cache_options": {"mode": "explicit", "ttl": "30m"}, + "input": [{ + "type": "message", + "role": "user", + "content": [{ + "type": "input_file", + "file_id": "file_123", + "prompt_cache_breakpoint": {"mode": "explicit"} + }] + }] + }), + ), + ] { + validate_openai_prompt_cache_request(format, "gpt-5.6-sol", &body) + .expect("GPT-5.6 prompt cache contract should be accepted"); + } + } + + #[test] + fn gpt_5_6_uses_prompt_cache_options_and_rejects_invalid_enums() { + let retention_error = validate_openai_prompt_cache_request( + "openai:chat", + "gpt-5.6-sol", + &json!({"prompt_cache_retention": "24h"}), + ) + .expect_err("GPT-5.6 uses prompt_cache_options.ttl"); + assert_eq!( + retention_error.kind, + OpenAiPromptCacheViolationKind::UnsupportedForModel + ); + + let cases = [ + ( + json!({"prompt_cache_options": {"mode": "automatic"}}), + "prompt_cache_options.mode", + OpenAiPromptCacheViolationKind::InvalidEnum, + ), + ( + json!({"prompt_cache_options": {"ttl": "1h"}}), + "prompt_cache_options.ttl", + OpenAiPromptCacheViolationKind::InvalidEnum, + ), + ( + json!({ + "messages": [{ + "role": "user", + "content": [{ + "type": "text", + "text": "stable", + "prompt_cache_breakpoint": {"mode": "implicit"} + }] + }] + }), + "messages[0].content[0].prompt_cache_breakpoint.mode", + OpenAiPromptCacheViolationKind::InvalidEnum, + ), + ]; + + for (body, field, kind) in cases { + let error = validate_openai_prompt_cache_request("openai:chat", "gpt-5.6-sol", &body) + .expect_err("invalid GPT-5.6 cache contract should fail"); + assert_eq!(error.field, field); + assert_eq!(error.kind, kind); + } + } + + #[test] + fn prompt_cache_capability_uses_the_provider_model() { + let options = json!({ + "model": "gpt-5.5", + "prompt_cache_options": {"mode": "explicit", "ttl": "30m"}, + "messages": [{"role": "user", "content": "hello"}] + }); + validate_openai_prompt_cache_request("openai:chat", "gpt-5.6-terra", &options) + .expect("mapped GPT-5.6 provider model should enable prompt_cache_options"); + let error = validate_openai_prompt_cache_request("openai:chat", "gpt-5.5", &options) + .expect_err("mapped earlier provider model should reject prompt_cache_options"); + assert_eq!( + error.kind, + OpenAiPromptCacheViolationKind::UnsupportedForModel + ); + let error = + validate_openai_prompt_cache_request("openai:chat", "deployment-alias", &options) + .expect_err("opaque provider model must not inherit source model capability"); + assert_eq!( + error.kind, + OpenAiPromptCacheViolationKind::UnsupportedForModel + ); + + let retention = json!({ + "model": "gpt-5.6-sol", + "prompt_cache_retention": "24h", + "messages": [{"role": "user", "content": "hello"}] + }); + validate_openai_prompt_cache_request("openai:chat", "gpt-5.5", &retention) + .expect("mapped earlier provider model should retain its retention contract"); + let error = validate_openai_prompt_cache_request("openai:chat", "gpt-5.6-luna", &retention) + .expect_err("GPT-5.6 uses prompt_cache_options.ttl"); + assert_eq!( + error.kind, + OpenAiPromptCacheViolationKind::UnsupportedForModel + ); + } + + #[test] + fn prompt_cache_breakpoints_validate_supported_blocks_per_api() { + let cases = [ + ( + "openai:chat", + json!({ + "messages": [{ + "role": "user", + "content": [{ + "type": "video_url", + "video_url": {"url": "https://example.com/video.mp4"}, + "prompt_cache_breakpoint": {"mode": "explicit"} + }] + }] + }), + "messages[0].content[0].prompt_cache_breakpoint", + ), + ( + "openai:responses", + json!({ + "input": [{ + "type": "message", + "role": "user", + "content": [{ + "type": "input_audio", + "input_audio": {"data": "ZmFrZQ==", "format": "mp3"}, + "prompt_cache_breakpoint": {"mode": "explicit"} + }] + }] + }), + "input[0].content[0].prompt_cache_breakpoint", + ), + ]; + + for (format, body, field) in cases { + let error = validate_openai_prompt_cache_request(format, "gpt-5.6-sol", &body) + .expect_err("unsupported cache breakpoint block should fail"); + assert_eq!(error.field, field); + assert_eq!( + error.kind, + OpenAiPromptCacheViolationKind::UnsupportedContentBlock + ); + } + } + + #[test] + fn earlier_models_reject_breakpoint_options_and_validate_retention() { + let options = json!({"prompt_cache_options": {"mode": "explicit"}}); + let error = validate_openai_prompt_cache_request("openai:responses", "gpt-5.5", &options) + .expect_err("earlier model should reject prompt_cache_options"); + assert_eq!( + error.kind, + OpenAiPromptCacheViolationKind::UnsupportedForModel + ); + + validate_openai_prompt_cache_request( + "openai:responses", + "gpt-5.5-pro", + &json!({"prompt_cache_retention": "24h"}), + ) + .expect("GPT-5.5 supports 24h retention"); + let error = validate_openai_prompt_cache_request( + "openai:responses", + "gpt-5.5", + &json!({"prompt_cache_retention": "in_memory"}), + ) + .expect_err("GPT-5.5 only supports 24h retention"); + assert_eq!(error.kind, OpenAiPromptCacheViolationKind::InvalidEnum); + } + + #[test] + fn nullable_prompt_cache_fields_are_treated_as_unconfigured() { + for format in [ + "openai:chat", + "openai:responses", + "openai:responses:compact", + ] { + validate_openai_prompt_cache_request( + format, + "gpt-5.6-sol", + &json!({ + "model": "gpt-5.6-sol", + "prompt_cache_options": null, + "prompt_cache_retention": null + }), + ) + .expect("nullable prompt cache fields should be omitted semantically"); + } + } +} diff --git a/crates/aether-ai-formats/src/formats/openai/reasoning.rs b/crates/aether-ai-formats/src/formats/openai/reasoning.rs new file mode 100644 index 000000000..e0632bdb5 --- /dev/null +++ b/crates/aether-ai-formats/src/formats/openai/reasoning.rs @@ -0,0 +1,505 @@ +use serde_json::Value; + +use crate::formats::shared::model_directives::ReasoningEffort; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum OpenAiReasoningViolationKind { + InvalidType, + InvalidEnum, + UnsupportedForModel, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct OpenAiReasoningContractViolation { + pub kind: OpenAiReasoningViolationKind, + pub field: String, + pub value: Option, + pub reason: String, +} + +pub fn validate_openai_reasoning_request( + source_api_format: &str, + provider_api_format: &str, + provider_model: &str, + body: &Value, +) -> Result<(), OpenAiReasoningContractViolation> { + let source_model = body + .get("model") + .and_then(Value::as_str) + .unwrap_or_default(); + validate_openai_reasoning_request_with_source_model( + source_api_format, + provider_api_format, + provider_model, + source_model, + body, + ) +} + +pub(crate) fn validate_openai_reasoning_request_with_source_model( + source_api_format: &str, + provider_api_format: &str, + provider_model: &str, + source_model: &str, + body: &Value, +) -> Result<(), OpenAiReasoningContractViolation> { + validate_openai_reasoning_request_with_model_profile( + source_api_format, + provider_api_format, + provider_model, + source_model, + body, + None, + None, + ) +} + +pub(crate) fn validate_openai_reasoning_request_with_model_profile( + source_api_format: &str, + provider_api_format: &str, + provider_model: &str, + source_model: &str, + body: &Value, + model_card_reasoning_efforts: Option<&[String]>, + supports_reasoning_mode: Option, +) -> Result<(), OpenAiReasoningContractViolation> { + let Some(object) = body.as_object() else { + return Ok(()); + }; + let source_api_format = crate::normalize_api_format_alias(source_api_format); + let reasoning = match source_api_format.as_str() { + "openai:responses" | "openai:responses:compact" | "openai:search" => { + match object.get("reasoning") { + Some(Value::Object(reasoning)) => Some(reasoning), + Some(Value::Null) => None, + Some(value) => { + return Err(OpenAiReasoningContractViolation { + kind: OpenAiReasoningViolationKind::InvalidType, + field: "reasoning".to_string(), + value: Some(value.to_string()), + reason: "reasoning must be an object".to_string(), + }); + } + None => None, + } + } + "openai:chat" => None, + _ => return Ok(()), + }; + + let provider_model = provider_model.trim(); + let provider_model = if provider_model.is_empty() { + source_model + } else { + provider_model + }; + + let effort = match source_api_format.as_str() { + "openai:chat" => object.get("reasoning_effort"), + "openai:responses" | "openai:responses:compact" | "openai:search" => { + reasoning.and_then(|reasoning| reasoning.get("effort")) + } + _ => None, + }; + if let Some(value) = effort.filter(|value| !value.is_null()) { + validate_reasoning_effort( + value, + source_api_format.as_str(), + provider_api_format, + provider_model, + source_model, + model_card_reasoning_efforts, + )?; + } + + if source_api_format != "openai:search" { + if let Some(mode) = reasoning + .and_then(|reasoning| reasoning.get("mode")) + .filter(|value| !value.is_null()) + { + validate_reasoning_mode(mode, provider_model, source_model, supports_reasoning_mode)?; + } + } + if let Some(context) = reasoning + .and_then(|reasoning| reasoning.get("context")) + .filter(|value| !value.is_null()) + { + validate_reasoning_context(context)?; + } + if let Some(summary) = reasoning + .and_then(|reasoning| reasoning.get("summary")) + .filter(|value| !value.is_null()) + { + validate_reasoning_summary(summary)?; + } + + Ok(()) +} + +fn validate_reasoning_effort( + value: &Value, + source_api_format: &str, + provider_api_format: &str, + provider_model: &str, + source_model: &str, + model_card_reasoning_efforts: Option<&[String]>, +) -> Result<(), OpenAiReasoningContractViolation> { + let field = if source_api_format == "openai:chat" { + "reasoning_effort" + } else { + "reasoning.effort" + }; + let Some(raw) = value.as_str() else { + return Err(OpenAiReasoningContractViolation { + kind: OpenAiReasoningViolationKind::InvalidType, + field: field.to_string(), + value: Some(value.to_string()), + reason: "reasoning effort must be a string".to_string(), + }); + }; + if raw.trim().is_empty() { + return Err(OpenAiReasoningContractViolation { + kind: OpenAiReasoningViolationKind::InvalidEnum, + field: field.to_string(), + value: Some(raw.to_string()), + reason: "reasoning effort must not be empty".to_string(), + }); + } + if raw.trim() == "ultra" { + return Err(OpenAiReasoningContractViolation { + kind: OpenAiReasoningViolationKind::InvalidEnum, + field: field.to_string(), + value: Some(raw.to_string()), + reason: "ultra is a Codex client preset, not an OpenAI wire effort".to_string(), + }); + } + if let Some(supported_efforts) = + model_card_reasoning_efforts.filter(|values| !values.is_empty()) + { + if supported_efforts + .iter() + .any(|effort| effort == raw.trim() || (raw.trim() == "max" && effort == "ultra")) + { + return Ok(()); + } + return Err(OpenAiReasoningContractViolation { + kind: OpenAiReasoningViolationKind::UnsupportedForModel, + field: field.to_string(), + value: Some(raw.to_string()), + reason: "provider model card does not support the requested reasoning effort" + .to_string(), + }); + } + let Some(effort) = ReasoningEffort::parse(raw) else { + return Ok(()); + }; + if crate::reasoning_effort_supported_for_model( + provider_api_format, + provider_model, + source_model, + effort, + ) { + return Ok(()); + } + Err(OpenAiReasoningContractViolation { + kind: OpenAiReasoningViolationKind::UnsupportedForModel, + field: field.to_string(), + value: Some(raw.to_string()), + reason: "provider model does not support the requested reasoning effort".to_string(), + }) +} + +fn validate_reasoning_mode( + value: &Value, + provider_model: &str, + source_model: &str, + supports_reasoning_mode: Option, +) -> Result<(), OpenAiReasoningContractViolation> { + let Some(mode) = value.as_str() else { + return Err(OpenAiReasoningContractViolation { + kind: OpenAiReasoningViolationKind::InvalidType, + field: "reasoning.mode".to_string(), + value: Some(value.to_string()), + reason: "reasoning mode must be a string".to_string(), + }); + }; + if mode.trim().is_empty() { + return Err(OpenAiReasoningContractViolation { + kind: OpenAiReasoningViolationKind::InvalidEnum, + field: "reasoning.mode".to_string(), + value: Some(mode.to_string()), + reason: "reasoning mode must not be empty".to_string(), + }); + } + if !matches!(mode, "standard" | "pro") { + return Err(OpenAiReasoningContractViolation { + kind: OpenAiReasoningViolationKind::InvalidEnum, + field: "reasoning.mode".to_string(), + value: Some(mode.to_string()), + reason: "reasoning mode supports standard or pro".to_string(), + }); + } + let supported = supports_reasoning_mode.unwrap_or_else(|| { + crate::formats::shared::model_directives::openai_model_resolves_to_gpt_5_6( + provider_model, + source_model, + ) + }); + if supported { + return Ok(()); + } + Err(OpenAiReasoningContractViolation { + kind: OpenAiReasoningViolationKind::UnsupportedForModel, + field: "reasoning.mode".to_string(), + value: Some(mode.to_string()), + reason: "provider model does not support reasoning mode".to_string(), + }) +} + +fn validate_reasoning_context(value: &Value) -> Result<(), OpenAiReasoningContractViolation> { + let Some(context) = value.as_str() else { + return Err(OpenAiReasoningContractViolation { + kind: OpenAiReasoningViolationKind::InvalidType, + field: "reasoning.context".to_string(), + value: Some(value.to_string()), + reason: "reasoning context must be a string".to_string(), + }); + }; + if matches!(context, "auto" | "current_turn" | "all_turns") { + return Ok(()); + } + Err(OpenAiReasoningContractViolation { + kind: OpenAiReasoningViolationKind::InvalidEnum, + field: "reasoning.context".to_string(), + value: Some(context.to_string()), + reason: "reasoning context is not a supported wire value".to_string(), + }) +} + +fn validate_reasoning_summary(value: &Value) -> Result<(), OpenAiReasoningContractViolation> { + let Some(summary) = value.as_str() else { + return Err(OpenAiReasoningContractViolation { + kind: OpenAiReasoningViolationKind::InvalidType, + field: "reasoning.summary".to_string(), + value: Some(value.to_string()), + reason: "reasoning summary must be a string".to_string(), + }); + }; + if matches!(summary, "auto" | "concise" | "detailed") { + return Ok(()); + } + Err(OpenAiReasoningContractViolation { + kind: OpenAiReasoningViolationKind::InvalidEnum, + field: "reasoning.summary".to_string(), + value: Some(summary.to_string()), + reason: "reasoning summary is not a supported wire value".to_string(), + }) +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::{validate_openai_reasoning_request, OpenAiReasoningViolationKind}; + + #[test] + fn mapped_model_is_authoritative_for_openai_reasoning_effort() { + let alias = json!({ + "model": "deployment-alias", + "reasoning": {"effort": "max"} + }); + validate_openai_reasoning_request( + "openai:responses", + "openai:responses", + "gpt-5.6-sol", + &alias, + ) + .expect("GPT-5.6 should accept max"); + + let error = validate_openai_reasoning_request( + "openai:responses", + "openai:responses", + "gpt-5.4", + &alias, + ) + .expect_err("GPT-5.4 should reject max"); + assert_eq!( + error.kind, + OpenAiReasoningViolationKind::UnsupportedForModel + ); + } + + #[test] + fn gpt_5_6_rejects_known_unsupported_effort_and_preserves_custom_effort() { + let unsupported = json!({ + "model": "gpt-5.6-terra", + "reasoning_effort": "minimal" + }); + let error = validate_openai_reasoning_request( + "openai:chat", + "openai:chat", + "gpt-5.6-terra", + &unsupported, + ) + .expect_err("known unsupported effort should be rejected"); + assert_eq!( + error.kind, + OpenAiReasoningViolationKind::UnsupportedForModel + ); + + validate_openai_reasoning_request( + "openai:responses", + "openai:responses", + "gpt-5.6-terra", + &json!({ + "model": "gpt-5.6-terra", + "reasoning": {"effort": "future"} + }), + ) + .expect("model-advertised custom effort should pass through"); + + let ultra = validate_openai_reasoning_request( + "openai:responses", + "openai:responses", + "gpt-5.6-terra", + &json!({"reasoning": {"effort": "ultra"}}), + ) + .expect_err("Codex local ultra preset should not enter the OpenAI wire contract"); + assert_eq!(ultra.kind, OpenAiReasoningViolationKind::InvalidEnum); + + let empty = validate_openai_reasoning_request( + "openai:responses", + "openai:responses", + "gpt-5.6-terra", + &json!({"reasoning": {"effort": " "}}), + ) + .expect_err("empty reasoning effort should be rejected"); + assert_eq!(empty.kind, OpenAiReasoningViolationKind::InvalidEnum); + } + + #[test] + fn reasoning_mode_is_responses_only_and_requires_gpt_5_6() { + for mode in ["standard", "pro"] { + validate_openai_reasoning_request( + "openai:responses", + "openai:responses", + "gpt-5.6-sol", + &json!({"model": "deployment-alias", "reasoning": {"mode": mode}}), + ) + .expect("GPT-5.6 should accept reasoning mode"); + } + + let unsupported = validate_openai_reasoning_request( + "openai:responses", + "openai:responses", + "gpt-5.4", + &json!({"reasoning": {"mode": "pro"}}), + ) + .expect_err("earlier GPT models should reject reasoning mode"); + assert_eq!( + unsupported.kind, + OpenAiReasoningViolationKind::UnsupportedForModel + ); + + let invalid = validate_openai_reasoning_request( + "openai:responses", + "openai:responses", + "gpt-5.6-sol", + &json!({"reasoning": {"mode": "fast"}}), + ) + .expect_err("unknown reasoning mode should be rejected"); + assert_eq!(invalid.kind, OpenAiReasoningViolationKind::InvalidEnum); + + let empty = validate_openai_reasoning_request( + "openai:responses", + "openai:responses", + "gpt-5.6-sol", + &json!({"reasoning": {"mode": ""}}), + ) + .expect_err("empty reasoning mode should be rejected"); + assert_eq!(empty.kind, OpenAiReasoningViolationKind::InvalidEnum); + + validate_openai_reasoning_request( + "openai:chat", + "openai:chat", + "gpt-5.4", + &json!({"reasoning": {"mode": "pro"}}), + ) + .expect("Chat Completions does not define reasoning.mode"); + } + + #[test] + fn reasoning_context_validates_wire_values_without_model_gating() { + for context in ["auto", "current_turn", "all_turns"] { + validate_openai_reasoning_request( + "openai:responses", + "openai:responses", + "gpt-5.4", + &json!({"reasoning": {"context": context}}), + ) + .expect("reasoning context should remain available to Codex Responses models"); + } + + let invalid = validate_openai_reasoning_request( + "openai:responses", + "openai:responses", + "gpt-5.6-sol", + &json!({"reasoning": {"context": "session"}}), + ) + .expect_err("unknown reasoning context should be rejected"); + assert_eq!(invalid.kind, OpenAiReasoningViolationKind::InvalidEnum); + } + + #[test] + fn reasoning_summary_accepts_only_openai_wire_values() { + for summary in ["auto", "concise", "detailed"] { + validate_openai_reasoning_request( + "openai:responses", + "openai:responses", + "gpt-5.6-sol", + &json!({"reasoning": {"summary": summary}}), + ) + .expect("documented reasoning summary should be accepted"); + } + + for (summary, expected_kind) in [ + (json!("none"), OpenAiReasoningViolationKind::InvalidEnum), + (json!(true), OpenAiReasoningViolationKind::InvalidType), + ] { + let error = validate_openai_reasoning_request( + "openai:responses", + "openai:responses", + "gpt-5.6-sol", + &json!({"reasoning": {"summary": summary}}), + ) + .expect_err("invalid reasoning summary should be rejected"); + assert_eq!(error.kind, expected_kind); + } + } + + #[test] + fn nullable_reasoning_fields_are_treated_as_unconfigured() { + for body in [ + json!({"model": "gpt-5.6-sol", "reasoning": null}), + json!({ + "model": "gpt-5.6-sol", + "reasoning": {"effort": null, "mode": null, "context": null} + }), + ] { + validate_openai_reasoning_request( + "openai:responses", + "openai:responses", + "gpt-5.6-sol", + &body, + ) + .expect("nullable Responses reasoning fields should be omitted semantically"); + } + validate_openai_reasoning_request( + "openai:chat", + "openai:chat", + "gpt-5.6-sol", + &json!({"model": "gpt-5.6-sol", "reasoning_effort": null}), + ) + .expect("nullable Chat reasoning effort should be omitted semantically"); + } +} diff --git a/crates/aether-ai-formats/src/formats/openai/request_contract.rs b/crates/aether-ai-formats/src/formats/openai/request_contract.rs new file mode 100644 index 000000000..722c4aa13 --- /dev/null +++ b/crates/aether-ai-formats/src/formats/openai/request_contract.rs @@ -0,0 +1,788 @@ +use serde_json::Value; + +use super::prompt_cache::OpenAiPromptCacheContractViolation; +use super::reasoning::OpenAiReasoningContractViolation; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum OpenAiProviderRequestContractViolation { + CodexCompact(super::responses::codex::CodexOpenAiCompactRequestContractViolation), + Responses(super::responses::request::OpenAiResponsesRequestContractViolation), + PromptCache(OpenAiPromptCacheContractViolation), + Reasoning(OpenAiReasoningContractViolation), +} + +#[derive(Clone, Copy, Debug)] +pub struct OpenAiProviderRequestFinalization<'a> { + pub source_api_format: &'a str, + pub provider_api_format: &'a str, + pub provider_type: &'a str, + pub provider_model: &'a str, + pub source_model: &'a str, + pub body_rules: Option<&'a Value>, + pub upstream_is_stream: bool, + pub require_body_stream_field: bool, +} + +pub fn finalize_openai_provider_request( + body: &mut Value, + finalization: OpenAiProviderRequestFinalization<'_>, +) -> Result<(), OpenAiProviderRequestContractViolation> { + finalize_openai_provider_request_with_codex_model_capabilities(body, finalization, None) +} + +pub fn finalize_openai_provider_request_with_codex_model_capabilities( + body: &mut Value, + finalization: OpenAiProviderRequestFinalization<'_>, + model_capabilities: Option<&super::responses::codex::CodexResponsesModelCapabilities>, +) -> Result<(), OpenAiProviderRequestContractViolation> { + let is_codex_reasoning_endpoint = finalization + .provider_type + .trim() + .eq_ignore_ascii_case("codex") + && (crate::is_openai_responses_family_format(finalization.provider_api_format) + || crate::api_format_alias_matches(finalization.provider_api_format, "openai:search")); + let resolved_model_capabilities = (is_codex_reasoning_endpoint && model_capabilities.is_none()) + .then(|| { + super::responses::codex::resolve_codex_responses_model_capabilities( + finalization.provider_model, + finalization.source_model, + None, + ) + }); + let model_capabilities = is_codex_reasoning_endpoint + .then(|| model_capabilities.or(resolved_model_capabilities.as_ref())) + .flatten(); + match crate::normalize_api_format_alias(finalization.source_api_format).as_str() { + "openai:responses" | "openai:responses:compact" => { + super::responses::codex::apply_codex_openai_responses_special_body_edits_with_source_model_and_capabilities( + body, + finalization.provider_type, + finalization.provider_api_format, + finalization.provider_model, + finalization.source_model, + model_capabilities, + finalization.body_rules, + ); + } + _ => { + super::responses::codex::apply_codex_openai_responses_chat_body_edits_with_source_model_and_capabilities( + body, + finalization.provider_type, + finalization.provider_api_format, + finalization.provider_model, + finalization.source_model, + model_capabilities, + finalization.body_rules, + ) + } + } + super::responses::codex::normalize_codex_openai_reasoning_wire_effort( + body, + finalization.provider_type, + finalization.provider_api_format, + ); + super::responses::codex::apply_openai_responses_compact_special_body_edits( + body, + finalization.provider_api_format, + ); + crate::enforce_request_body_stream_field( + body, + finalization.provider_api_format, + finalization.upstream_is_stream, + finalization.require_body_stream_field, + ); + super::search::apply_openai_search_request_projection(body, finalization.provider_api_format); + let provider_model = body + .get("model") + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .unwrap_or(finalization.provider_model); + super::responses::codex::validate_codex_openai_responses_compact_request_contract( + body, + finalization.provider_type, + finalization.provider_api_format, + ) + .map_err(OpenAiProviderRequestContractViolation::CodexCompact)?; + validate_openai_provider_request_contract_with_codex_model_capabilities( + finalization.provider_api_format, + provider_model, + finalization.source_model, + body, + model_capabilities, + ) +} + +pub fn validate_openai_provider_request_contract( + provider_api_format: &str, + provider_model: &str, + source_model: &str, + body: &Value, +) -> Result<(), OpenAiProviderRequestContractViolation> { + validate_openai_provider_request_contract_with_codex_model_capabilities( + provider_api_format, + provider_model, + source_model, + body, + None, + ) +} + +fn validate_openai_provider_request_contract_with_codex_model_capabilities( + provider_api_format: &str, + provider_model: &str, + source_model: &str, + body: &Value, + model_capabilities: Option<&super::responses::codex::CodexResponsesModelCapabilities>, +) -> Result<(), OpenAiProviderRequestContractViolation> { + super::responses::request::validate_openai_responses_request_contract( + body, + provider_api_format, + ) + .map_err(OpenAiProviderRequestContractViolation::Responses)?; + super::prompt_cache::validate_openai_prompt_cache_request_with_source_model( + provider_api_format, + provider_model, + source_model, + body, + ) + .map_err(OpenAiProviderRequestContractViolation::PromptCache)?; + super::reasoning::validate_openai_reasoning_request_with_model_profile( + provider_api_format, + provider_api_format, + provider_model, + source_model, + body, + model_capabilities.map(|capabilities| capabilities.supported_reasoning_efforts.as_slice()), + None, + ) + .map_err(OpenAiProviderRequestContractViolation::Reasoning) +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::{ + finalize_openai_provider_request, + finalize_openai_provider_request_with_codex_model_capabilities, + validate_openai_provider_request_contract, OpenAiProviderRequestFinalization, + }; + use crate::CodexResponsesModelCapabilities; + + #[test] + fn validates_reasoning_and_prompt_cache_against_the_final_provider_model() { + let body = json!({ + "model": "gpt-5.6-sol", + "input": [], + "reasoning": {"effort": "max"}, + "prompt_cache_options": {"mode": "explicit", "ttl": "30m"} + }); + validate_openai_provider_request_contract( + "openai:responses", + "gpt-5.6-sol", + "gpt-5.6-sol", + &body, + ) + .expect("GPT-5.6 request should satisfy the final provider contract"); + + assert!(validate_openai_provider_request_contract( + "openai:responses", + "gpt-5.4", + "gpt-5.6-sol", + &body, + ) + .is_err()); + } + + #[test] + fn opaque_provider_models_inherit_source_capabilities_but_concrete_models_do_not() { + let body = json!({ + "model": "azure-production", + "input": [], + "reasoning": {"effort": "max", "mode": "pro"}, + "prompt_cache_options": {"mode": "explicit", "ttl": "30m"} + }); + validate_openai_provider_request_contract( + "openai:responses", + "azure-production", + "gpt-5.6-sol-max", + &body, + ) + .expect("opaque deployments should inherit the concrete source model capability"); + assert!(validate_openai_provider_request_contract( + "openai:responses", + "gpt-5.4", + "gpt-5.6-sol", + &body, + ) + .is_err()); + } + + #[test] + fn codex_finalization_enforces_model_card_reasoning_efforts() { + let mut body = json!({ + "model": "gpt-5.6-sol", + "input": [], + "reasoning": {"effort": "minimal"} + }); + + let error = finalize_openai_provider_request( + &mut body, + OpenAiProviderRequestFinalization { + source_api_format: "openai:responses", + provider_api_format: "openai:responses", + provider_type: "codex", + provider_model: "gpt-5.6-sol", + source_model: "gpt-5.6-sol", + body_rules: None, + upstream_is_stream: false, + require_body_stream_field: true, + }, + ) + .expect_err("GPT-5.6 Codex model card should reject minimal"); + + assert!(matches!( + error, + super::OpenAiProviderRequestContractViolation::Reasoning( + super::OpenAiReasoningContractViolation { + kind: crate::formats::openai::reasoning::OpenAiReasoningViolationKind::UnsupportedForModel, + .. + } + ) + )); + } + + #[test] + fn finalization_reapplies_codex_and_compact_projection_after_mutations() { + let mut body = json!({ + "model": "gpt-5.6-sol", + "input": [], + "store": true, + "include": ["reasoning.encrypted_content"], + "client_metadata": {"source": "mapping"}, + "stream": true, + "stream_options": {"include_usage": true}, + "tool_choice": "auto", + "temperature": 0.5, + "previous_response_id": "resp_123" + }); + finalize_openai_provider_request( + &mut body, + OpenAiProviderRequestFinalization { + source_api_format: "openai:responses", + provider_api_format: "openai:responses:compact", + provider_type: "codex", + provider_model: "gpt-5.6-sol", + source_model: "gpt-5.6-sol", + body_rules: None, + upstream_is_stream: false, + require_body_stream_field: true, + }, + ) + .expect("final Compact request should satisfy its provider contract"); + + for field in [ + "store", + "include", + "client_metadata", + "stream", + "stream_options", + "tool_choice", + "temperature", + "previous_response_id", + ] { + assert!(body.get(field).is_none(), "{field} must not reach Compact"); + } + } + + #[test] + fn non_responses_sources_receive_codex_responses_reasoning_defaults() { + for source_api_format in ["openai:chat", "claude:messages", "gemini:generate_content"] { + let mut body = json!({ + "model": "gpt-5.6-sol", + "input": [] + }); + finalize_openai_provider_request( + &mut body, + OpenAiProviderRequestFinalization { + source_api_format, + provider_api_format: "openai:responses", + provider_type: "codex", + provider_model: "gpt-5.6-sol", + source_model: "gpt-5.6-sol", + body_rules: None, + upstream_is_stream: true, + require_body_stream_field: true, + }, + ) + .expect("Codex Responses request should satisfy the final provider contract"); + + assert_eq!(body["reasoning"]["effort"], "low"); + assert!(body["reasoning"].get("summary").is_none()); + } + } + + #[test] + fn codex_ultra_preset_uses_max_for_every_codex_model_on_the_wire() { + let mut sol = json!({ + "model": "gpt-5.6-sol", + "input": [], + "reasoning": {"effort": "ultra"} + }); + let mut luna = json!({ + "model": "gpt-5.6-luna", + "input": [], + "reasoning": {"effort": "ultra"} + }); + let finalization_for = |model| OpenAiProviderRequestFinalization { + source_api_format: "openai:responses", + provider_api_format: "openai:responses", + provider_type: "codex", + provider_model: model, + source_model: model, + body_rules: None, + upstream_is_stream: true, + require_body_stream_field: true, + }; + + finalize_openai_provider_request(&mut sol, finalization_for("gpt-5.6-sol")) + .expect("Sol ultra preset should map to the OpenAI wire contract"); + assert_eq!(sol["reasoning"]["effort"], "max"); + + finalize_openai_provider_request(&mut luna, finalization_for("gpt-5.6-luna")) + .expect("Luna ultra preset should map to the OpenAI wire contract"); + assert_eq!(luna["reasoning"]["effort"], "max"); + } + + #[test] + fn dynamic_codex_card_controls_default_effort_and_keeps_mode_model_specific() { + let finalization = OpenAiProviderRequestFinalization { + source_api_format: "openai:responses", + provider_api_format: "openai:responses", + provider_type: "codex", + provider_model: "gpt-5.7-sol", + source_model: "gpt-5.7-sol", + body_rules: None, + upstream_is_stream: true, + require_body_stream_field: true, + }; + let capabilities = CodexResponsesModelCapabilities { + use_responses_lite: true, + supports_reasoning_summary_parameter: true, + default_reasoning_effort: Some("ultra".to_string()), + default_reasoning_summary: None, + supported_reasoning_efforts: vec!["max".to_string(), "ultra".to_string()], + supports_parallel_tool_calls: true, + support_verbosity: true, + default_verbosity: Some("low".to_string()), + supported_service_tiers: vec!["priority".to_string()], + }; + + let mut default_body = json!({"model": "gpt-5.7-sol", "input": []}); + finalize_openai_provider_request_with_codex_model_capabilities( + &mut default_body, + finalization, + Some(&capabilities), + ) + .expect("card default ultra should use the max wire effort"); + assert_eq!(default_body["reasoning"]["effort"], "max"); + + let mut mode_body = json!({ + "model": "gpt-5.7-sol", + "input": [], + "reasoning": {"effort": "max", "mode": "pro"} + }); + let mode_error = finalize_openai_provider_request_with_codex_model_capabilities( + &mut mode_body, + finalization, + Some(&capabilities), + ) + .expect_err("Responses Lite alone must not enable GPT-5.6 reasoning modes"); + assert!(matches!( + mode_error, + super::OpenAiProviderRequestContractViolation::Reasoning(_) + )); + + let ultra_only = CodexResponsesModelCapabilities { + supported_reasoning_efforts: vec!["ultra".to_string()], + ..capabilities.clone() + }; + let mut ultra_only_body = json!({ + "model": "gpt-5.7-sol", + "input": [], + "reasoning": {"effort": "ultra"} + }); + finalize_openai_provider_request_with_codex_model_capabilities( + &mut ultra_only_body, + finalization, + Some(&ultra_only), + ) + .expect("Codex maps the Ultra preset to max without card-list wire validation"); + assert_eq!(ultra_only_body["reasoning"]["effort"], "max"); + } + + #[test] + fn codex_search_normalizes_reasoning_and_projects_the_typed_request() { + let finalization = OpenAiProviderRequestFinalization { + source_api_format: "openai:search", + provider_api_format: "openai:search", + provider_type: "codex", + provider_model: "gpt-5.6-sol", + source_model: "gpt-5.6-sol", + body_rules: None, + upstream_is_stream: false, + require_body_stream_field: false, + }; + let mut body = json!({ + "id": "session-1", + "model": "gpt-5.6-sol", + "reasoning": { + "effort": "ultra", + "summary": "auto", + "context": "current_turn", + "future_reasoning_field": true + }, + "commands": {"search_query": [{"q": "Aether"}]}, + "store": false, + "future_request_field": {"enabled": true}, + "stream": true + }); + + finalize_openai_provider_request(&mut body, finalization) + .expect("Codex Search request should finalize"); + + assert_eq!(body["reasoning"]["effort"], "max"); + assert_eq!(body["reasoning"]["summary"], "auto"); + assert_eq!(body["reasoning"]["context"], "current_turn"); + assert!(body["reasoning"].get("future_reasoning_field").is_none()); + assert_eq!(body["commands"]["search_query"][0]["q"], "Aether"); + assert!(body.get("store").is_none()); + assert!(body.get("future_request_field").is_none()); + assert!(body.get("stream").is_none()); + assert!(body.get("tool_choice").is_none()); + assert!(body.get("include").is_none()); + } + + #[test] + fn codex_search_validates_reasoning_effort_against_model_card() { + let finalization = OpenAiProviderRequestFinalization { + source_api_format: "openai:search", + provider_api_format: "openai:search", + provider_type: "codex", + provider_model: "gpt-5.6-sol", + source_model: "gpt-5.6-sol", + body_rules: None, + upstream_is_stream: false, + require_body_stream_field: false, + }; + let mut supported = json!({ + "id": "session-1", + "model": "gpt-5.6-sol", + "reasoning": {"effort": "high"} + }); + finalize_openai_provider_request(&mut supported, finalization) + .expect("published Search effort should pass"); + + let mut unsupported = json!({ + "id": "session-1", + "model": "gpt-5.6-sol", + "reasoning": {"effort": "none"} + }); + let error = finalize_openai_provider_request(&mut unsupported, finalization) + .expect_err("unpublished Search effort should be rejected"); + assert!(matches!( + error, + super::OpenAiProviderRequestContractViolation::Reasoning(_) + )); + } + + #[test] + fn dynamic_codex_card_preserves_custom_reasoning_effort_case() { + let finalization = OpenAiProviderRequestFinalization { + source_api_format: "openai:responses", + provider_api_format: "openai:responses", + provider_type: "codex", + provider_model: "codex-custom", + source_model: "codex-custom", + body_rules: None, + upstream_is_stream: true, + require_body_stream_field: true, + }; + let capabilities = CodexResponsesModelCapabilities { + use_responses_lite: false, + supports_reasoning_summary_parameter: true, + default_reasoning_effort: Some("VendorEffortX".to_string()), + default_reasoning_summary: None, + supported_reasoning_efforts: vec!["VendorEffortX".to_string()], + supports_parallel_tool_calls: true, + support_verbosity: true, + default_verbosity: Some("low".to_string()), + supported_service_tiers: vec![], + }; + + let mut body = json!({"model": "codex-custom", "input": []}); + finalize_openai_provider_request_with_codex_model_capabilities( + &mut body, + finalization, + Some(&capabilities), + ) + .expect("custom card effort should remain exact"); + assert_eq!(body["reasoning"]["effort"], "VendorEffortX"); + + let mut custom = json!({ + "model": "codex-custom", + "input": [], + "reasoning": {"effort": "vendoreffortx"} + }); + let error = finalize_openai_provider_request_with_codex_model_capabilities( + &mut custom, + finalization, + Some(&capabilities), + ) + .expect_err("custom reasoning efforts should match the model card exactly"); + assert!(matches!( + error, + super::OpenAiProviderRequestContractViolation::Reasoning(_) + )); + + let mut ultra = json!({ + "model": "codex-custom", + "input": [], + "reasoning": {"effort": "ultra"} + }); + let error = finalize_openai_provider_request_with_codex_model_capabilities( + &mut ultra, + finalization, + Some(&capabilities), + ) + .expect_err("ultra should require model-card support before mapping to max"); + assert!(matches!( + error, + super::OpenAiProviderRequestContractViolation::Reasoning(_) + )); + } + + #[test] + fn gpt_5_6_sol_uses_the_responses_lite_request_contract() { + let mut body = json!({ + "model": "gpt-5.6-sol", + "instructions": "Follow the project instructions.", + "input": [ + { + "type": "message", + "role": "user", + "content": [{ + "type": "input_image", + "image_url": "data:image/png;base64,aGVsbG8=", + "detail": "original" + }] + }, + { + "type": "function_call_output", + "call_id": "call-1", + "output": [{ + "type": "input_image", + "image_url": "data:image/png;base64,ZnVuY3Rpb24=", + "detail": "high" + }] + }, + { + "type": "custom_tool_call_output", + "call_id": "call-2", + "output": [{ + "type": "input_image", + "image_url": "data:image/png;base64,Y3VzdG9t", + "detail": "auto" + }] + } + ], + "tools": [{ + "type": "function", + "name": "lookup", + "parameters": { + "type": "object", + "properties": { + "detail": {"type": "string"} + } + } + }], + "parallel_tool_calls": true + }); + + let finalization = OpenAiProviderRequestFinalization { + source_api_format: "openai:responses", + provider_api_format: "openai:responses", + provider_type: "codex", + provider_model: "gpt-5.6-sol", + source_model: "gpt-5.6-sol", + body_rules: None, + upstream_is_stream: true, + require_body_stream_field: true, + }; + finalize_openai_provider_request(&mut body, finalization) + .expect("GPT-5.6 Sol should satisfy the Responses Lite contract"); + let first = body.clone(); + finalize_openai_provider_request(&mut body, finalization) + .expect("Responses Lite finalization should be idempotent"); + + assert_eq!(body, first); + assert!(body.get("instructions").is_none()); + assert!(body.get("tools").is_none()); + assert_eq!(body["input"][0]["type"], "additional_tools"); + assert_eq!(body["input"][0]["role"], "developer"); + assert_eq!(body["input"][0]["tools"][0]["name"], "lookup"); + assert_eq!(body["input"][1]["type"], "message"); + assert_eq!(body["input"][1]["role"], "developer"); + assert_eq!( + body["input"][1]["content"][0]["text"], + "Follow the project instructions." + ); + assert!(body["input"][2]["content"][0].get("detail").is_none()); + assert!(body["input"][3]["output"][0].get("detail").is_none()); + assert!(body["input"][4]["output"][0].get("detail").is_none()); + assert_eq!( + body["input"][0]["tools"][0]["parameters"]["properties"]["detail"]["type"], + "string" + ); + assert_eq!(body["parallel_tool_calls"], false); + assert_eq!(body["reasoning"]["effort"], "low"); + assert_eq!(body["reasoning"]["context"], "all_turns"); + assert!(body["reasoning"].get("summary").is_none()); + } + + #[test] + fn opaque_codex_deployments_use_the_exact_source_model_card() { + let mut body = json!({ + "model": "azure-production", + "instructions": "Use the configured tools.", + "input": [], + "tools": [], + "parallel_tool_calls": true + }); + + finalize_openai_provider_request( + &mut body, + OpenAiProviderRequestFinalization { + source_api_format: "openai:responses", + provider_api_format: "openai:responses:compact", + provider_type: "codex", + provider_model: "azure-production", + source_model: "gpt-5.6-terra", + body_rules: None, + upstream_is_stream: false, + require_body_stream_field: true, + }, + ) + .expect("opaque Codex deployment should use the exact source model card"); + + assert!(body.get("instructions").is_none()); + assert!(body.get("tools").is_none()); + assert_eq!(body["input"][0]["type"], "additional_tools"); + assert_eq!(body["input"][1]["role"], "developer"); + assert_eq!(body["parallel_tool_calls"], false); + assert_eq!(body["reasoning"]["effort"], "medium"); + assert_eq!(body["reasoning"]["context"], "all_turns"); + for field in body.as_object().expect("object").keys() { + assert!( + [ + "model", + "input", + "instructions", + "tools", + "parallel_tool_calls", + "reasoning", + "service_tier", + "prompt_cache_key", + "text", + ] + .contains(&field.as_str()), + "unexpected Compact field: {field}" + ); + } + } + + #[test] + fn gpt_5_4_keeps_the_standard_codex_responses_shape() { + let mut body = json!({ + "model": "gpt-5.4", + "instructions": "Keep this top-level instruction.", + "input": [], + "tools": [{"type": "function", "name": "lookup", "parameters": {}}], + "parallel_tool_calls": true + }); + + finalize_openai_provider_request( + &mut body, + OpenAiProviderRequestFinalization { + source_api_format: "openai:responses", + provider_api_format: "openai:responses", + provider_type: "codex", + provider_model: "gpt-5.4", + source_model: "gpt-5.4", + body_rules: None, + upstream_is_stream: true, + require_body_stream_field: true, + }, + ) + .expect("GPT-5.4 should satisfy the standard Codex Responses contract"); + + assert_eq!(body["instructions"], "Keep this top-level instruction."); + assert_eq!(body["tools"][0]["name"], "lookup"); + assert_eq!(body["parallel_tool_calls"], true); + assert_eq!(body["reasoning"]["effort"], "medium"); + assert!(body["reasoning"].get("context").is_none()); + assert!(body["reasoning"].get("summary").is_none()); + } + + #[test] + fn cross_format_compact_finalization_removes_post_conversion_fields() { + for source_api_format in ["claude:messages", "gemini:generate_content"] { + let mut body = json!({ + "model": "gpt-5.6-sol", + "input": [], + "client_metadata": {"source": "mapping"}, + "include": ["reasoning.encrypted_content"], + "store": true, + "stream": true, + "stream_options": {"include_usage": true}, + "tool_choice": "auto", + "parallel_tool_calls": true, + "reasoning": {"effort": "max"}, + "text": {"verbosity": "medium"}, + "tools": [{"type": "function", "name": "lookup", "parameters": {}}] + }); + finalize_openai_provider_request( + &mut body, + OpenAiProviderRequestFinalization { + source_api_format, + provider_api_format: "openai:responses:compact", + provider_type: "codex", + provider_model: "gpt-5.6-sol", + source_model: "gpt-5.6-sol", + body_rules: None, + upstream_is_stream: false, + require_body_stream_field: true, + }, + ) + .expect("cross-format Compact request should satisfy its final contract"); + + for field in [ + "client_metadata", + "include", + "store", + "stream", + "stream_options", + "tool_choice", + ] { + assert!(body.get(field).is_none(), "{field} must not reach Compact"); + } + assert_eq!(body["parallel_tool_calls"], false); + assert_eq!(body["reasoning"]["effort"], "max"); + assert_eq!(body["reasoning"]["context"], "all_turns"); + assert_eq!(body["text"]["verbosity"], "medium"); + assert!(body.get("tools").is_none()); + assert_eq!(body["input"][0]["tools"][0]["name"], "lookup"); + } + } +} diff --git a/crates/aether-ai-formats/src/formats/openai/responses/codex.rs b/crates/aether-ai-formats/src/formats/openai/responses/codex.rs index d50a165bc..116b978a7 100644 --- a/crates/aether-ai-formats/src/formats/openai/responses/codex.rs +++ b/crates/aether-ai-formats/src/formats/openai/responses/codex.rs @@ -1,17 +1,14 @@ use std::collections::BTreeMap; -use std::fmt::Write; +use std::sync::OnceLock; use aether_ai_formats::provider_compat::proxy::rules::body_rules_handle_path; +use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; -use sha1::{Digest as Sha1Digest, Sha1}; -use sha2::Sha256; -use uuid::Uuid; -const CODEX_PROMPT_CACHE_NAMESPACE_VERSION: &str = "v3"; -const CODEX_DEFAULT_INSTRUCTIONS: &str = ""; const CODEX_DEFAULT_REASONING_EFFORT: &str = "medium"; -const CODEX_DEFAULT_REASONING_SUMMARY: &str = "auto"; const CODEX_REASONING_ENCRYPTED_CONTENT_INCLUDE: &str = "reasoning.encrypted_content"; +pub const CODEX_RESPONSES_LITE_HEADER: &str = "x-openai-internal-codex-responses-lite"; +pub const CODEX_MODEL_CATALOG_METADATA_FIELD: &str = "codex_models"; const CODEX_OPENAI_RESPONSES_UNSUPPORTED_BODY_FIELDS: &[&str] = &[ "max_output_tokens", "max_completion_tokens", @@ -21,14 +18,25 @@ const CODEX_OPENAI_RESPONSES_UNSUPPORTED_BODY_FIELDS: &[&str] = &[ "presence_penalty", "user", "metadata", + "prompt_cache_options", "prompt_cache_retention", "safety_identifier", - "stream_options", "previous_response_id", ]; -const CODEX_DEFAULT_USER_AGENT: &str = - "codex-tui/0.122.0 (Mac OS 15.2.0; arm64) vscode/2.6.11 (codex-tui; 0.122.0)"; -const CODEX_DEFAULT_ORIGINATOR: &str = "codex-tui"; +const CODEX_OPENAI_RESPONSES_COMPACT_BODY_FIELDS: &[&str] = &[ + "model", + "input", + "instructions", + "tools", + "parallel_tool_calls", + "reasoning", + "service_tier", + "prompt_cache_key", + "text", +]; +pub const CODEX_CLIENT_VERSION: &str = "0.144.1"; +pub const CODEX_CLIENT_USER_AGENT: &str = "codex_cli_rs/0.144.1"; +pub const CODEX_CLIENT_ORIGINATOR: &str = "codex_cli_rs"; pub const CODEX_OPENAI_IMAGE_INTERNAL_MODEL: &str = "gpt-5.4-mini"; pub const CODEX_OPENAI_IMAGE_DEFAULT_MODEL: &str = "gpt-image-2"; pub const CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_MODEL: &str = "dall-e-2"; @@ -38,14 +46,9 @@ pub const CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_PROMPT: &str = const CODEX_IMAGE_TOOL_DEFAULT_SIZE: &str = "1024x1024"; const CODEX_IMAGE_TOOL_DEFAULT_QUALITY: &str = "high"; const CODEX_IMAGE_TOOL_DEFAULT_BACKGROUND: &str = "auto"; -const UUID_NAMESPACE_OID_BYTES: [u8; 16] = [ - 0x6b, 0xa7, 0xb8, 0x12, 0x9d, 0xad, 0x11, 0xd1, 0x80, 0xb4, 0x00, 0xc0, 0x4f, 0xd4, 0x30, 0xc8, -]; - fn is_codex_openai_responses_request(provider_type: &str, provider_api_format: &str) -> bool { provider_type.trim().eq_ignore_ascii_case("codex") - && (aether_ai_formats::is_openai_responses_family_format(provider_api_format) - || is_openai_image_request(provider_api_format)) + && aether_ai_formats::is_openai_responses_family_format(provider_api_format) } fn is_openai_responses_compact_request(provider_api_format: &str) -> bool { @@ -58,16 +61,558 @@ fn is_openai_image_request(provider_api_format: &str) -> bool { .eq_ignore_ascii_case("openai:image") } -/// Returns true only when `tool_choice` *explicitly* targets image_generation, -/// matching either the `"image_generation"` string form or the -/// `{"type":"image_generation"}` object form. -/// -/// Intentionally does NOT inspect the `tools` array: tools merely advertise -/// what is available, while only `tool_choice` expresses the caller's -/// selection. Treating "image_generation present in tools" as a trigger -/// caused codex CLI requests (which list image_generation alongside ~20 -/// other tools under `tool_choice: "auto"`) to be incorrectly rewritten -/// into image-generation-only requests, leading to upstream 400s. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum CodexOpenAiEndpointKind { + Responses, + Compact, + Search, + Images, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct CodexResponsesModelCapabilities { + pub use_responses_lite: bool, + pub supports_reasoning_summary_parameter: bool, + pub default_reasoning_effort: Option, + pub default_reasoning_summary: Option, + pub supported_reasoning_efforts: Vec, + pub supports_parallel_tool_calls: bool, + pub support_verbosity: bool, + pub default_verbosity: Option, + pub supported_service_tiers: Vec, +} + +impl CodexResponsesModelCapabilities { + pub fn supports_reasoning_effort(&self, effort: &str) -> bool { + self.supported_reasoning_efforts + .iter() + .any(|candidate| candidate == effort.trim()) + } + + fn supports_service_tier(&self, service_tier: &str) -> bool { + self.supported_service_tiers + .iter() + .any(|candidate| candidate == service_tier) + } +} + +fn codex_namespaced_model_suffix(model: &str) -> Option<&str> { + let (namespace, suffix) = model.split_once('/')?; + if suffix.contains('/') + || namespace.is_empty() + || !namespace + .chars() + .all(|character| character.is_ascii_alphanumeric() || matches!(character, '_' | '-')) + { + return None; + } + Some(suffix) +} + +fn conservative_codex_responses_model_capabilities() -> CodexResponsesModelCapabilities { + CodexResponsesModelCapabilities { + use_responses_lite: false, + supports_reasoning_summary_parameter: true, + default_reasoning_effort: None, + default_reasoning_summary: None, + supported_reasoning_efforts: Vec::new(), + supports_parallel_tool_calls: false, + support_verbosity: false, + default_verbosity: None, + supported_service_tiers: Vec::new(), + } +} + +pub fn codex_responses_model_capabilities_from_card( + card: &Value, +) -> Option<(String, CodexResponsesModelCapabilities)> { + let card = card.as_object()?; + let model_id = card + .get("slug") + .or_else(|| card.get("id")) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty())?; + let mut capabilities = conservative_codex_responses_model_capabilities(); + if let Some(value) = card.get("use_responses_lite").and_then(Value::as_bool) { + capabilities.use_responses_lite = value; + } + capabilities.supports_reasoning_summary_parameter = card + .get("supports_reasoning_summary_parameter") + .and_then(Value::as_bool) + .unwrap_or(true); + if let Some(value) = card + .get("default_reasoning_level") + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + { + capabilities.default_reasoning_effort = Some(value.to_string()); + } + if let Some(value) = card + .get("default_reasoning_summary") + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + { + capabilities.default_reasoning_summary = + (!value.eq_ignore_ascii_case("none")).then(|| value.to_ascii_lowercase()); + } + if let Some(levels) = card + .get("supported_reasoning_levels") + .and_then(Value::as_array) + { + capabilities.supported_reasoning_efforts = levels + .iter() + .filter_map(|level| { + level + .get("effort") + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_string) + }) + .collect(); + } + if let Some(value) = card + .get("supports_parallel_tool_calls") + .and_then(Value::as_bool) + { + capabilities.supports_parallel_tool_calls = value; + } + if let Some(value) = card.get("support_verbosity").and_then(Value::as_bool) { + capabilities.support_verbosity = value; + } + if let Some(value) = card + .get("default_verbosity") + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + { + capabilities.default_verbosity = Some(value.to_ascii_lowercase()); + } + if let Some(service_tiers) = card.get("service_tiers").and_then(Value::as_array) { + capabilities.supported_service_tiers = service_tiers + .iter() + .filter_map(|tier| { + tier.get("id") + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_string) + }) + .collect(); + } + Some((model_id.to_string(), capabilities)) +} + +pub fn build_codex_model_catalog_metadata(cards: &[Value]) -> Value { + let cards = cards + .iter() + .filter_map(|card| { + let model_id = card + .get("slug") + .or_else(|| card.get("id")) + .and_then(Value::as_str)?; + Some((model_id.to_string(), codex_execution_model_card(card))) + }) + .collect::>(); + json!({ + CODEX_MODEL_CATALOG_METADATA_FIELD: { + "cards": cards, + }, + }) +} + +fn codex_execution_model_card(card: &Value) -> Value { + const EXCLUDED_FIELDS: [&str; 3] = + ["base_instructions", "model_messages", "available_in_plans"]; + let Some(card) = card.as_object() else { + return card.clone(); + }; + Value::Object( + card.iter() + .filter(|(key, _)| !EXCLUDED_FIELDS.contains(&key.as_str())) + .map(|(key, value)| (key.clone(), value.clone())) + .collect(), + ) +} + +pub fn effective_codex_model_cards(remote_cards: &[Value]) -> Vec { + if remote_cards.iter().any(|card| { + card.get("visibility") + .and_then(Value::as_str) + .is_some_and(|visibility| visibility == "list") + }) { + return remote_cards.to_vec(); + } + + let mut cards = bundled_codex_model_cards().to_vec(); + for remote_card in remote_cards { + let Some(remote_id) = remote_card + .get("slug") + .or_else(|| remote_card.get("id")) + .and_then(Value::as_str) + else { + continue; + }; + if let Some(index) = cards.iter().position(|card| { + card.get("slug") + .or_else(|| card.get("id")) + .and_then(Value::as_str) + .is_some_and(|model_id| model_id == remote_id) + }) { + let Some(base) = cards[index].as_object() else { + cards[index] = remote_card.clone(); + continue; + }; + let Some(overlay) = remote_card.as_object() else { + cards[index] = remote_card.clone(); + continue; + }; + let mut merged = base.clone(); + merged.extend(overlay.clone()); + cards[index] = Value::Object(merged); + } else { + cards.push(remote_card.clone()); + } + } + cards +} + +fn reasoning_level_cards(efforts: &[&str]) -> Vec { + efforts + .iter() + .map(|effort| { + let description = match *effort { + "low" => "Fast responses with lighter reasoning", + "medium" => "Balances speed and reasoning depth for everyday tasks", + "high" => "Greater reasoning depth for complex problems", + "xhigh" => "Extra high reasoning depth for complex problems", + "max" => "Maximum reasoning depth for the hardest problems", + "ultra" => "Maximum reasoning with automatic task delegation", + _ => "", + }; + json!({ "effort": effort, "description": description }) + }) + .collect() +} + +struct BundledCodexModelCardSpec<'a> { + model_id: &'a str, + display_name: &'a str, + description: &'a str, + default_reasoning_level: &'a str, + default_reasoning_summary: &'a str, + use_responses_lite: bool, + efforts: &'a [&'a str], + default_verbosity: &'a str, + supports_priority_tier: bool, +} + +fn bundled_codex_model_card(spec: BundledCodexModelCardSpec<'_>) -> Value { + let service_tiers = spec + .supports_priority_tier + .then(|| { + json!({ + "id": "priority", + "name": "Fast", + "description": "1.5x speed, increased usage", + }) + }) + .into_iter() + .collect::>(); + json!({ + "id": spec.model_id, + "slug": spec.model_id, + "object": "model", + "owned_by": "openai", + "display_name": spec.display_name, + "description": spec.description, + "api_formats": ["openai:responses"], + "default_reasoning_level": spec.default_reasoning_level, + "supported_reasoning_levels": reasoning_level_cards(spec.efforts), + "default_reasoning_summary": spec.default_reasoning_summary, + "support_verbosity": true, + "default_verbosity": spec.default_verbosity, + "supports_parallel_tool_calls": true, + "service_tiers": service_tiers, + "use_responses_lite": spec.use_responses_lite, + }) +} + +fn bundled_gpt_5_6_codex_model_card( + model_id: &str, + display_name: &str, + description: &str, + default_reasoning_level: &str, + priority: u64, + multi_agent_version: &str, + supports_ultra: bool, +) -> Value { + let mut efforts = vec!["low", "medium", "high", "xhigh", "max"]; + if supports_ultra { + efforts.push("ultra"); + } + let mut card = bundled_codex_model_card(BundledCodexModelCardSpec { + model_id, + display_name, + description, + default_reasoning_level, + default_reasoning_summary: "none", + use_responses_lite: true, + efforts: &efforts, + default_verbosity: "low", + supports_priority_tier: true, + }); + let object = card + .as_object_mut() + .expect("bundled Codex model card must be an object"); + object.extend( + json!({ + "shell_type": "shell_command", + "supports_image_detail_original": true, + "supports_search_tool": true, + "input_modalities": ["text", "image"], + "context_window": 372_000, + "max_context_window": 372_000, + "comp_hash": "3000", + "experimental_supported_tools": [], + "visibility": "list", + "supported_in_api": true, + "priority": priority, + "additional_speed_tiers": ["fast"], + "multi_agent_version": multi_agent_version, + "tool_mode": "code_mode_only", + "prefer_websockets": true, + "reasoning_summary_format": "experimental", + "include_skills_usage_instructions": false, + "apply_patch_tool_type": "freeform", + "web_search_tool_type": "text_and_image", + "truncation_policy": { "mode": "tokens", "limit": 10_000 }, + "minimal_client_version": "0.144.0", + }) + .as_object() + .expect("bundled Codex model card extension must be an object") + .clone(), + ); + card +} + +fn bundled_codex_auto_review_model_card() -> Value { + let mut card = bundled_codex_model_card(BundledCodexModelCardSpec { + model_id: "codex-auto-review", + display_name: "Codex Auto Review", + description: "Automatic approval review model for Codex.", + default_reasoning_level: "medium", + default_reasoning_summary: "none", + use_responses_lite: false, + efforts: &["low", "medium", "high", "xhigh"], + default_verbosity: "low", + supports_priority_tier: false, + }); + let object = card + .as_object_mut() + .expect("bundled Codex model card must be an object"); + object.extend( + json!({ + "shell_type": "shell_command", + "supports_image_detail_original": true, + "supports_search_tool": true, + "input_modalities": ["text", "image"], + "context_window": 272_000, + "max_context_window": 1_000_000, + "experimental_supported_tools": [], + "visibility": "hide", + "supported_in_api": true, + "priority": 43, + "additional_speed_tiers": [], + "prefer_websockets": true, + "reasoning_summary_format": "experimental", + "include_skills_usage_instructions": false, + "apply_patch_tool_type": "freeform", + "web_search_tool_type": "text_and_image", + "truncation_policy": { "mode": "tokens", "limit": 10_000 }, + "minimal_client_version": "0.98.0", + }) + .as_object() + .expect("bundled Codex model card extension must be an object") + .clone(), + ); + card +} + +pub fn bundled_codex_model_cards() -> &'static [Value] { + static CARDS: OnceLock> = OnceLock::new(); + CARDS.get_or_init(|| { + vec![ + bundled_gpt_5_6_codex_model_card( + "gpt-5.6-sol", + "GPT-5.6-Sol", + "Latest frontier agentic coding model.", + "low", + 1, + "v2", + true, + ), + bundled_gpt_5_6_codex_model_card( + "gpt-5.6-terra", + "GPT-5.6-Terra", + "Balanced agentic coding model for everyday work.", + "medium", + 2, + "v2", + true, + ), + bundled_gpt_5_6_codex_model_card( + "gpt-5.6-luna", + "GPT-5.6-Luna", + "Fast and affordable agentic coding model.", + "medium", + 3, + "v1", + false, + ), + bundled_codex_model_card(BundledCodexModelCardSpec { + model_id: "gpt-5.5", + display_name: "GPT-5.5", + description: "Frontier model for complex coding, research, and real-world work.", + default_reasoning_level: "medium", + default_reasoning_summary: "none", + use_responses_lite: false, + efforts: &["low", "medium", "high", "xhigh"], + default_verbosity: "low", + supports_priority_tier: true, + }), + bundled_codex_model_card(BundledCodexModelCardSpec { + model_id: "gpt-5.4", + display_name: "GPT-5.4", + description: "Strong model for everyday coding.", + default_reasoning_level: "medium", + default_reasoning_summary: "none", + use_responses_lite: false, + efforts: &["low", "medium", "high", "xhigh"], + default_verbosity: "low", + supports_priority_tier: true, + }), + bundled_codex_model_card(BundledCodexModelCardSpec { + model_id: "gpt-5.4-mini", + display_name: "GPT-5.4 Mini", + description: "Small, fast, and cost-efficient model for simpler coding tasks.", + default_reasoning_level: "medium", + default_reasoning_summary: "none", + use_responses_lite: false, + efforts: &["low", "medium", "high", "xhigh"], + default_verbosity: "medium", + supports_priority_tier: false, + }), + bundled_codex_model_card(BundledCodexModelCardSpec { + model_id: "gpt-5.2", + display_name: "GPT-5.2", + description: "Optimized for professional work and long-running agents.", + default_reasoning_level: "medium", + default_reasoning_summary: "auto", + use_responses_lite: false, + efforts: &["low", "medium", "high", "xhigh"], + default_verbosity: "low", + supports_priority_tier: false, + }), + bundled_codex_auto_review_model_card(), + ] + }) +} + +fn bundled_codex_model_card_for_identity(model_id: &str) -> Option<&'static Value> { + let find = |model: &str| { + bundled_codex_model_cards() + .iter() + .filter_map(|card| { + let slug = card.get("slug").and_then(Value::as_str)?; + model.starts_with(slug).then_some((slug.len(), card)) + }) + .max_by_key(|(slug_len, _)| *slug_len) + .map(|(_, card)| card) + }; + find(model_id).or_else(|| codex_namespaced_model_suffix(model_id).and_then(find)) +} + +fn catalog_codex_model_card_for_identity<'a>( + cards: &'a serde_json::Map, + model_id: &str, +) -> Option<&'a Value> { + let find = |model: &str| { + cards + .iter() + .filter_map(|(catalog_id, card)| { + model + .starts_with(catalog_id) + .then_some((catalog_id.len(), card)) + }) + .max_by_key(|(slug_len, _)| *slug_len) + .map(|(_, card)| card) + }; + find(model_id).or_else(|| codex_namespaced_model_suffix(model_id).and_then(find)) +} + +pub fn resolve_codex_responses_model_capabilities( + provider_model: &str, + source_model: &str, + upstream_metadata: Option<&Value>, +) -> CodexResponsesModelCapabilities { + let provider_model = provider_model.trim(); + let source_model = source_model.trim(); + let source_fallback = (provider_model.is_empty() + || crate::formats::shared::model_directives::openai_model_capability_is_opaque( + provider_model, + provider_model, + )) + .then_some(source_model) + .filter(|model| !model.is_empty() && *model != provider_model); + if let Some(cards) = upstream_metadata + .and_then(|metadata| metadata.get(CODEX_MODEL_CATALOG_METADATA_FIELD)) + .and_then(|catalog| catalog.get("cards")) + .and_then(Value::as_object) + { + if !cards.is_empty() { + return catalog_codex_model_card_for_identity(cards, provider_model) + .or_else(|| { + source_fallback + .and_then(|model| catalog_codex_model_card_for_identity(cards, model)) + }) + .and_then(codex_responses_model_capabilities_from_card) + .map(|(_, capabilities)| capabilities) + .unwrap_or_else(conservative_codex_responses_model_capabilities); + } + } + bundled_codex_model_card_for_identity(provider_model) + .or_else(|| source_fallback.and_then(bundled_codex_model_card_for_identity)) + .and_then(codex_responses_model_capabilities_from_card) + .map(|(_, capabilities)| capabilities) + .unwrap_or_else(conservative_codex_responses_model_capabilities) +} + +fn codex_openai_endpoint_kind( + provider_type: &str, + provider_api_format: &str, +) -> Option { + if !provider_type.trim().eq_ignore_ascii_case("codex") { + return None; + } + if is_openai_responses_compact_request(provider_api_format) { + Some(CodexOpenAiEndpointKind::Compact) + } else if aether_ai_formats::is_openai_responses_format(provider_api_format) { + Some(CodexOpenAiEndpointKind::Responses) + } else if aether_ai_formats::api_format_alias_matches(provider_api_format, "openai:search") { + Some(CodexOpenAiEndpointKind::Search) + } else if is_openai_image_request(provider_api_format) { + Some(CodexOpenAiEndpointKind::Images) + } else { + None + } +} + +/// Matches an explicit image-generation selection in `tool_choice`. +/// The `tools` array describes availability and does not select a tool. fn codex_openai_responses_tool_choice_references_image_generation( body_object: &serde_json::Map, ) -> bool { @@ -179,440 +724,137 @@ fn inject_codex_default_variation_prompt(body_object: &mut serde_json::Map Option { - let normalized = seed.trim(); - if normalized.is_empty() { - return None; - } - - let normalized_kind = kind - .trim() - .to_ascii_lowercase() - .chars() - .filter(|ch| ch.is_ascii_alphanumeric() || *ch == '_' || *ch == '-') - .collect::(); - let normalized_kind = if normalized_kind.is_empty() { - "seed".to_string() - } else { - normalized_kind - }; - let namespace = format!( - "aether:codex:prompt-cache:{CODEX_PROMPT_CACHE_NAMESPACE_VERSION}:{normalized_kind}:{normalized}" - ); - let mut hasher = Sha1::new(); - hasher.update(UUID_NAMESPACE_OID_BYTES); - hasher.update(namespace.as_bytes()); - - let digest = hasher.finalize(); - let mut bytes = [0u8; 16]; - bytes.copy_from_slice(&digest[..16]); - bytes[6] = (bytes[6] & 0x0f) | 0x50; - bytes[8] = (bytes[8] & 0x3f) | 0x80; - Some(Uuid::from_bytes(bytes).to_string()) -} - -fn build_stable_codex_prompt_cache_key(user_api_key_id: &str) -> Option { - build_stable_codex_prompt_cache_key_from_seed("user", user_api_key_id) -} - -fn extract_codex_prompt_cache_session_seed(provider_request_body: &Value) -> Option { - fn non_empty_str(value: Option<&Value>) -> Option<&str> { - value - .and_then(Value::as_str) - .map(str::trim) - .filter(|value| !value.is_empty()) - } - - fn session_seed_from_metadata(metadata: &Value) -> Option { - let object = metadata.as_object()?; - non_empty_str(object.get("session_id")) - .or_else(|| non_empty_str(object.get("sessionId"))) - .or_else(|| non_empty_str(object.get("conversation_id"))) - .or_else(|| non_empty_str(object.get("conversationId"))) - .map(|value| format!("metadata:{value}")) - .or_else(|| { - let user_id = non_empty_str(object.get("user_id"))?; - serde_json::from_str::(user_id) - .ok() - .and_then(|decoded| { - non_empty_str(decoded.get("session_id")) - .or_else(|| non_empty_str(decoded.get("sessionId"))) - .or_else(|| non_empty_str(decoded.get("conversation_id"))) - .or_else(|| non_empty_str(decoded.get("conversationId"))) - .map(|value| format!("metadata.user_id:{value}")) - }) - }) - } - - let object = provider_request_body.as_object()?; - non_empty_str(object.get("session_id")) - .or_else(|| non_empty_str(object.get("sessionId"))) - .or_else(|| non_empty_str(object.get("conversation_id"))) - .or_else(|| non_empty_str(object.get("conversationId"))) - .map(|value| format!("body:{value}")) - .or_else(|| object.get("metadata").and_then(session_seed_from_metadata)) -} - -fn sha256_hex(input: &[u8]) -> String { - let digest = Sha256::digest(input); - let mut output = String::with_capacity(digest.len() * 2); - for byte in digest { - let _ = write!(&mut output, "{byte:02x}"); - } - output -} - -fn stable_json_digest(value: &Value) -> Option { - serde_json::to_vec(value) - .ok() - .map(|serialized| sha256_hex(&serialized)) -} - -fn compact_prompt_cache_text(value: &str) -> Option { - const MAX_PROMPT_CACHE_TEXT_CHARS: usize = 4096; - let normalized = value.trim(); - if normalized.is_empty() { - return None; - } - let mut text = normalized - .chars() - .take(MAX_PROMPT_CACHE_TEXT_CHARS) - .collect::(); - if normalized.chars().count() > MAX_PROMPT_CACHE_TEXT_CHARS { - text.push_str("..."); - } - Some(Value::String(text)) -} - -fn compact_prompt_cache_anchor(value: &Value) -> Value { - match value { - Value::String(text) => compact_prompt_cache_text(text).unwrap_or(Value::Null), - Value::Array(items) => Value::Array( - items - .iter() - .take(16) - .map(compact_prompt_cache_anchor) - .filter(|value| !value.is_null()) - .collect(), - ), - Value::Object(object) => { - let mut compacted = serde_json::Map::new(); - for key in [ - "type", - "role", - "id", - "name", - "description", - "text", - "input_text", - "output_text", - "content", - "call_id", - "arguments", - "output", - "parameters", - "strict", - "function", - "effort", - "summary", - ] { - let Some(value) = object.get(key) else { - continue; - }; - let value = compact_prompt_cache_anchor(value); - if !value.is_null() { - compacted.insert(key.to_string(), value); +fn strip_codex_content_cache_control_fields(content: &mut Value) { + match content { + Value::Array(parts) => { + for part in parts { + if let Some(part) = part.as_object_mut() { + part.remove("cache_control"); } } - Value::Object(compacted) - } - Value::Null | Value::Bool(_) | Value::Number(_) => value.clone(), - } -} - -fn compact_prompt_cache_json_anchor(value: &Value) -> Value { - match value { - Value::String(text) => compact_prompt_cache_text(text).unwrap_or(Value::Null), - Value::Array(items) => Value::Array( - items - .iter() - .take(16) - .map(compact_prompt_cache_json_anchor) - .filter(|value| !value.is_null()) - .collect(), - ), - Value::Object(object) => { - let mut compacted = serde_json::Map::new(); - let mut keys = object.keys().collect::>(); - keys.sort(); - for key in keys { - if key == "cache_control" { - continue; - } - let Some(value) = object.get(key) else { - continue; - }; - let value = compact_prompt_cache_json_anchor(value); - if !value.is_null() { - compacted.insert(key.clone(), value); - } - } - Value::Object(compacted) - } - Value::Null | Value::Bool(_) | Value::Number(_) => value.clone(), - } -} - -fn collect_codex_prompt_cache_control_anchors(value: &Value, anchors: &mut Vec) { - const MAX_PROMPT_CACHE_CONTROL_ANCHORS: usize = 16; - if anchors.len() >= MAX_PROMPT_CACHE_CONTROL_ANCHORS { - return; - } - - match value { - Value::Object(object) => { - if object.contains_key("cache_control") { - let mut anchor = object.clone(); - anchor.remove("cache_control"); - let anchor = compact_prompt_cache_anchor(&Value::Object(anchor)); - if !anchor.is_null() { - anchors.push(anchor); - } - } - for child in object.values() { - if anchors.len() >= MAX_PROMPT_CACHE_CONTROL_ANCHORS { - break; - } - collect_codex_prompt_cache_control_anchors(child, anchors); - } } - Value::Array(items) => { - for child in items { - if anchors.len() >= MAX_PROMPT_CACHE_CONTROL_ANCHORS { - break; - } - collect_codex_prompt_cache_control_anchors(child, anchors); - } + Value::Object(part) => { + part.remove("cache_control"); } _ => {} } } +fn strip_codex_tool_cache_control_fields(tools: &mut Value) { + let Some(tools) = tools.as_array_mut() else { + return; + }; + for tool in tools { + let Some(tool) = tool.as_object_mut() else { + continue; + }; + tool.remove("cache_control"); + if let Some(function) = tool.get_mut("function").and_then(Value::as_object_mut) { + function.remove("cache_control"); + } + } +} + fn strip_codex_cache_control_fields(value: &mut Value) { - match value { - Value::Object(object) => { - object.remove("cache_control"); - for child in object.values_mut() { - strip_codex_cache_control_fields(child); - } - } - Value::Array(items) => { - for child in items { - strip_codex_cache_control_fields(child); - } - } - _ => {} + let Some(body) = value.as_object_mut() else { + return; + }; + body.remove("cache_control"); + if let Some(tools) = body.get_mut("tools") { + strip_codex_tool_cache_control_fields(tools); } -} - -fn extract_codex_prompt_cache_control_seed(provider_request_body: &Value) -> Option { - let mut anchors = Vec::new(); - collect_codex_prompt_cache_control_anchors(provider_request_body, &mut anchors); - if anchors.is_empty() { - return None; - } - - let seed = json!({ - "model": provider_request_body.get("model"), - "anchors": anchors, - }); - stable_json_digest(&seed).map(|digest| format!("cache_control:{digest}")) -} - -fn first_responses_input_anchor(input: &Value) -> Option { - let items = input.as_array()?; - let first_user_message = items.iter().find(|item| { - item.get("type") - .and_then(Value::as_str) - .is_some_and(|value| value == "message") - && item - .get("role") - .and_then(Value::as_str) - .is_some_and(|value| value == "user") - }); - let first_item = first_user_message.or_else(|| items.first())?; - let anchor = compact_prompt_cache_anchor(first_item); - (!anchor.is_null()).then_some(anchor) -} - -fn extract_codex_stable_request_prompt_cache_seed( - provider_request_body: &Value, - user_api_key_id: Option<&str>, -) -> Option { - let object = provider_request_body.as_object()?; - let mut seed = serde_json::Map::new(); - - for key in [ - "model", - "instructions", - "reasoning", - "tools", - "tool_choice", - "parallel_tool_calls", - ] { - if let Some(value) = object.get(key).filter(|value| !value.is_null()) { - let value = if key == "tools" { - compact_prompt_cache_json_anchor(value) - } else { - compact_prompt_cache_anchor(value) + if let Some(input) = body.get_mut("input").and_then(Value::as_array_mut) { + for item in input { + let Some(item) = item.as_object_mut() else { + continue; }; - seed.insert(key.to_string(), value); + item.remove("cache_control"); + if let Some(content) = item.get_mut("content") { + strip_codex_content_cache_control_fields(content); + } + if item + .get("type") + .and_then(Value::as_str) + .is_some_and(|item_type| item_type == "additional_tools") + { + if let Some(tools) = item.get_mut("tools") { + strip_codex_tool_cache_control_fields(tools); + } + } } } - if let Some(input_anchor) = object.get("input").and_then(first_responses_input_anchor) { - seed.insert("first_input".to_string(), input_anchor); - } - if let Some(user_api_key_id) = user_api_key_id +} + +fn remove_btree_header(headers: &mut BTreeMap, header_name: &str) { + headers.retain(|name, _| !name.trim().eq_ignore_ascii_case(header_name)); +} + +fn header_value_contains_media_type(value: &str, media_type: &str) -> bool { + value.split(',').any(|media_range| { + media_range + .split(';') + .next() + .map(str::trim) + .is_some_and(|value| value.eq_ignore_ascii_case(media_type)) + }) +} + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct CodexAuthIdentity { + pub account_id: Option, + pub is_fedramp: bool, + pub uses_codex_backend: bool, +} + +pub fn parse_codex_auth_identity(decrypted_auth_config_raw: Option<&str>) -> CodexAuthIdentity { + let Some(raw) = decrypted_auth_config_raw .map(str::trim) .filter(|value| !value.is_empty()) - { - seed.insert( - "api_key_id".to_string(), - Value::String(user_api_key_id.to_string()), - ); - } - - if seed.len() < 2 { - return None; - } - stable_json_digest(&Value::Object(seed)).map(|digest| format!("stable_request:{digest}")) -} - -fn build_short_codex_header_id(seed: &str) -> Option { - let normalized = seed.trim(); - if normalized.is_empty() { - return None; - } - - let digest = Sha256::digest(normalized.as_bytes()); - let mut short_id = String::with_capacity(16); - for byte in digest.iter().take(8) { - let _ = write!(&mut short_id, "{byte:02x}"); - } - Some(short_id) -} - -fn header_map_has_non_empty_value(headers: &http::HeaderMap, header_name: &str) -> bool { - let target = header_name.trim().to_ascii_lowercase(); - if target.is_empty() { - return false; - } - - headers.iter().any(|(name, value)| { - if name.as_str().trim().to_ascii_lowercase() != target { - return false; - } - value - .to_str() - .ok() - .map(str::trim) - .map(|value| !value.is_empty()) - .unwrap_or(false) - }) -} - -fn btree_map_has_non_empty_value(headers: &BTreeMap, header_name: &str) -> bool { - let target = header_name.trim().to_ascii_lowercase(); - if target.is_empty() { - return false; - } - - headers - .iter() - .any(|(name, value)| name.trim().eq_ignore_ascii_case(&target) && !value.trim().is_empty()) -} - -fn extract_codex_account_id(decrypted_auth_config_raw: Option<&str>) -> Option { - let raw = decrypted_auth_config_raw?.trim(); - if raw.is_empty() { - return None; - } - - serde_json::from_str::(raw).ok().and_then(|value| { - value - .get("account_id") + else { + return CodexAuthIdentity::default(); + }; + let Ok(value) = serde_json::from_str::(raw) else { + return CodexAuthIdentity::default(); + }; + let namespaced_auth = value + .get("https://api.openai.com/auth") + .and_then(Value::as_object); + let account_id = value + .get("account_id") + .or_else(|| value.get("chatgpt_account_id")) + .or_else(|| namespaced_auth.and_then(|auth| auth.get("chatgpt_account_id"))) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned); + let is_fedramp = value + .get("is_fedramp") + .or_else(|| value.get("chatgpt_account_is_fedramp")) + .or_else(|| namespaced_auth.and_then(|auth| auth.get("chatgpt_account_is_fedramp"))) + .and_then(Value::as_bool) + .unwrap_or(false); + let uses_codex_backend = account_id.is_some() + || value + .get("provider_type") .and_then(Value::as_str) - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(ToOwned::to_owned) - }) + .is_some_and(|provider_type| provider_type.trim().eq_ignore_ascii_case("codex")); + + CodexAuthIdentity { + account_id, + is_fedramp, + uses_codex_backend, + } } -fn maybe_insert_default_codex_header( +fn set_codex_client_header( provider_request_headers: &mut BTreeMap, - original_headers: &http::HeaderMap, header_name: &str, header_value: &str, ) { - if header_map_has_non_empty_value(original_headers, header_name) - || btree_map_has_non_empty_value(provider_request_headers, header_name) - { - return; - } - + remove_btree_header(provider_request_headers, header_name); provider_request_headers.insert(header_name.to_string(), header_value.to_string()); } -fn codex_prompt_cache_key_to_insert( - provider_request_body: &Value, - provider_type: &str, - provider_api_format: &str, - user_api_key_id: Option<&str>, -) -> Option { - if !is_codex_openai_responses_request(provider_type, provider_api_format) { - return None; - } - - let existing = provider_request_body - .get("prompt_cache_key") - .and_then(Value::as_str) - .map(str::trim) - .unwrap_or_default(); - if !existing.is_empty() { - return None; - } - - extract_codex_prompt_cache_session_seed(provider_request_body) - .and_then(|seed| build_stable_codex_prompt_cache_key_from_seed("session", &seed)) - .or_else(|| { - extract_codex_prompt_cache_control_seed(provider_request_body) - .and_then(|seed| build_stable_codex_prompt_cache_key_from_seed("anchor", &seed)) - }) - .or_else(|| { - extract_codex_stable_request_prompt_cache_seed(provider_request_body, user_api_key_id) - .and_then(|seed| build_stable_codex_prompt_cache_key_from_seed("request", &seed)) - }) - .or_else(|| user_api_key_id.and_then(build_stable_codex_prompt_cache_key)) -} - -fn insert_codex_prompt_cache_key( - provider_request_body: &mut Value, - prompt_cache_key: Option, -) { - let Some(prompt_cache_key) = prompt_cache_key else { - return; - }; - - let Some(body_object) = provider_request_body.as_object_mut() else { - return; - }; - - body_object.insert( - "prompt_cache_key".to_string(), - Value::String(prompt_cache_key), - ); -} - pub fn apply_openai_responses_compact_special_body_edits( provider_request_body: &mut Value, provider_api_format: &str, @@ -624,81 +866,529 @@ pub fn apply_openai_responses_compact_special_body_edits( let Some(body_object) = provider_request_body.as_object_mut() else { return; }; - - // `/v1/responses/compact` does not accept `include`, `store`, or body-level `stream`. - body_object.remove("include"); - body_object.remove("store"); - body_object.remove("stream"); + super::request::apply_compact_request_projection(body_object); } -fn ensure_codex_responses_passthrough_fields( - body_object: &mut serde_json::Map, +pub fn apply_codex_openai_responses_compact_body_edits( + provider_request_body: &mut Value, + provider_type: &str, provider_api_format: &str, - body_rules: Option<&Value>, ) { - if is_openai_responses_compact_request(provider_api_format) - || is_openai_image_request(provider_api_format) + if !is_codex_openai_responses_request(provider_type, provider_api_format) + || !is_openai_responses_compact_request(provider_api_format) { return; } - if !body_rules_handle_path(body_rules, "parallel_tool_calls") { - body_object - .entry("parallel_tool_calls".to_string()) - .or_insert_with(|| json!(true)); + strip_codex_cache_control_fields(provider_request_body); + let Some(body_object) = provider_request_body.as_object_mut() else { + return; + }; + body_object + .retain(|field, _| CODEX_OPENAI_RESPONSES_COMPACT_BODY_FIELDS.contains(&field.as_str())); + body_object + .entry("parallel_tool_calls".to_string()) + .or_insert_with(|| json!(true)); + for field in [ + "tools", + "reasoning", + "service_tier", + "prompt_cache_key", + "text", + ] { + if body_object.get(field).is_some_and(Value::is_null) { + body_object.remove(field); + } } - if !body_rules_handle_path(body_rules, "include") { - match body_object.get_mut("include") { - Some(Value::Array(include)) => { - let has_reasoning_encrypted_content = include - .iter() - .any(|value| value.as_str() == Some(CODEX_REASONING_ENCRYPTED_CONTENT_INCLUDE)); - if !has_reasoning_encrypted_content { - include.push(json!(CODEX_REASONING_ENCRYPTED_CONTENT_INCLUDE)); + if body_object + .get("instructions") + .is_some_and(|value| value.is_null() || value.as_str().is_some_and(str::is_empty)) + { + body_object.remove("instructions"); + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct CodexOpenAiCompactRequestContractViolation { + pub field: &'static str, + pub reason: &'static str, +} + +pub fn validate_codex_openai_responses_compact_request_contract( + provider_request_body: &Value, + provider_type: &str, + provider_api_format: &str, +) -> Result<(), CodexOpenAiCompactRequestContractViolation> { + if codex_openai_endpoint_kind(provider_type, provider_api_format) + != Some(CodexOpenAiEndpointKind::Compact) + { + return Ok(()); + } + let Some(body) = provider_request_body.as_object() else { + return Err(CodexOpenAiCompactRequestContractViolation { + field: "$", + reason: "Codex Compact request body must be an object", + }); + }; + if body + .get("model") + .and_then(Value::as_str) + .is_none_or(|model| model.trim().is_empty()) + { + return Err(CodexOpenAiCompactRequestContractViolation { + field: "model", + reason: "Codex Compact model must be a non-empty string", + }); + } + let Some(input) = body.get("input").and_then(Value::as_array) else { + return Err(CodexOpenAiCompactRequestContractViolation { + field: "input", + reason: "Codex Compact input must be an array", + }); + }; + if input.iter().any(|item| { + !item.as_object().is_some_and(|item| { + item.get("type") + .and_then(Value::as_str) + .is_some_and(|item_type| !item_type.trim().is_empty()) + }) + }) { + return Err(CodexOpenAiCompactRequestContractViolation { + field: "input", + reason: "Codex Compact input items must be typed objects", + }); + } + for (field, valid) in [ + ( + "instructions", + body.get("instructions").is_none_or(Value::is_string), + ), + ("tools", body.get("tools").is_none_or(Value::is_array)), + ( + "reasoning", + body.get("reasoning").is_none_or(Value::is_object), + ), + ( + "service_tier", + body.get("service_tier").is_none_or(Value::is_string), + ), + ( + "prompt_cache_key", + body.get("prompt_cache_key").is_none_or(Value::is_string), + ), + ("text", body.get("text").is_none_or(Value::is_object)), + ] { + if !valid { + return Err(CodexOpenAiCompactRequestContractViolation { + field, + reason: "Codex Compact optional field has an invalid type", + }); + } + } + if !body + .get("parallel_tool_calls") + .is_some_and(Value::is_boolean) + { + return Err(CodexOpenAiCompactRequestContractViolation { + field: "parallel_tool_calls", + reason: "Codex Compact parallel_tool_calls must be a boolean", + }); + } + Ok(()) +} + +pub fn apply_codex_openai_compact_terminal_headers( + provider_request_headers: &mut BTreeMap, + provider_type: &str, + provider_api_format: &str, +) { + if codex_openai_endpoint_kind(provider_type, provider_api_format) + != Some(CodexOpenAiEndpointKind::Compact) + { + return; + } + remove_btree_header(provider_request_headers, "x-client-request-id"); + remove_btree_header(provider_request_headers, "accept"); + remove_btree_header(provider_request_headers, "content-encoding"); +} + +fn apply_codex_model_request_capabilities( + body_object: &mut serde_json::Map, + provider_api_format: &str, + capabilities: &CodexResponsesModelCapabilities, + body_rules: Option<&Value>, +) { + if is_openai_image_request(provider_api_format) { + return; + } + if !body_rules_handle_path(body_rules, "parallel_tool_calls") { + let requested = body_object + .get("parallel_tool_calls") + .and_then(Value::as_bool) + .unwrap_or(capabilities.supports_parallel_tool_calls); + body_object.insert( + "parallel_tool_calls".to_string(), + json!( + requested + && capabilities.supports_parallel_tool_calls + && !capabilities.use_responses_lite + ), + ); + } + + if !is_openai_responses_compact_request(provider_api_format) + && !body_rules_handle_path(body_rules, "include") + { + let include = if body_object.get("reasoning").is_some_and(Value::is_object) { + json!([CODEX_REASONING_ENCRYPTED_CONTENT_INCLUDE]) + } else { + json!([]) + }; + body_object.insert("include".to_string(), include); + } + + if !body_rules_handle_path(body_rules, "text") { + if capabilities.support_verbosity { + if let Some(default_verbosity) = capabilities.default_verbosity.as_deref() { + match body_object.get_mut("text") { + Some(Value::Object(text)) => { + text.entry("verbosity".to_string()) + .or_insert_with(|| json!(default_verbosity)); + } + None | Some(Value::Null) => { + body_object.insert( + "text".to_string(), + json!({ "verbosity": default_verbosity }), + ); + } + Some(_) => {} } } - Some(_) | None => { - body_object.insert( - "include".to_string(), - json!([CODEX_REASONING_ENCRYPTED_CONTENT_INCLUDE]), - ); + } else if let Some(Value::Object(text)) = body_object.get_mut("text") { + text.remove("verbosity"); + if text.is_empty() { + body_object.remove("text"); } } } + + if !body_rules_handle_path(body_rules, "service_tier") { + let service_tier = body_object + .get("service_tier") + .and_then(Value::as_str) + .map(str::to_string); + if !service_tier.as_deref().is_some_and(|service_tier| { + service_tier != "default" && capabilities.supports_service_tier(service_tier) + }) { + body_object.remove("service_tier"); + } + } } -fn ensure_codex_chat_reasoning_defaults( +fn ensure_codex_reasoning_defaults( body_object: &mut serde_json::Map, - provider_api_format: &str, + capabilities: &CodexResponsesModelCapabilities, + supports_reasoning_mode: bool, body_rules: Option<&Value>, ) { - if is_openai_responses_compact_request(provider_api_format) - || is_openai_image_request(provider_api_format) - { - return; - } - if body_rules_handle_path(body_rules, "reasoning") { + if body_rules_handle_path(body_rules, "reasoning") && !capabilities.use_responses_lite { + let has_summary = body_object + .get_mut("reasoning") + .and_then(Value::as_object_mut) + .and_then(|reasoning| { + if !capabilities.supports_reasoning_summary_parameter + || reasoning + .get("summary") + .is_some_and(codex_reasoning_summary_is_disabled) + { + reasoning.remove("summary"); + } + reasoning.get("summary") + }) + .is_some_and(|summary| !summary.is_null()); + if !has_summary { + remove_codex_reasoning_summary_delivery(body_object); + } return; } let reasoning = body_object .entry("reasoning".to_string()) .or_insert_with(|| json!({})); + if reasoning.is_null() { + *reasoning = json!({}); + } let Some(reasoning_object) = reasoning.as_object_mut() else { return; }; - reasoning_object - .entry("effort".to_string()) - .or_insert_with(|| json!(CODEX_DEFAULT_REASONING_EFFORT)); - reasoning_object - .entry("summary".to_string()) - .or_insert_with(|| json!(CODEX_DEFAULT_REASONING_SUMMARY)); + if reasoning_object.get("effort").is_none_or(Value::is_null) { + let default_effort = if supports_reasoning_mode + && reasoning_object + .get("mode") + .and_then(Value::as_str) + .is_some_and(|mode| matches!(mode, "standard" | "pro")) + { + Some(CODEX_DEFAULT_REASONING_EFFORT) + } else { + capabilities.default_reasoning_effort.as_deref() + }; + if let Some(default_effort) = default_effort { + reasoning_object.insert("effort".to_string(), json!(default_effort)); + } + } + if !capabilities.supports_reasoning_summary_parameter + || reasoning_object + .get("summary") + .is_some_and(codex_reasoning_summary_is_disabled) + { + reasoning_object.remove("summary"); + } else if !reasoning_object.contains_key("summary") { + if let Some(summary) = capabilities.default_reasoning_summary.as_deref() { + reasoning_object.insert("summary".to_string(), json!(summary)); + } + } + if capabilities.use_responses_lite { + reasoning_object.insert("context".to_string(), json!("all_turns")); + } + let has_summary = reasoning_object + .get("summary") + .is_some_and(|summary| !summary.is_null()); + if !has_summary { + remove_codex_reasoning_summary_delivery(body_object); + } +} + +fn codex_reasoning_summary_is_disabled(value: &Value) -> bool { + value.is_null() + || value + .as_str() + .is_some_and(|summary| summary.eq_ignore_ascii_case("none")) +} + +fn remove_codex_reasoning_summary_delivery(body_object: &mut serde_json::Map) { + let remove_stream_options = body_object + .get_mut("stream_options") + .and_then(Value::as_object_mut) + .is_some_and(|stream_options| { + stream_options.remove("reasoning_summary_delivery"); + stream_options.is_empty() + }); + if remove_stream_options { + body_object.remove("stream_options"); + } +} + +fn normalize_codex_reasoning_effort(body_object: &mut serde_json::Map) { + let Some(reasoning) = body_object + .get_mut("reasoning") + .and_then(Value::as_object_mut) + else { + return; + }; + let is_ultra = reasoning + .get("effort") + .and_then(Value::as_str) + .is_some_and(|effort| effort == "ultra"); + if is_ultra { + reasoning.insert("effort".to_string(), json!("max")); + } +} + +pub fn normalize_codex_openai_reasoning_wire_effort( + provider_request_body: &mut Value, + provider_type: &str, + provider_api_format: &str, +) { + if !provider_type.trim().eq_ignore_ascii_case("codex") + || !(aether_ai_formats::is_openai_responses_family_format(provider_api_format) + || aether_ai_formats::api_format_alias_matches(provider_api_format, "openai:search")) + { + return; + } + let Some(body_object) = provider_request_body.as_object_mut() else { + return; + }; + normalize_codex_reasoning_effort(body_object); +} + +fn is_codex_responses_lite_additional_tools_item(value: &Value) -> bool { + value + .get("type") + .and_then(Value::as_str) + .is_some_and(|item_type| item_type == "additional_tools") +} + +fn codex_tool_type_accepts_top_level_name(tool_type: &str) -> bool { + matches!(tool_type, "function" | "custom" | "namespace") +} + +fn is_codex_client_executed_tool(tool: &Value) -> bool { + match tool.get("type").and_then(Value::as_str) { + Some("function" | "custom" | "namespace") => true, + Some("tool_search") => tool.get("execution").and_then(Value::as_str) == Some("client"), + _ => false, + } +} + +fn retain_codex_client_executed_tools(additional_tools: &mut Value) { + let Some(tools) = additional_tools + .get_mut("tools") + .and_then(Value::as_array_mut) + else { + return; + }; + tools.retain(is_codex_client_executed_tool); +} + +fn is_codex_responses_lite_instruction_item(value: &Value) -> bool { + value + .get("type") + .and_then(Value::as_str) + .is_some_and(|item_type| item_type == "message") + && value + .get("role") + .and_then(Value::as_str) + .is_some_and(|role| role == "developer") + && value + .get("content") + .and_then(Value::as_array) + .is_some_and(|content| { + content.len() == 1 + && content[0] + .get("type") + .and_then(Value::as_str) + .is_some_and(|item_type| item_type == "input_text") + && content[0].get("text").is_some_and(Value::is_string) + }) +} + +fn strip_codex_responses_lite_image_details(item: &mut Value) { + let Some(item_object) = item.as_object_mut() else { + return; + }; + let item_type = item_object + .get("type") + .and_then(Value::as_str) + .unwrap_or_default(); + let content = match item_type { + "message" => item_object.get_mut("content"), + "function_call_output" | "custom_tool_call_output" => item_object.get_mut("output"), + _ => None, + }; + let Some(content) = content.and_then(Value::as_array_mut) else { + return; + }; + for content_item in content { + let Some(content_object) = content_item.as_object_mut() else { + continue; + }; + if content_object + .get("type") + .and_then(Value::as_str) + .is_some_and(|content_type| content_type == "input_image") + { + content_object.remove("detail"); + } + } +} + +fn apply_codex_responses_lite_body_contract( + body_object: &mut serde_json::Map, + capabilities: &CodexResponsesModelCapabilities, +) { + if !capabilities.use_responses_lite { + return; + } + + let tools_are_valid = body_object + .get("tools") + .is_none_or(|tools| tools.is_null() || tools.is_array()); + let instructions_are_valid = body_object + .get("instructions") + .is_none_or(|instructions| instructions.is_null() || instructions.is_string()); + if !tools_are_valid || !instructions_are_valid { + return; + } + if !body_object.get("input").is_some_and(Value::is_array) { + return; + } + + let top_level_tools = body_object + .remove("tools") + .and_then(|tools| tools.as_array().cloned()) + .map(|tools| { + tools + .into_iter() + .filter(is_codex_client_executed_tool) + .collect::>() + }); + let top_level_instructions = body_object + .remove("instructions") + .and_then(|instructions| instructions.as_str().map(ToOwned::to_owned)) + .filter(|instructions| !instructions.is_empty()); + let input = body_object + .get_mut("input") + .and_then(Value::as_array_mut) + .expect("Responses Lite input was validated as an array"); + + let existing_additional_tools = input + .iter() + .position(is_codex_responses_lite_additional_tools_item) + .map(|index| input.remove(index)); + let mut additional_tools = existing_additional_tools.unwrap_or_else(|| { + json!({ + "type": "additional_tools", + "role": "developer", + "tools": [], + }) + }); + if let Some(tools) = top_level_tools { + if let Some(object) = additional_tools.as_object_mut() { + object.insert("tools".to_string(), Value::Array(tools)); + } + } + retain_codex_client_executed_tools(&mut additional_tools); + input.insert(0, additional_tools); + + if let Some(instructions) = top_level_instructions { + if input + .get(1) + .is_some_and(is_codex_responses_lite_instruction_item) + { + input.remove(1); + } + input.insert( + 1, + json!({ + "type": "message", + "role": "developer", + "content": [{ + "type": "input_text", + "text": instructions, + }], + }), + ); + } + + for item in input + .iter_mut() + .filter(|item| !is_codex_responses_lite_additional_tools_item(item)) + { + strip_codex_responses_lite_image_details(item); + } + body_object.insert("parallel_tool_calls".to_string(), json!(false)); +} + +fn remove_empty_codex_instructions(body_object: &mut serde_json::Map) { + if body_object + .get("instructions") + .is_some_and(|value| value.is_null() || value.as_str().is_some_and(str::is_empty)) + { + body_object.remove("instructions"); + } } fn codex_tool_type_rejects_top_level_name(tool_type: &str) -> bool { let normalized = tool_type.trim().to_ascii_lowercase(); - !normalized.is_empty() - && normalized != "function" - && normalized != "custom" - && normalized != "namespace" + !normalized.is_empty() && !codex_tool_type_accepts_top_level_name(normalized.as_str()) } fn strip_codex_hosted_tool_names_for_backend(body_object: &mut serde_json::Map) { @@ -773,16 +1463,54 @@ pub fn apply_codex_openai_responses_special_body_edits( body_rules: Option<&Value>, user_api_key_id: Option<&str>, ) { - if !is_codex_openai_responses_request(provider_type, provider_api_format) { - return; - } - - let prompt_cache_key = codex_prompt_cache_key_to_insert( + let provider_model = provider_request_body + .get("model") + .and_then(Value::as_str) + .unwrap_or_default() + .to_string(); + apply_codex_openai_responses_special_body_edits_with_source_model( provider_request_body, provider_type, provider_api_format, + provider_model.as_str(), + provider_model.as_str(), + body_rules, user_api_key_id, ); +} + +pub fn apply_codex_openai_responses_special_body_edits_with_source_model( + provider_request_body: &mut Value, + provider_type: &str, + provider_api_format: &str, + provider_model: &str, + source_model: &str, + body_rules: Option<&Value>, + _user_api_key_id: Option<&str>, +) { + apply_codex_openai_responses_special_body_edits_with_source_model_and_capabilities( + provider_request_body, + provider_type, + provider_api_format, + provider_model, + source_model, + None, + body_rules, + ); +} + +pub fn apply_codex_openai_responses_special_body_edits_with_source_model_and_capabilities( + provider_request_body: &mut Value, + provider_type: &str, + provider_api_format: &str, + provider_model: &str, + source_model: &str, + model_capabilities: Option<&CodexResponsesModelCapabilities>, + body_rules: Option<&Value>, +) { + if !is_codex_openai_responses_request(provider_type, provider_api_format) { + return; + } let Some(body_object) = provider_request_body.as_object_mut() else { return; @@ -799,23 +1527,16 @@ pub fn apply_codex_openai_responses_special_body_edits( } else if !body_rules_handle_path(body_rules, "store") { body_object.insert("store".to_string(), json!(false)); } - ensure_codex_responses_passthrough_fields(body_object, provider_api_format, body_rules); - if !body_rules_handle_path(body_rules, "instructions") - && !body_object.contains_key("instructions") - { - body_object.insert( - "instructions".to_string(), - json!(CODEX_DEFAULT_INSTRUCTIONS), - ); - } else if body_object.contains_key("instructions") - && body_object.get("instructions").is_some_and(|v| v.is_null()) - { - body_object.insert("instructions".to_string(), json!("")); - } + remove_empty_codex_instructions(body_object); strip_codex_hosted_tool_names_for_backend(body_object); strip_codex_hosted_tool_choice_name_for_backend(body_object); - if is_openai_image_request(provider_api_format) - || codex_openai_responses_tool_choice_references_image_generation(body_object) + if !is_openai_responses_compact_request(provider_api_format) { + body_object + .entry("tool_choice".to_string()) + .or_insert_with(|| json!("auto")); + } + if !is_openai_responses_compact_request(provider_api_format) + && codex_openai_responses_tool_choice_references_image_generation(body_object) { body_object.insert( "model".to_string(), @@ -825,9 +1546,56 @@ pub fn apply_codex_openai_responses_special_body_edits( apply_codex_openai_image_tool_overrides(body_object); inject_codex_default_variation_prompt(body_object); } - + let effective_provider_model = body_object + .get("model") + .and_then(Value::as_str) + .map(str::trim) + .filter(|model| !model.is_empty()) + .unwrap_or(provider_model) + .to_string(); + let bundled_capabilities; + let capabilities = if effective_provider_model == CODEX_OPENAI_IMAGE_INTERNAL_MODEL { + bundled_capabilities = resolve_codex_responses_model_capabilities( + effective_provider_model.as_str(), + source_model, + None, + ); + &bundled_capabilities + } else if let Some(capabilities) = model_capabilities { + capabilities + } else { + bundled_capabilities = resolve_codex_responses_model_capabilities( + effective_provider_model.as_str(), + source_model, + None, + ); + &bundled_capabilities + }; + let supports_reasoning_mode = + crate::formats::shared::model_directives::openai_model_resolves_to_gpt_5_6( + effective_provider_model.as_str(), + source_model, + ); + ensure_codex_reasoning_defaults( + body_object, + capabilities, + supports_reasoning_mode, + body_rules, + ); + normalize_codex_reasoning_effort(body_object); + apply_codex_model_request_capabilities( + body_object, + provider_api_format, + capabilities, + body_rules, + ); + apply_codex_responses_lite_body_contract(body_object, capabilities); strip_codex_cache_control_fields(provider_request_body); - insert_codex_prompt_cache_key(provider_request_body, prompt_cache_key); + apply_codex_openai_responses_compact_body_edits( + provider_request_body, + provider_type, + provider_api_format, + ); } pub fn apply_codex_openai_responses_chat_body_edits( @@ -837,13 +1605,60 @@ pub fn apply_codex_openai_responses_chat_body_edits( body_rules: Option<&Value>, user_api_key_id: Option<&str>, ) { - apply_codex_openai_responses_special_body_edits( + let provider_model = provider_request_body + .get("model") + .and_then(Value::as_str) + .unwrap_or_default() + .to_string(); + apply_codex_openai_responses_chat_body_edits_with_source_model( provider_request_body, provider_type, provider_api_format, + provider_model.as_str(), + provider_model.as_str(), body_rules, user_api_key_id, ); +} + +pub fn apply_codex_openai_responses_chat_body_edits_with_source_model( + provider_request_body: &mut Value, + provider_type: &str, + provider_api_format: &str, + provider_model: &str, + source_model: &str, + body_rules: Option<&Value>, + _user_api_key_id: Option<&str>, +) { + apply_codex_openai_responses_chat_body_edits_with_source_model_and_capabilities( + provider_request_body, + provider_type, + provider_api_format, + provider_model, + source_model, + None, + body_rules, + ); +} + +pub fn apply_codex_openai_responses_chat_body_edits_with_source_model_and_capabilities( + provider_request_body: &mut Value, + provider_type: &str, + provider_api_format: &str, + provider_model: &str, + source_model: &str, + model_capabilities: Option<&CodexResponsesModelCapabilities>, + body_rules: Option<&Value>, +) { + apply_codex_openai_responses_special_body_edits_with_source_model_and_capabilities( + provider_request_body, + provider_type, + provider_api_format, + provider_model, + source_model, + model_capabilities, + body_rules, + ); if !is_codex_openai_responses_request(provider_type, provider_api_format) { return; @@ -851,93 +1666,742 @@ pub fn apply_codex_openai_responses_chat_body_edits( let Some(body_object) = provider_request_body.as_object_mut() else { return; }; - ensure_codex_chat_reasoning_defaults(body_object, provider_api_format, body_rules); if let Some(prompt_cache_key) = body_object.remove("prompt_cache_key") { body_object.insert("prompt_cache_key".to_string(), prompt_cache_key); } } -pub fn apply_codex_openai_responses_special_headers( +pub fn apply_codex_openai_responses_lite_header( provider_request_headers: &mut BTreeMap, - provider_request_body: &Value, - original_headers: &http::HeaderMap, provider_type: &str, provider_api_format: &str, - request_id: Option<&str>, - decrypted_auth_config_raw: Option<&str>, + provider_model: &str, + source_model: &str, +) { + apply_codex_openai_responses_lite_header_with_capabilities( + provider_request_headers, + provider_type, + provider_api_format, + provider_model, + source_model, + None, + ); +} + +pub fn apply_codex_openai_responses_lite_header_with_capabilities( + provider_request_headers: &mut BTreeMap, + provider_type: &str, + provider_api_format: &str, + provider_model: &str, + source_model: &str, + model_capabilities: Option<&CodexResponsesModelCapabilities>, ) { if !is_codex_openai_responses_request(provider_type, provider_api_format) { return; } + remove_btree_header(provider_request_headers, CODEX_RESPONSES_LITE_HEADER); + let bundled_capabilities; + let capabilities = if let Some(capabilities) = model_capabilities { + capabilities + } else { + bundled_capabilities = + resolve_codex_responses_model_capabilities(provider_model, source_model, None); + &bundled_capabilities + }; + if capabilities.use_responses_lite { + provider_request_headers + .insert(CODEX_RESPONSES_LITE_HEADER.to_string(), "true".to_string()); + } +} - let prompt_cache_key = provider_request_body - .get("prompt_cache_key") +pub fn apply_codex_openai_special_headers( + provider_request_headers: &mut BTreeMap, + provider_request_body: &Value, + _original_headers: &http::HeaderMap, + provider_type: &str, + provider_api_format: &str, + _request_id: Option<&str>, + decrypted_auth_config_raw: Option<&str>, +) { + let Some(endpoint_kind) = codex_openai_endpoint_kind(provider_type, provider_api_format) else { + return; + }; + + let auth_identity = parse_codex_auth_identity(decrypted_auth_config_raw); + + remove_btree_header(provider_request_headers, "chatgpt-account-id"); + remove_btree_header(provider_request_headers, "x-openai-fedramp"); + if let Some(account_id) = auth_identity.account_id { + provider_request_headers.insert("chatgpt-account-id".to_string(), account_id); + } + + if auth_identity.is_fedramp { + provider_request_headers.insert("x-openai-fedramp".to_string(), "true".to_string()); + } + + set_codex_client_header( + provider_request_headers, + "user-agent", + CODEX_CLIENT_USER_AGENT, + ); + set_codex_client_header( + provider_request_headers, + "originator", + CODEX_CLIENT_ORIGINATOR, + ); + if endpoint_kind == CodexOpenAiEndpointKind::Search { + remove_btree_header(provider_request_headers, CODEX_RESPONSES_LITE_HEADER); + remove_btree_header(provider_request_headers, "openai-beta"); + if provider_request_headers.iter().any(|(name, value)| { + name.eq_ignore_ascii_case("accept") + && header_value_contains_media_type(value, "text/event-stream") + }) { + remove_btree_header(provider_request_headers, "accept"); + } + return; + } + if endpoint_kind == CodexOpenAiEndpointKind::Images { + return; + } + + let provider_model = provider_request_body + .get("model") .and_then(Value::as_str) - .map(str::trim) - .filter(|value| !value.is_empty()); + .unwrap_or_default(); + apply_codex_openai_responses_lite_header( + provider_request_headers, + provider_type, + provider_api_format, + provider_model, + provider_model, + ); - if !header_map_has_non_empty_value(original_headers, "chatgpt-account-id") - && !btree_map_has_non_empty_value(provider_request_headers, "chatgpt-account-id") - { - if let Some(account_id) = extract_codex_account_id(decrypted_auth_config_raw) { - provider_request_headers.insert("chatgpt-account-id".to_string(), account_id); - } - } - - if !header_map_has_non_empty_value(original_headers, "x-client-request-id") - && !btree_map_has_non_empty_value(provider_request_headers, "x-client-request-id") - { - if let Some(request_id) = request_id.map(str::trim).filter(|value| !value.is_empty()) { - provider_request_headers - .insert("x-client-request-id".to_string(), request_id.to_string()); - } - } - - if !is_openai_image_request(provider_api_format) { - maybe_insert_default_codex_header( - provider_request_headers, - original_headers, - "user-agent", - CODEX_DEFAULT_USER_AGENT, - ); - maybe_insert_default_codex_header( - provider_request_headers, - original_headers, - "originator", - CODEX_DEFAULT_ORIGINATOR, - ); - } - - let short_session_id = prompt_cache_key.and_then(build_short_codex_header_id); - - if !header_map_has_non_empty_value(original_headers, "session_id") - && !btree_map_has_non_empty_value(provider_request_headers, "session_id") - { - if let Some(short_session_id) = short_session_id.as_deref() { - provider_request_headers.insert("session_id".to_string(), short_session_id.to_string()); - } - } - - if aether_ai_formats::is_openai_responses_format(provider_api_format) - && !header_map_has_non_empty_value(original_headers, "conversation_id") - && !btree_map_has_non_empty_value(provider_request_headers, "conversation_id") - { - if let Some(short_session_id) = short_session_id.as_deref() { - provider_request_headers - .insert("conversation_id".to_string(), short_session_id.to_string()); - } - } + apply_codex_openai_compact_terminal_headers( + provider_request_headers, + provider_type, + provider_api_format, + ); } #[cfg(test)] mod tests { use super::{ apply_codex_openai_responses_chat_body_edits, + apply_codex_openai_responses_compact_body_edits, + apply_codex_openai_responses_lite_header_with_capabilities, apply_codex_openai_responses_special_body_edits, - apply_openai_responses_compact_special_body_edits, CODEX_OPENAI_IMAGE_INTERNAL_MODEL, - CODEX_OPENAI_RESPONSES_UNSUPPORTED_BODY_FIELDS, + apply_codex_openai_responses_special_body_edits_with_source_model_and_capabilities, + apply_codex_openai_special_headers, apply_openai_responses_compact_special_body_edits, + build_codex_model_catalog_metadata, bundled_codex_model_cards, effective_codex_model_cards, + resolve_codex_responses_model_capabilities, + validate_codex_openai_responses_compact_request_contract, CODEX_CLIENT_ORIGINATOR, + CODEX_CLIENT_USER_AGENT, CODEX_OPENAI_IMAGE_INTERNAL_MODEL, + CODEX_OPENAI_RESPONSES_UNSUPPORTED_BODY_FIELDS, CODEX_RESPONSES_LITE_HEADER, }; - use serde_json::json; + use serde_json::{json, Value}; + + #[test] + fn model_card_drives_unknown_codex_model_body_and_header_contracts() { + let card = json!({ + "id": "gpt-future-agent", + "slug": "gpt-future-agent", + "use_responses_lite": true, + "default_reasoning_level": "low", + "default_reasoning_summary": "auto", + "supports_parallel_tool_calls": true, + "support_verbosity": true, + "default_verbosity": "low", + "service_tiers": [{"id": "priority"}], + "supported_reasoning_levels": [ + {"effort": "low"}, + {"effort": "max"}, + {"effort": "ultra"} + ], + "base_instructions": "large prompt", + "model_messages": {"instructions_template": "large template"}, + "available_in_plans": ["pro"], + "future_capability": {"mode": "native"} + }); + let metadata = build_codex_model_catalog_metadata(&[card]); + assert_eq!( + metadata["codex_models"]["cards"]["gpt-future-agent"]["future_capability"]["mode"], + "native" + ); + let stored_card = &metadata["codex_models"]["cards"]["gpt-future-agent"]; + assert!(stored_card.get("base_instructions").is_none()); + assert!(stored_card.get("model_messages").is_none()); + assert!(stored_card.get("available_in_plans").is_none()); + let capabilities = resolve_codex_responses_model_capabilities( + "gpt-future-agent", + "gpt-future-agent", + Some(&metadata), + ); + assert!(capabilities.supports_reasoning_summary_parameter); + assert_eq!( + capabilities.default_reasoning_effort.as_deref(), + Some("low") + ); + assert_eq!( + capabilities.default_reasoning_summary.as_deref(), + Some("auto") + ); + + let mut body = json!({ + "model": "gpt-future-agent", + "instructions": "Use the tools.", + "input": [{"id": "msg-1", "type": "message", "role": "user", "content": []}], + "tools": [{"type": "function", "name": "lookup"}], + "reasoning": {"effort": "ultra"}, + "parallel_tool_calls": true, + "service_tier": "priority", + "text": {"format": {"type": "json_schema"}}, + "include": ["file_search_call.results"] + }); + apply_codex_openai_responses_special_body_edits_with_source_model_and_capabilities( + &mut body, + "codex", + "openai:responses", + "gpt-future-agent", + "gpt-future-agent", + Some(&capabilities), + None, + ); + assert_eq!(body["reasoning"]["effort"], "max"); + assert_eq!(body["reasoning"]["summary"], "auto"); + assert_eq!(body["reasoning"]["context"], "all_turns"); + assert!(body.get("instructions").is_none()); + assert!(body.get("tools").is_none()); + assert_eq!(body["input"][0]["type"], "additional_tools"); + assert!(body["input"][1].get("id").is_none()); + assert_eq!(body["input"][2]["id"], "msg-1"); + assert_eq!(body["parallel_tool_calls"], false); + assert_eq!(body["service_tier"], "priority"); + assert_eq!(body["text"]["verbosity"], "low"); + assert_eq!(body["text"]["format"]["type"], "json_schema"); + assert_eq!(body["include"], json!(["reasoning.encrypted_content"])); + + let mut headers = std::collections::BTreeMap::new(); + apply_codex_openai_responses_lite_header_with_capabilities( + &mut headers, + "codex", + "openai:responses", + "gpt-future-agent", + "gpt-future-agent", + Some(&capabilities), + ); + assert_eq!( + headers.get("x-openai-internal-codex-responses-lite"), + Some(&"true".to_string()) + ); + } + + #[test] + fn model_card_without_summary_parameter_support_keeps_reasoning_and_ids() { + let metadata = build_codex_model_catalog_metadata(&[json!({ + "slug": "gpt-no-summary", + "supports_reasoning_summary_parameter": false, + "default_reasoning_level": "high", + "default_reasoning_summary": "detailed", + "supported_reasoning_levels": [{"effort": "high"}], + "supports_parallel_tool_calls": true + })]); + let capabilities = resolve_codex_responses_model_capabilities( + "gpt-no-summary", + "gpt-no-summary", + Some(&metadata), + ); + assert!(!capabilities.supports_reasoning_summary_parameter); + + let mut body = json!({ + "model": "gpt-no-summary", + "input": [{ + "id": "msg-client", + "type": "message", + "role": "user", + "content": [] + }], + "reasoning": {"effort": "high", "summary": "detailed"}, + "stream_options": { + "reasoning_summary_delivery": "sequential_cutoff", + "future_option": true + } + }); + apply_codex_openai_responses_special_body_edits_with_source_model_and_capabilities( + &mut body, + "codex", + "openai:responses", + "gpt-no-summary", + "gpt-no-summary", + Some(&capabilities), + None, + ); + + assert_eq!(body["reasoning"], json!({"effort": "high"})); + assert_eq!(body["include"], json!(["reasoning.encrypted_content"])); + assert_eq!(body["stream_options"], json!({"future_option": true})); + assert_eq!(body["input"][0]["id"], "msg-client"); + } + + #[test] + fn codex_request_normalizes_the_reasoning_envelope_without_model_defaults() { + let capabilities = resolve_codex_responses_model_capabilities( + "gpt-future-agent", + "gpt-future-agent", + None, + ); + for initial_reasoning in [None, Some(Value::Null)] { + let mut body = json!({ + "model": "gpt-future-agent", + "input": [], + "stream_options": { + "reasoning_summary_delivery": "sequential_cutoff" + } + }); + if let Some(initial_reasoning) = initial_reasoning { + body["reasoning"] = initial_reasoning; + } + + apply_codex_openai_responses_special_body_edits_with_source_model_and_capabilities( + &mut body, + "codex", + "openai:responses", + "gpt-future-agent", + "gpt-future-agent", + Some(&capabilities), + None, + ); + + assert_eq!(body["reasoning"], json!({})); + assert_eq!(body["include"], json!(["reasoning.encrypted_content"])); + assert!(body.get("stream_options").is_none()); + } + + let mut compact = json!({ + "model": "gpt-future-agent", + "input": [], + "reasoning": null, + "include": ["reasoning.encrypted_content"], + "store": true + }); + apply_codex_openai_responses_special_body_edits_with_source_model_and_capabilities( + &mut compact, + "codex", + "openai:responses:compact", + "gpt-future-agent", + "gpt-future-agent", + Some(&capabilities), + None, + ); + + assert_eq!(compact["reasoning"], json!({})); + assert!(compact.get("include").is_none()); + assert!(compact.get("store").is_none()); + } + + #[test] + fn empty_codex_model_catalog_clears_the_card_map() { + assert_eq!( + build_codex_model_catalog_metadata(&[])["codex_models"]["cards"], + json!({}) + ); + } + + #[test] + fn ultra_reasoning_always_uses_max_on_the_provider_wire() { + for provider_api_format in ["openai:responses", "openai:responses:compact"] { + let mut body = json!({ + "model": "gpt-5.6-luna", + "input": [], + "reasoning": {"effort": "ultra"} + }); + apply_codex_openai_responses_special_body_edits_with_source_model_and_capabilities( + &mut body, + "codex", + provider_api_format, + "gpt-5.6-luna", + "gpt-5.6-luna", + None, + None, + ); + assert_eq!(body["reasoning"]["effort"], "max"); + } + + let mut custom = json!({ + "model": "gpt-5.6-luna", + "input": [], + "reasoning": {"effort": " ultra "} + }); + apply_codex_openai_responses_special_body_edits_with_source_model_and_capabilities( + &mut custom, + "codex", + "openai:responses", + "gpt-5.6-luna", + "gpt-5.6-luna", + None, + None, + ); + assert_eq!(custom["reasoning"]["effort"], " ultra "); + } + + #[test] + fn populated_remote_catalog_is_authoritative_and_matches_model_variants() { + let metadata = build_codex_model_catalog_metadata(&[json!({ + "slug": "gpt-future-agent", + "supports_reasoning_summary_parameter": true, + "default_reasoning_level": "high", + "supported_reasoning_levels": [{"effort": "high"}], + "supports_parallel_tool_calls": true + })]); + + let variant = resolve_codex_responses_model_capabilities( + "gpt-future-agent-2026-07-10", + "gpt-future-agent-2026-07-10", + Some(&metadata), + ); + assert_eq!(variant.default_reasoning_effort.as_deref(), Some("high")); + assert!(variant.supports_parallel_tool_calls); + + for model in [ + "gpt-future-agentpreview", + "provider_1/gpt-future-agent-2026-07-10", + ] { + let capabilities = + resolve_codex_responses_model_capabilities(model, model, Some(&metadata)); + assert_eq!( + capabilities.default_reasoning_effort.as_deref(), + Some("high"), + "model: {model}" + ); + } + + for model in [ + "GPT-FUTURE-AGENT", + "gpt_future_agent", + "org/team/gpt-future-agent", + "provider!/gpt-future-agent", + ] { + let capabilities = + resolve_codex_responses_model_capabilities(model, model, Some(&metadata)); + assert_eq!( + capabilities.default_reasoning_effort, None, + "model: {model}" + ); + } + + let missing = resolve_codex_responses_model_capabilities( + "gpt-5.6-sol", + "gpt-5.6-sol", + Some(&metadata), + ); + assert!(!missing.use_responses_lite); + assert!(missing.supports_reasoning_summary_parameter); + assert!(!missing.supports_parallel_tool_calls); + assert_eq!(missing.default_reasoning_effort, None); + + let mut body = json!({ + "model": "gpt-5.6-sol", + "input": [], + "reasoning": {"effort": "high"}, + "parallel_tool_calls": true, + "service_tier": "priority", + "text": { + "verbosity": "high", + "format": {"type": "json_schema"} + } + }); + apply_codex_openai_responses_special_body_edits_with_source_model_and_capabilities( + &mut body, + "codex", + "openai:responses", + "gpt-5.6-sol", + "gpt-5.6-sol", + Some(&missing), + None, + ); + assert_eq!(body["reasoning"]["effort"], "high"); + assert_eq!(body["include"], json!(["reasoning.encrypted_content"])); + assert_eq!(body["parallel_tool_calls"], false); + assert!(body.get("service_tier").is_none()); + assert!(body["text"].get("verbosity").is_none()); + assert_eq!(body["text"]["format"]["type"], "json_schema"); + + let mut compact = json!({ + "model": "gpt-5.6-sol", + "input": [], + "reasoning": {"effort": "high"} + }); + apply_codex_openai_responses_special_body_edits_with_source_model_and_capabilities( + &mut compact, + "codex", + "openai:responses:compact", + "gpt-5.6-sol", + "gpt-5.6-sol", + Some(&missing), + None, + ); + assert_eq!(compact["reasoning"]["effort"], "high"); + } + + #[test] + fn model_catalog_keys_and_hidden_card_overlays_use_exact_slugs() { + let metadata = build_codex_model_catalog_metadata(&[ + json!({"slug": "gpt-card"}), + json!({"slug": "GPT-CARD"}), + json!({"slug": "gpt_card"}), + ]); + let cards = metadata["codex_models"]["cards"] + .as_object() + .expect("card map"); + assert_eq!(cards.len(), 3); + assert!(cards.contains_key("gpt-card")); + assert!(cards.contains_key("GPT-CARD")); + assert!(cards.contains_key("gpt_card")); + + let effective = effective_codex_model_cards(&[json!({ + "slug": "GPT-5.6-SOL", + "visibility": "hide" + })]); + assert!(effective.iter().any(|card| card["slug"] == "gpt-5.6-sol")); + assert!(effective.iter().any(|card| card["slug"] == "GPT-5.6-SOL")); + + let effective = effective_codex_model_cards(&[json!({ + "id": "gpt-5.6-sol", + "default_reasoning_level": "high", + "future_capability": {"mode": "native"} + })]); + let sol = effective + .iter() + .find(|card| card["slug"] == "gpt-5.6-sol") + .expect("merged Sol card"); + assert_eq!(sol["default_reasoning_level"], "high"); + assert_eq!(sol["future_capability"]["mode"], "native"); + assert_eq!(sol["use_responses_lite"], true); + assert_eq!(sol["default_reasoning_summary"], "none"); + assert!(sol["supported_reasoning_levels"] + .as_array() + .is_some_and(|levels| levels.iter().any(|level| level["effort"] == "ultra"))); + } + + #[test] + fn bare_gpt_5_6_requires_an_explicit_model_card() { + let conservative = resolve_codex_responses_model_capabilities("gpt-5.6", "gpt-5.6", None); + assert!(!conservative.use_responses_lite); + assert!(conservative.supports_reasoning_summary_parameter); + assert!(!conservative.supports_parallel_tool_calls); + assert_eq!(conservative.default_reasoning_effort, None); + assert_eq!(conservative.default_verbosity, None); + assert!(conservative.supported_service_tiers.is_empty()); + + let metadata = build_codex_model_catalog_metadata(&[json!({ + "slug": "gpt-5.6", + "use_responses_lite": true, + "supports_reasoning_summary_parameter": true, + "default_reasoning_level": "medium", + "supported_reasoning_levels": [{"effort": "medium"}], + "supports_parallel_tool_calls": true, + "support_verbosity": true, + "default_verbosity": "low", + "service_tiers": [{"id": "priority"}] + })]); + let explicit = + resolve_codex_responses_model_capabilities("gpt-5.6", "gpt-5.6", Some(&metadata)); + assert!(explicit.use_responses_lite); + assert_eq!(explicit.default_reasoning_effort.as_deref(), Some("medium")); + assert_eq!(explicit.default_verbosity.as_deref(), Some("low")); + assert_eq!(explicit.supported_service_tiers, vec!["priority"]); + + let opaque = resolve_codex_responses_model_capabilities( + "deployment-production", + "gpt-5.6-sol", + None, + ); + assert!(opaque.use_responses_lite); + assert_eq!(opaque.default_reasoning_effort.as_deref(), Some("low")); + + for provider_model in ["GPT-5.6-SOL", "gpt_5.6_sol", "org/team/gpt-5.6-sol"] { + let capabilities = + resolve_codex_responses_model_capabilities(provider_model, "gpt-5.6-sol", None); + assert!(!capabilities.use_responses_lite, "model: {provider_model}"); + } + } + + #[test] + fn model_card_preserves_custom_reasoning_effort_values() { + let metadata = build_codex_model_catalog_metadata(&[json!({ + "slug": "codex-custom", + "supports_reasoning_summary_parameter": true, + "default_reasoning_level": "VendorEffortX", + "supported_reasoning_levels": [ + {"effort": "VendorEffortX"}, + {"effort": "MAX"} + ] + })]); + let capabilities = resolve_codex_responses_model_capabilities( + "codex-custom", + "codex-custom", + Some(&metadata), + ); + assert_eq!( + capabilities.default_reasoning_effort.as_deref(), + Some("VendorEffortX") + ); + assert_eq!( + capabilities.supported_reasoning_efforts, + vec!["VendorEffortX", "MAX"] + ); + assert!(capabilities.supports_reasoning_effort("VendorEffortX")); + assert!(!capabilities.supports_reasoning_effort("vendoreffortx")); + assert!(!capabilities.supports_reasoning_effort("max")); + } + + #[test] + fn bundled_auto_review_card_matches_the_codex_request_profile() { + let card = bundled_codex_model_cards() + .iter() + .find(|card| card["slug"] == "codex-auto-review") + .expect("Codex auto review card"); + assert_eq!(card["visibility"], "hide"); + assert_eq!(card["supported_in_api"], true); + assert_eq!(card["priority"], 43); + + let capabilities = resolve_codex_responses_model_capabilities( + "codex-auto-review", + "codex-auto-review", + None, + ); + assert!(!capabilities.use_responses_lite); + assert_eq!( + capabilities.default_reasoning_effort.as_deref(), + Some("medium") + ); + assert_eq!(capabilities.default_reasoning_summary, None); + assert!(capabilities.supports_parallel_tool_calls); + assert_eq!(capabilities.default_verbosity.as_deref(), Some("low")); + assert!(capabilities.supported_service_tiers.is_empty()); + } + + #[test] + fn codex_identity_headers_are_derived_only_from_auth_config() { + let mut headers = std::collections::BTreeMap::from([ + ("chatgpt-account-id".to_string(), "spoofed".to_string()), + ("x-openai-fedramp".to_string(), "true".to_string()), + ]); + + apply_codex_openai_special_headers( + &mut headers, + &json!({"model": "gpt-5.6-sol", "input": []}), + &http::HeaderMap::new(), + "codex", + "openai:responses", + Some("request-1"), + Some(r#"{"is_fedramp":false}"#), + ); + + assert!(!headers.contains_key("chatgpt-account-id")); + assert!(!headers.contains_key("x-openai-fedramp")); + assert_eq!( + headers + .get("x-openai-internal-codex-responses-lite") + .map(String::as_str), + Some("true") + ); + } + + #[test] + fn codex_search_uses_identity_headers_without_responses_protocol_headers() { + let mut headers = std::collections::BTreeMap::from([ + ( + "x-openai-internal-codex-responses-lite".to_string(), + "true".to_string(), + ), + ("openai-beta".to_string(), "responses=v1".to_string()), + ( + "accept".to_string(), + "application/json, Text/Event-Stream; q=0.9".to_string(), + ), + ]); + + apply_codex_openai_special_headers( + &mut headers, + &json!({"id": "session-1", "model": "gpt-5.6-luna"}), + &http::HeaderMap::new(), + "codex", + "openai:search", + Some("request-search"), + Some(r#"{"account_id":"account-1","is_fedramp":true}"#), + ); + + assert_eq!( + headers.get("chatgpt-account-id").map(String::as_str), + Some("account-1") + ); + assert_eq!( + headers.get("x-openai-fedramp").map(String::as_str), + Some("true") + ); + assert_eq!( + headers.get("user-agent").map(String::as_str), + Some(CODEX_CLIENT_USER_AGENT) + ); + assert_eq!( + headers.get("originator").map(String::as_str), + Some(CODEX_CLIENT_ORIGINATOR) + ); + assert!(!headers.contains_key(CODEX_RESPONSES_LITE_HEADER)); + assert!(!headers.contains_key("openai-beta")); + assert!(!headers.contains_key("accept")); + } + + #[test] + fn standard_codex_models_do_not_send_the_responses_lite_header() { + let mut headers = std::collections::BTreeMap::from([( + "x-openai-internal-codex-responses-lite".to_string(), + "true".to_string(), + )]); + + apply_codex_openai_special_headers( + &mut headers, + &json!({"model": "gpt-5.4", "input": []}), + &http::HeaderMap::new(), + "codex", + "openai:responses", + Some("request-1"), + None, + ); + + assert!(!headers.contains_key("x-openai-internal-codex-responses-lite")); + } + + #[test] + fn codex_compact_headers_preserve_session_identity_and_remove_stream_headers() { + let mut headers = std::collections::BTreeMap::from([ + ("session-id".to_string(), "session-1".to_string()), + ("thread-id".to_string(), "thread-1".to_string()), + ("x-client-request-id".to_string(), "thread-1".to_string()), + ("accept".to_string(), "text/event-stream".to_string()), + ]); + + apply_codex_openai_special_headers( + &mut headers, + &json!({"model": "gpt-5.6", "input": []}), + &http::HeaderMap::new(), + "codex", + "openai:responses:compact", + Some("request-1"), + None, + ); + + assert_eq!( + headers.get("session-id").map(String::as_str), + Some("session-1") + ); + assert_eq!( + headers.get("thread-id").map(String::as_str), + Some("thread-1") + ); + assert!(!headers.contains_key("x-client-request-id")); + assert!(!headers.contains_key("accept")); + } #[test] fn codex_responses_body_edits_inject_passthrough_fields_without_reasoning_summary() { @@ -958,17 +2422,47 @@ mod tests { None, ); - assert!(provider_request_body.get("reasoning").is_none()); + assert_eq!(provider_request_body["reasoning"]["effort"], "medium"); + assert!(provider_request_body["reasoning"].get("summary").is_none()); assert_eq!( provider_request_body["include"], json!(["reasoning.encrypted_content"]) ); assert_eq!(provider_request_body["parallel_tool_calls"], json!(true)); - assert_eq!(provider_request_body["instructions"], json!("")); + assert_eq!(provider_request_body["tool_choice"], json!("auto")); + assert!(provider_request_body.get("instructions").is_none()); } #[test] - fn codex_responses_body_edits_preserve_existing_include_and_parallel_tool_calls() { + fn codex_responses_body_edits_omit_disabled_reasoning_summary_and_delivery() { + let mut provider_request_body = json!({ + "input": [{"role": "user", "content": "hello"}], + "model": "gpt-5.6-sol", + "stream": true, + "reasoning": {"effort": "high", "summary": "none"}, + "stream_options": { + "reasoning_summary_delivery": "sequential_cutoff", + "future_option": true + } + }); + + apply_codex_openai_responses_special_body_edits( + &mut provider_request_body, + "codex", + "openai:responses", + None, + None, + ); + + assert!(provider_request_body["reasoning"].get("summary").is_none()); + assert_eq!( + provider_request_body["stream_options"], + json!({"future_option": true}) + ); + } + + #[test] + fn codex_responses_body_edits_project_include_and_preserve_disabled_parallel_calls() { let mut provider_request_body = json!( { "input": [], "model": "gpt-5.4", @@ -1000,16 +2494,7 @@ mod tests { ); assert_eq!( provider_request_body["include"], - json!([ - "file_search_call.results", - "web_search_call.results", - "web_search_call.action.sources", - "message.input_image.image_url", - "computer_call_output.output.image_url", - "code_interpreter_call.outputs", - "message.output_text.logprobs", - "reasoning.encrypted_content" - ]) + json!(["reasoning.encrypted_content"]) ); assert_eq!(provider_request_body["parallel_tool_calls"], json!(false)); } @@ -1093,9 +2578,14 @@ mod tests { } #[test] - fn codex_responses_body_edits_strip_sub2api_unsupported_fields() { + fn codex_responses_body_edits_apply_backend_request_contract() { let mut provider_request_body = json!({ - "input": [{"role": "user", "content": "hello"}], + "input": [{ + "id": "msg-1", + "type": "message", + "role": "user", + "content": [{"id": "content-1", "type": "input_text", "text": "hello"}] + }], "model": "gpt-5.4", "max_output_tokens": 1024, "max_completion_tokens": 1024, @@ -1105,9 +2595,10 @@ mod tests { "presence_penalty": 0.1, "user": "user-123", "metadata": {"client": "cursor"}, + "prompt_cache_options": {"mode": "explicit", "ttl": "30m"}, "prompt_cache_retention": "24h", "safety_identifier": "safe-user-123", - "stream_options": {"include_usage": true}, + "stream_options": {"reasoning_summary_delivery": "sequential_cutoff"}, "previous_response_id": "resp_123" }); @@ -1125,7 +2616,105 @@ mod tests { "{field} must be stripped" ); } - assert_eq!(provider_request_body["input"][0]["content"], json!("hello")); + assert!(provider_request_body.get("stream_options").is_none()); + assert_eq!(provider_request_body["input"][0]["id"], "msg-1"); + assert_eq!( + provider_request_body["input"][0]["content"][0]["id"], + json!("content-1") + ); + } + + #[test] + fn codex_responses_body_edits_preserve_client_selected_input_item_ids() { + let body_rules = json!([{"action":"set","path":"store","value":true}]); + let mut provider_request_body = json!({ + "model": "gpt-5.4", + "store": true, + "input": [{ + "id": "msg-1", + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "hello"}] + }] + }); + + apply_codex_openai_responses_special_body_edits( + &mut provider_request_body, + "codex", + "openai:responses", + Some(&body_rules), + None, + ); + + assert_eq!(provider_request_body["store"], true); + assert_eq!(provider_request_body["input"][0]["id"], "msg-1"); + } + + #[test] + fn codex_responses_lite_additional_tools_contain_only_client_executed_specs() { + let mut provider_request_body = json!({ + "model": "gpt-5.6-sol", + "input": [{ + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "Use a tool"}] + }], + "tools": [ + {"type": "function", "name": "lookup", "parameters": {}}, + {"type": "custom", "name": "apply_patch", "description": "Apply a patch", "format": {}}, + {"type": "namespace", "name": "web", "description": "Web tools", "tools": []}, + {"type": "web_search"}, + {"type": "image_generation"}, + {"type": "tool_search", "execution": "client"}, + {"type": "tool_search", "execution": "server"}, + {"type": "tool_search"}, + {"type": "future_tool"}, + {"name": "missing_type"} + ] + }); + + apply_codex_openai_responses_special_body_edits( + &mut provider_request_body, + "codex", + "openai:responses", + None, + None, + ); + + assert!(provider_request_body.get("tools").is_none()); + assert_eq!( + provider_request_body["input"][0]["tools"], + json!([ + {"type": "function", "name": "lookup", "parameters": {}}, + {"type": "custom", "name": "apply_patch", "description": "Apply a patch", "format": {}}, + {"type": "namespace", "name": "web", "description": "Web tools", "tools": []}, + {"type": "tool_search", "execution": "client"} + ]) + ); + + let mut existing_additional_tools = json!({ + "model": "gpt-5.6-sol", + "input": [{ + "type": "additional_tools", + "role": "developer", + "tools": [ + {"type": "function", "name": "shell", "parameters": {}}, + {"type": "web_search"}, + {"type": "image_generation"} + ] + }] + }); + apply_codex_openai_responses_special_body_edits( + &mut existing_additional_tools, + "codex", + "openai:responses:compact", + None, + None, + ); + assert_eq!( + existing_additional_tools["input"][0]["tools"], + json!([{"type": "function", "name": "shell", "parameters": {}}]) + ); } #[test] @@ -1160,59 +2749,53 @@ mod tests { } #[test] - fn codex_responses_body_edits_derive_prompt_cache_key_from_session_metadata() { - let mut body_a = json!({ + fn codex_responses_body_edits_do_not_derive_prompt_cache_key_from_metadata() { + let mut body = json!({ "input": [{"role": "user", "content": "hello"}], "model": "gpt-5.4", "metadata": { "user_id": "{\"session_id\":\"session-a\",\"device_id\":\"device-a\"}" } }); - let mut body_b = json!({ - "input": [{"role": "user", "content": "hello again"}], - "model": "gpt-5.4", - "metadata": { - "user_id": "{\"session_id\":\"session-a\",\"device_id\":\"device-b\"}" - } - }); - let mut body_c = json!({ - "input": [{"role": "user", "content": "hello"}], - "model": "gpt-5.4", - "metadata": {"session_id": "session-b"} - }); apply_codex_openai_responses_special_body_edits( - &mut body_a, - "codex", - "openai:responses", - None, - Some("key-123"), - ); - apply_codex_openai_responses_special_body_edits( - &mut body_b, - "codex", - "openai:responses", - None, - Some("different-key"), - ); - apply_codex_openai_responses_special_body_edits( - &mut body_c, + &mut body, "codex", "openai:responses", None, Some("key-123"), ); - assert_eq!(body_a["prompt_cache_key"], body_b["prompt_cache_key"]); - assert_ne!(body_a["prompt_cache_key"], body_c["prompt_cache_key"]); - assert!(body_a.get("metadata").is_none()); - assert!(body_b.get("metadata").is_none()); - assert!(body_c.get("metadata").is_none()); + assert!(body.get("prompt_cache_key").is_none()); + assert!(body.get("metadata").is_none()); } #[test] - fn codex_responses_body_edits_derive_prompt_cache_key_from_cache_control_anchor() { - let mut body_a = json!({ + fn codex_responses_body_edits_do_not_treat_client_metadata_as_a_cache_key_source() { + let mut body = json!({ + "model": "gpt-5.6-sol", + "input": [], + "client_metadata": { + "session_id": "session-123", + "thread_id": "thread-123" + } + }); + + apply_codex_openai_responses_special_body_edits( + &mut body, + "codex", + "openai:responses", + None, + Some("key-123"), + ); + + assert!(body.get("prompt_cache_key").is_none()); + assert_eq!(body["client_metadata"]["thread_id"], "thread-123"); + } + + #[test] + fn codex_responses_body_edits_strip_cache_control_without_deriving_a_cache_key() { + let mut body = json!({ "input": [{ "type": "message", "role": "user", @@ -1228,62 +2811,17 @@ mod tests { }], "model": "gpt-5.4" }); - let mut body_b = json!({ - "input": [{ - "type": "message", - "role": "user", - "content": [{ - "type": "input_text", - "text": "stable project brief", - "cache_control": {"type": "ephemeral"} - }] - }, { - "type": "message", - "role": "user", - "content": [{"type": "input_text", "text": "new turn B"}] - }], - "model": "gpt-5.4" - }); - let mut body_c = json!({ - "input": [{ - "type": "message", - "role": "user", - "content": [{ - "type": "input_text", - "text": "different project brief", - "cache_control": {"type": "ephemeral"} - }] - }], - "model": "gpt-5.4" - }); apply_codex_openai_responses_special_body_edits( - &mut body_a, - "codex", - "openai:responses", - None, - Some("key-a"), - ); - apply_codex_openai_responses_special_body_edits( - &mut body_b, - "codex", - "openai:responses", - None, - Some("key-b"), - ); - apply_codex_openai_responses_special_body_edits( - &mut body_c, + &mut body, "codex", "openai:responses", None, Some("key-a"), ); - assert_eq!(body_a["prompt_cache_key"], body_b["prompt_cache_key"]); - assert_ne!(body_a["prompt_cache_key"], body_c["prompt_cache_key"]); - assert!(!body_a.to_string().contains("\"cache_control\"")); - assert!(!body_b.to_string().contains("\"cache_control\"")); - assert!(!body_c.to_string().contains("\"cache_control\"")); + assert!(body.get("prompt_cache_key").is_none()); + assert!(!body.to_string().contains("\"cache_control\"")); } #[test] @@ -1313,10 +2851,7 @@ mod tests { Some("key-a"), ); - assert!(provider_request_body - .get("prompt_cache_key") - .and_then(|value| value.as_str()) - .is_some_and(|value| !value.trim().is_empty())); + assert!(provider_request_body.get("prompt_cache_key").is_none()); assert!(!provider_request_body .to_string() .contains("\"cache_control\"")); @@ -1327,8 +2862,107 @@ mod tests { } #[test] - fn codex_responses_body_edits_derive_prompt_cache_key_from_stable_request_anchor() { - let mut body_a = json!({ + fn codex_responses_body_edits_preserve_cache_control_named_tool_schema_property() { + let mut provider_request_body = json!({ + "model": "gpt-5.4", + "input": [{ + "type": "message", + "role": "user", + "content": [{ + "type": "input_text", + "text": "inspect the payload", + "cache_control": {"type": "ephemeral"} + }] + }], + "tools": [{ + "type": "function", + "name": "inspect", + "cache_control": {"type": "ephemeral"}, + "parameters": { + "type": "object", + "properties": { + "cache_control": {"type": "string"} + } + } + }] + }); + + apply_codex_openai_responses_special_body_edits( + &mut provider_request_body, + "codex", + "openai:responses", + None, + Some("key-a"), + ); + + assert!(provider_request_body["input"][0]["content"][0] + .get("cache_control") + .is_none()); + assert!(provider_request_body["tools"][0] + .get("cache_control") + .is_none()); + assert_eq!( + provider_request_body["tools"][0]["parameters"]["properties"]["cache_control"]["type"], + "string" + ); + } + + #[test] + fn codex_responses_lite_compact_preserves_cache_control_named_tool_schema_property() { + let mut provider_request_body = json!({ + "model": "gpt-5.6-sol", + "instructions": "Use the inspection tool.", + "input": [{ + "type": "message", + "role": "user", + "content": [{ + "type": "input_text", + "text": "inspect the payload", + "cache_control": {"type": "ephemeral"} + }] + }], + "tools": [{ + "type": "function", + "name": "inspect", + "cache_control": {"type": "ephemeral"}, + "parameters": { + "type": "object", + "properties": { + "cache_control": {"type": "string"} + } + } + }] + }); + + apply_codex_openai_responses_special_body_edits( + &mut provider_request_body, + "codex", + "openai:responses:compact", + None, + Some("key-a"), + ); + + assert!(provider_request_body.get("tools").is_none()); + assert_eq!( + provider_request_body["input"][0]["type"], + "additional_tools" + ); + assert!(provider_request_body["input"][0]["tools"][0] + .get("cache_control") + .is_none()); + assert_eq!( + provider_request_body["input"][0]["tools"][0]["parameters"]["properties"] + ["cache_control"]["type"], + "string" + ); + assert!(provider_request_body["input"][2]["content"][0] + .get("cache_control") + .is_none()); + } + + #[test] + fn codex_responses_body_edits_do_not_synthesize_a_cache_key_from_request_content() { + let mut body = json!({ "input": [{ "type": "message", "role": "user", @@ -1347,71 +2981,33 @@ mod tests { }], "reasoning": {"effort": "medium"} }); - let mut body_b = json!({ - "input": [{ - "type": "message", - "role": "user", - "content": [{"type": "input_text", "text": "open workspace"}] - }, { - "type": "message", - "role": "user", - "content": [{"type": "input_text", "text": "new turn B"}] - }], - "model": "gpt-5.4", - "instructions": "Be concise.", - "tools": [{ - "type": "function", - "name": "shell", - "parameters": {"type": "object", "properties": {}} - }], - "reasoning": {"effort": "medium"} - }); - let mut body_c = json!({ - "input": [{ - "type": "message", - "role": "user", - "content": [{"type": "input_text", "text": "open another workspace"}] - }], - "model": "gpt-5.4", - "instructions": "Be concise.", - "tools": [{"type": "function", "name": "shell"}], - "reasoning": {"effort": "medium"} - }); - apply_codex_openai_responses_special_body_edits( - &mut body_a, - "codex", - "openai:responses", - None, - Some("key-a"), - ); - apply_codex_openai_responses_special_body_edits( - &mut body_b, - "codex", - "openai:responses", - None, - Some("key-a"), - ); - apply_codex_openai_responses_special_body_edits( - &mut body_c, + &mut body, "codex", "openai:responses", None, Some("key-a"), ); - assert_eq!(body_a["prompt_cache_key"], body_b["prompt_cache_key"]); - assert_ne!(body_a["prompt_cache_key"], body_c["prompt_cache_key"]); + assert!(body.get("prompt_cache_key").is_none()); } #[test] - fn compact_body_edits_strip_include_store_and_stream() { + fn compact_body_edits_apply_the_codex_request_projection() { let mut provider_request_body = json!({ "input": [], "model": "gpt-5.4", + "client_metadata": {"origin": "codex"}, "include": ["reasoning.encrypted_content"], "store": true, "stream": true, + "stream_options": {"reasoning_summary_delivery": "sequential_cutoff"}, + "tool_choice": "auto", + "parallel_tool_calls": true, + "reasoning": {"effort": "high"}, + "text": {"verbosity": "medium"}, + "tools": [{"type": "function", "name": "lookup"}], + "prompt_cache_key": "session:compact", }); apply_openai_responses_compact_special_body_edits( @@ -1419,37 +3015,230 @@ mod tests { "openai:responses:compact", ); - assert!(provider_request_body.get("include").is_none()); - assert!(provider_request_body.get("store").is_none()); - assert!(provider_request_body.get("stream").is_none()); + for field in [ + "client_metadata", + "include", + "store", + "stream", + "stream_options", + "tool_choice", + ] { + assert!(provider_request_body.get(field).is_none()); + } assert_eq!(provider_request_body["model"], json!("gpt-5.4")); + assert_eq!(provider_request_body["input"], json!([])); + assert_eq!(provider_request_body["parallel_tool_calls"], json!(true)); + assert_eq!(provider_request_body["reasoning"]["effort"], json!("high")); + assert_eq!(provider_request_body["text"]["verbosity"], json!("medium")); + assert_eq!(provider_request_body["tools"][0]["name"], json!("lookup")); + assert_eq!( + provider_request_body["prompt_cache_key"], + json!("session:compact") + ); } #[test] - fn codex_chat_body_edits_inject_reasoning_summary_defaults() { - let mut provider_request_body = json!({ - "input": [], - "model": "gpt-5.4" - }); + fn codex_responses_body_edits_omit_empty_instructions() { + for instructions in [None, Some(Value::Null), Some(json!(""))] { + for api_format in ["openai:responses", "openai:responses:compact"] { + let mut provider_request_body = json!({ + "model": "gpt-5.6-sol", + "input": [] + }); + if let Some(instructions) = instructions.clone() { + provider_request_body["instructions"] = instructions; + } - apply_codex_openai_responses_chat_body_edits( - &mut provider_request_body, + apply_codex_openai_responses_special_body_edits( + &mut provider_request_body, + "codex", + api_format, + None, + None, + ); + + assert!(provider_request_body.get("instructions").is_none()); + } + } + + let mut whitespace = json!({ + "model": "gpt-5.6-sol", + "input": [], + "instructions": " " + }); + apply_codex_openai_responses_special_body_edits( + &mut whitespace, "codex", "openai:responses", None, None, ); + assert!(whitespace.get("instructions").is_none()); + assert_eq!(whitespace["input"][0]["type"], "additional_tools"); + assert_eq!(whitespace["input"][1]["role"], "developer"); + assert_eq!(whitespace["input"][1]["content"][0]["text"], " "); + } + + #[test] + fn codex_compact_body_matches_the_typed_client_payload() { + let mut provider_request_body = json!({ + "model": "gpt-5.6-sol", + "input": [ + { + "id": "msg-1", + "type": "message", + "role": "user", + "content": [{ + "id": "content-id", + "type": "input_text", + "text": "hello", + "cache_control": {"type": "ephemeral"} + }] + }, + {"id": "call-1", "type": "function_call", "name": "lookup"}, + {"id": "future-1", "type": "future_item", "value": true} + ], + "top_logprobs": 5, + "max_output_tokens": 100, + "previous_response_id": "resp_123", + "prompt_cache_options": {"ttl": "30m"}, + "custom_extension": true, + "service_tier": "priority", + "text": {"verbosity": "medium"} + }); + + apply_codex_openai_responses_compact_body_edits( + &mut provider_request_body, + "codex", + "openai:responses:compact", + ); assert_eq!( - provider_request_body["reasoning"]["effort"], - json!("medium") + provider_request_body, + json!({ + "model": "gpt-5.6-sol", + "input": [ + { + "id": "msg-1", + "type": "message", + "role": "user", + "content": [{ + "id": "content-id", + "type": "input_text", + "text": "hello" + }] + }, + {"id": "call-1", "type": "function_call", "name": "lookup"}, + {"id": "future-1", "type": "future_item", "value": true} + ], + "parallel_tool_calls": true, + "service_tier": "priority", + "text": {"verbosity": "medium"} + }) ); - assert_eq!(provider_request_body["reasoning"]["summary"], json!("auto")); - assert_eq!( - provider_request_body["include"], - json!(["reasoning.encrypted_content"]) - ); - assert_eq!(provider_request_body["parallel_tool_calls"], json!(true)); + validate_codex_openai_responses_compact_request_contract( + &provider_request_body, + "codex", + "openai:responses:compact", + ) + .expect("projected Compact request should match the typed contract"); + } + + #[test] + fn codex_compact_ignores_image_generation_tool_choice_before_projection() { + let base = json!({ + "model": "gpt-5.4", + "input": [{"role": "user", "content": "generate image"}], + "tools": [{"type": "image_generation"}] + }); + let mut without_tool_choice = base.clone(); + let mut with_tool_choice = base; + with_tool_choice["tool_choice"] = json!({"type": "image_generation"}); + + for body in [&mut without_tool_choice, &mut with_tool_choice] { + apply_codex_openai_responses_special_body_edits( + body, + "codex", + "openai:responses:compact", + None, + None, + ); + } + + assert_eq!(with_tool_choice, without_tool_choice); + assert_eq!(with_tool_choice["model"], json!("gpt-5.4")); + assert!(with_tool_choice.get("tool_choice").is_none()); + } + + #[test] + fn codex_compact_contract_rejects_invalid_routing_mutations() { + for (field, value) in [ + ("input", json!("hello")), + ("parallel_tool_calls", json!("true")), + ("reasoning", json!("high")), + ("text", json!("medium")), + ] { + let mut body = json!({ + "model": "gpt-5.6-sol", + "input": [], + "parallel_tool_calls": true + }); + body[field] = value; + let violation = validate_codex_openai_responses_compact_request_contract( + &body, + "codex", + "openai:responses:compact", + ) + .expect_err("invalid Compact field type should be rejected"); + assert_eq!(violation.field, field); + } + } + + #[test] + fn codex_chat_body_edits_apply_model_reasoning_defaults() { + for (model, effort, summary) in [ + ("gpt-5.6-sol", "low", None), + ("gpt-5.6-terra", "medium", None), + ("gpt-5.6-luna", "medium", None), + ("gpt-5.4", "medium", None), + ("gpt-5.2", "medium", Some("auto")), + ] { + let mut provider_request_body = json!({ + "input": [], + "model": model + }); + + apply_codex_openai_responses_chat_body_edits( + &mut provider_request_body, + "codex", + "openai:responses", + None, + None, + ); + + assert_eq!(provider_request_body["reasoning"]["effort"], effort); + assert_eq!( + provider_request_body["reasoning"] + .get("summary") + .and_then(Value::as_str), + summary + ); + assert_eq!( + provider_request_body["include"], + json!(["reasoning.encrypted_content"]) + ); + let uses_responses_lite = model.starts_with("gpt-5.6"); + assert_eq!( + provider_request_body["parallel_tool_calls"], + json!(!uses_responses_lite) + ); + assert_eq!( + provider_request_body["reasoning"] + .get("context") + .and_then(Value::as_str), + uses_responses_lite.then_some("all_turns") + ); + } } #[test] @@ -1469,56 +3258,29 @@ mod tests { ); assert_eq!(provider_request_body["reasoning"]["effort"], json!("low")); - assert_eq!(provider_request_body["reasoning"]["summary"], json!("auto")); + assert!(provider_request_body["reasoning"].get("summary").is_none()); } #[test] - fn codex_image_body_edits_force_tool_choice_and_default_generate_tool_fields() { - let mut provider_request_body = json!({ - "input": [{ - "role": "user", - "content": "generate image" - }], - "tools": [{ - "type": "image_generation" - }], - "tool_choice": "auto" - }); + fn codex_gpt_5_6_public_reasoning_modes_default_to_medium_effort() { + for mode in ["standard", "pro"] { + let mut provider_request_body = json!({ + "input": [], + "model": "gpt-5.6-sol", + "reasoning": {"mode": mode} + }); - apply_codex_openai_responses_special_body_edits( - &mut provider_request_body, - "codex", - "openai:image", - None, - None, - ); + apply_codex_openai_responses_chat_body_edits( + &mut provider_request_body, + "codex", + "openai:responses", + None, + None, + ); - assert_eq!( - provider_request_body["tools"][0]["size"], - json!("1024x1024") - ); - assert_eq!(provider_request_body["tools"][0]["quality"], json!("high")); - assert_eq!( - provider_request_body["tools"][0]["background"], - json!("auto") - ); - assert_eq!( - provider_request_body["tools"][0]["output_format"], - json!("png") - ); - assert_eq!( - provider_request_body["tools"][0]["action"], - json!("generate") - ); - assert_eq!( - provider_request_body["model"], - json!(CODEX_OPENAI_IMAGE_INTERNAL_MODEL) - ); - assert_eq!(provider_request_body["stream"], json!(true)); - assert_eq!( - provider_request_body["tool_choice"]["type"], - json!("image_generation") - ); + assert_eq!(provider_request_body["reasoning"]["effort"], "medium"); + assert_eq!(provider_request_body["reasoning"]["mode"], mode); + } } #[test] @@ -1634,10 +3396,7 @@ mod tests { ); assert_eq!(provider_request_body["model"], json!(original_model)); - assert!( - provider_request_body.get("tool_choice").is_none(), - "tool_choice should not be injected when caller did not set it" - ); + assert_eq!(provider_request_body["tool_choice"], json!("auto")); assert_eq!( provider_request_body["tools"][0]["type"], json!("image_generation") @@ -1679,50 +3438,4 @@ mod tests { json!({"type": "function", "name": "shell"}) ); } - - #[test] - fn codex_image_body_edits_preserve_edit_action_without_generate_defaults() { - let mut provider_request_body = json!({ - "tools": [{ - "type": "image_generation", - "action": "edit", - "input_image_mask": { "image_url": "data:image/png;base64,mask" } - }], - "input": [{ - "role": "user", - "content": [{ - "type": "input_image", - "image_url": "data:image/png;base64,image" - }] - }], - "tool_choice": "auto" - }); - - apply_codex_openai_responses_special_body_edits( - &mut provider_request_body, - "codex", - "openai:image", - None, - None, - ); - - assert_eq!(provider_request_body["tools"][0]["action"], json!("edit")); - assert!(provider_request_body["tools"][0].get("size").is_none()); - assert!(provider_request_body["tools"][0].get("quality").is_none()); - assert!(provider_request_body["tools"][0] - .get("background") - .is_none()); - assert_eq!( - provider_request_body["tools"][0]["output_format"], - json!("png") - ); - assert_eq!( - provider_request_body["input"][0]["content"][0]["text"], - json!("Create a faithful variation of the provided image.") - ); - assert_eq!( - provider_request_body["tool_choice"]["type"], - json!("image_generation") - ); - } } diff --git a/crates/aether-ai-formats/src/formats/openai/responses/request.rs b/crates/aether-ai-formats/src/formats/openai/responses/request.rs index 25546f905..ef70667fc 100644 --- a/crates/aether-ai-formats/src/formats/openai/responses/request.rs +++ b/crates/aether-ai-formats/src/formats/openai/responses/request.rs @@ -4,22 +4,22 @@ use serde_json::{json, Map, Value}; use crate::{ formats::context::FormatContext, - formats::openai::shared::{ - map_thinking_budget_to_openai_reasoning_effort, OpenAiResponsesReasoningEffort, - }, + formats::openai::shared::map_thinking_budget_to_openai_reasoning_effort, protocol::canonical::{ canonical_response_format_to_openai_responses, canonical_tool_is_openai_custom, canonical_tool_use_to_openai_responses_input_item, is_claude_messages_request, is_claude_system_instruction, is_claude_thinking_block, is_claude_tool_result, - is_openai_responses_input_message, is_openai_thinking_block, media_data_or_url, - namespace_extension_object, openai_content_text, openai_extensions, + is_openai_responses_content_block, is_openai_responses_input_message, + is_openai_responses_raw_block, is_openai_responses_raw_content_block, + is_openai_thinking_block, media_data_or_url, namespace_extension_object, + openai_content_text, openai_extensions, openai_prompt_cache_breakpoint_from_extensions, openai_response_format_to_canonical, openai_responses_extension, openai_responses_generation_config, openai_responses_input_to_canonical_messages, - openai_responses_tool_choice_to_canonical, openai_responses_tools_to_canonical, - openai_tool_choice_raw_to_responses, strip_claude_billing_header, CanonicalContentBlock, - CanonicalInstruction, CanonicalRequest, CanonicalRole, CanonicalThinkingConfig, - CanonicalToolChoice, CanonicalToolDefinition, OPENAI_RESPONSES_EXTENSION_NAMESPACE, - OPENAI_RESPONSES_LEGACY_EXTENSION_NAMESPACE, + openai_responses_item_extension_object, openai_responses_tool_choice_to_canonical, + openai_responses_tools_to_canonical, openai_tool_choice_raw_to_responses, + strip_claude_billing_header, CanonicalContentBlock, CanonicalInstruction, CanonicalRequest, + CanonicalRole, CanonicalThinkingConfig, CanonicalToolChoice, CanonicalToolDefinition, + OPENAI_RESPONSES_EXTENSION_NAMESPACE, OPENAI_RESPONSES_LEGACY_EXTENSION_NAMESPACE, }, }; @@ -45,6 +45,77 @@ pub fn to_compact(request: &CanonicalRequest, ctx: &FormatContext) -> Option Result<(), OpenAiResponsesRequestContractViolation> { + if !crate::is_openai_responses_family_format(target_api_format) { + return Ok(()); + } + let Some(object) = body.as_object() else { + return Ok(()); + }; + let multi_agent_enabled = object + .get("multi_agent") + .and_then(Value::as_object) + .and_then(|multi_agent| multi_agent.get("enabled")) + .and_then(Value::as_bool) + == Some(true); + if !multi_agent_enabled { + return Ok(()); + } + if crate::is_openai_responses_compact_format(target_api_format) { + return Err(OpenAiResponsesRequestContractViolation { + field: "multi_agent", + reason: "OpenAI multi-agent requests are incompatible with Responses Compact", + }); + } + if object + .get("reasoning") + .and_then(Value::as_object) + .is_some_and(|reasoning| { + reasoning + .get("summary") + .is_some_and(|value| !value.is_null()) + }) + { + return Err(OpenAiResponsesRequestContractViolation { + field: "reasoning.summary", + reason: "OpenAI multi-agent requests do not support reasoning summaries", + }); + } + if object + .get("max_tool_calls") + .is_some_and(|value| !value.is_null()) + { + return Err(OpenAiResponsesRequestContractViolation { + field: "max_tool_calls", + reason: "OpenAI multi-agent requests do not support max_tool_calls", + }); + } + Ok(()) +} + pub fn from_raw(body_json: &Value) -> Option { let request = body_json.as_object()?; let mut canonical = CanonicalRequest { @@ -205,10 +276,26 @@ pub fn to_raw( )); apply_claude_responses_request_defaults(canonical, mapped_model, &mut output); if compact { - output.remove("stream"); + apply_compact_request_projection(&mut output); } output.remove("verbosity"); - Some(Value::Object(output)) + let output = Value::Object(output); + validate_openai_responses_request_contract( + &output, + if compact { + "openai:responses:compact" + } else { + "openai:responses" + }, + ) + .ok()?; + Some(output) +} + +pub(super) fn apply_compact_request_projection(output: &mut Map) { + for field in COMPACT_OMITTED_REQUEST_FIELDS { + output.remove(*field); + } } fn chat_openai_extension_object_to_responses( @@ -221,6 +308,9 @@ fn chat_openai_extension_object_to_responses( "service_tier", "safety_identifier", "prompt_cache_key", + "prompt_cache_options", + "prompt_cache_retention", + "user", ]; extensions .get("openai") @@ -337,7 +427,7 @@ fn canonical_messages_to_responses_input(canonical: &CanonicalRequest) -> Option input: arguments, extensions, } => { - flush_responses_message(&mut input, role, &mut content); + flush_responses_message(&mut input, role, &mut content, &message.extensions); saw_tool_item = true; let call_id = responses_tool_call_id(id, &mut next_generated_tool_call_index); let tool_name = responses_tool_name(name); @@ -350,10 +440,11 @@ fn canonical_messages_to_responses_input(canonical: &CanonicalRequest) -> Option tool_use_id, output, content_text, + is_error, extensions, .. } => { - flush_responses_message(&mut input, role, &mut content); + flush_responses_message(&mut input, role, &mut content, &message.extensions); saw_tool_item = true; let (tool_output, extra_user_content) = responses_tool_result_payload( output.as_ref(), @@ -362,12 +453,24 @@ fn canonical_messages_to_responses_input(canonical: &CanonicalRequest) -> Option )?; let call_id = responses_tool_result_call_id(tool_use_id, &mut pending_tool_call_ids)?; - input.push(json!({ - "type": responses_tool_result_item_type(extensions) - .unwrap_or("function_call_output"), - "call_id": call_id, - "output": tool_output, - })); + let mut item = Map::new(); + item.insert( + "type".to_string(), + Value::String( + responses_tool_result_item_type(extensions) + .unwrap_or("function_call_output") + .to_string(), + ), + ); + item.insert("call_id".to_string(), Value::String(call_id)); + item.insert("output".to_string(), tool_output); + if *is_error { + item.insert("is_error".to_string(), Value::Bool(true)); + } + let extension_fields = + openai_responses_item_extension_object(extensions, &item); + item.extend(extension_fields); + input.push(Value::Object(item)); if !extra_user_content.is_empty() { input.push(json!({ "type": "message", @@ -388,7 +491,12 @@ fn canonical_messages_to_responses_input(canonical: &CanonicalRequest) -> Option if role == "assistant" && is_openai_responses_reasoning_history_block(extensions) { - flush_responses_message(&mut input, role, &mut content); + flush_responses_message( + &mut input, + role, + &mut content, + &message.extensions, + ); if let Some(reasoning_item) = canonical_thinking_to_responses_reasoning_item( text, encrypted_content.as_deref(), @@ -406,6 +514,15 @@ fn canonical_messages_to_responses_input(canonical: &CanonicalRequest) -> Option })); } } + CanonicalContentBlock::Unknown { + payload, + extensions, + .. + } if is_openai_responses_raw_block(extensions) => { + flush_responses_message(&mut input, role, &mut content, &message.extensions); + input.push(payload.clone()); + saw_tool_item = true; + } other => { if let Some(part) = canonical_block_to_responses_input_part( other, @@ -418,25 +535,25 @@ fn canonical_messages_to_responses_input(canonical: &CanonicalRequest) -> Option } } if content.is_empty() && !saw_tool_item { - if role == "assistant" { - input.push(json!({ - "type": "message", - "role": role, - "content": [{ - "type": "output_text", - "text": "", - }], - })); + let content = if role == "assistant" { + json!([{ + "type": "output_text", + "text": "", + }]) } else { - input.push(json!({ - "type": "message", - "role": role, - "content": "", - })); - } + Value::String(String::new()) + }; + let mut item = Map::new(); + item.insert("type".to_string(), Value::String("message".to_string())); + item.insert("role".to_string(), Value::String(role.to_string())); + item.insert("content".to_string(), content); + let extension_fields = + openai_responses_item_extension_object(&message.extensions, &item); + item.extend(extension_fields); + input.push(Value::Object(item)); continue; } - flush_responses_message(&mut input, role, &mut content); + flush_responses_message(&mut input, role, &mut content, &message.extensions); } Some(input) } @@ -556,15 +673,22 @@ fn value_contains_json_word(value: &Value) -> bool { } } -fn flush_responses_message(input: &mut Vec, role: &str, content: &mut Vec) { +fn flush_responses_message( + input: &mut Vec, + role: &str, + content: &mut Vec, + extensions: &BTreeMap, +) { if content.is_empty() { return; } - input.push(json!({ - "type": "message", - "role": role, - "content": std::mem::take(content), - })); + let mut item = Map::new(); + item.insert("type".to_string(), Value::String("message".to_string())); + item.insert("role".to_string(), Value::String(role.to_string())); + item.insert("content".to_string(), Value::Array(std::mem::take(content))); + let extension_fields = openai_responses_item_extension_object(extensions, &item); + item.extend(extension_fields); + input.push(Value::Object(item)); } fn canonical_thinking_to_responses_reasoning_item( @@ -610,26 +734,36 @@ fn canonical_block_to_responses_input_part( strip_claude_billing_header_from_text: bool, ) -> Option { match block { - CanonicalContentBlock::Text { text, .. } => { + CanonicalContentBlock::Text { text, extensions } => { let text = if strip_claude_billing_header_from_text { strip_claude_billing_header(text) } else { text.clone() }; - if text.is_empty() { + if text.is_empty() && !is_openai_responses_content_block(extensions) { return None; } - Some(json!({ - "type": if role == "assistant" { "output_text" } else { "input_text" }, - "text": text, - })) + let mut part = Map::new(); + part.insert( + "type".to_string(), + Value::String(if role == "assistant" { + "output_text".to_string() + } else { + "input_text".to_string() + }), + ); + part.insert("text".to_string(), Value::String(text)); + insert_prompt_cache_breakpoint(&mut part, extensions); + let extension_fields = openai_responses_item_extension_object(extensions, &part); + part.extend(extension_fields); + Some(Value::Object(part)) } CanonicalContentBlock::Image { data, url, media_type, detail, - .. + extensions, } => { let mut item = Map::new(); item.insert( @@ -647,6 +781,9 @@ fn canonical_block_to_responses_input_part( if let Some(detail) = detail { item.insert("detail".to_string(), Value::String(detail.clone())); } + insert_prompt_cache_breakpoint(&mut item, extensions); + let extension_fields = openai_responses_item_extension_object(extensions, &item); + item.extend(extension_fields); Some(Value::Object(item)) } CanonicalContentBlock::File { @@ -655,7 +792,7 @@ fn canonical_block_to_responses_input_part( file_url, media_type, filename, - .. + extensions, } => { let mut item = Map::new(); item.insert("type".to_string(), Value::String("input_file".to_string())); @@ -675,15 +812,35 @@ fn canonical_block_to_responses_input_part( if let Some(value) = filename { item.insert("filename".to_string(), Value::String(value.clone())); } + insert_prompt_cache_breakpoint(&mut item, extensions); + let extension_fields = openai_responses_item_extension_object(extensions, &item); + item.extend(extension_fields); (item.len() > 1).then_some(Value::Object(item)) } - CanonicalContentBlock::Audio { data, format, .. } => Some(json!({ - "type": "input_audio", - "input_audio": { - "data": data.clone().unwrap_or_default(), - "format": format.clone().unwrap_or_else(|| "mp3".to_string()), - } - })), + CanonicalContentBlock::Audio { + data, + format, + extensions, + .. + } => { + let mut item = Map::new(); + item.insert("type".to_string(), Value::String("input_audio".to_string())); + item.insert( + "input_audio".to_string(), + json!({ + "data": data.clone().unwrap_or_default(), + "format": format.clone().unwrap_or_else(|| "mp3".to_string()), + }), + ); + let extension_fields = openai_responses_item_extension_object(extensions, &item); + item.extend(extension_fields); + Some(Value::Object(item)) + } + CanonicalContentBlock::Unknown { + payload, + extensions, + .. + } if is_openai_responses_raw_content_block(extensions) => Some(payload.clone()), CanonicalContentBlock::Unknown { raw_type, payload, .. } if raw_type == "refusal" => payload @@ -698,6 +855,15 @@ fn canonical_block_to_responses_input_part( } } +fn insert_prompt_cache_breakpoint( + part: &mut Map, + extensions: &BTreeMap, +) { + if let Some(value) = openai_prompt_cache_breakpoint_from_extensions(extensions) { + part.insert("prompt_cache_breakpoint".to_string(), value); + } +} + fn canonical_tools_to_responses(canonical: &CanonicalRequest) -> Vec { let mut tools = canonical .tools @@ -788,13 +954,8 @@ fn reasoning_config_to_responses(thinking: &CanonicalThinkingConfig) -> Option Option<&'static str> { - match effort.trim().to_ascii_lowercase().as_str() { - "max" => Some("xhigh"), - value => { - OpenAiResponsesReasoningEffort::parse(value).map(OpenAiResponsesReasoningEffort::as_str) - } - } +fn openai_responses_reasoning_effort(effort: &str) -> Option<&str> { + (!effort.trim().is_empty()).then_some(effort) } fn canonical_text_config_to_responses(canonical: &CanonicalRequest) -> Option { @@ -834,6 +995,11 @@ fn canonical_tool_to_responses(tool: &CanonicalToolDefinition) -> Value { tool.extensions .get(OPENAI_RESPONSES_LEGACY_EXTENSION_NAMESPACE) }) + .filter(|raw| { + raw.get("type") + .and_then(Value::as_str) + .is_some_and(|tool_type| !tool_type.eq_ignore_ascii_case("function")) + }) { return raw.clone(); } @@ -1304,7 +1470,7 @@ fn insert_number(output: &mut Map, key: &str, value: Option) #[cfg(test)] mod tests { - use super::to_raw; + use super::{from_raw, to_raw, COMPACT_OMITTED_REQUEST_FIELDS}; use crate::protocol::canonical::{ CanonicalContentBlock, CanonicalMessage, CanonicalRequest, CanonicalResponseFormat, CanonicalRole, @@ -1389,6 +1555,85 @@ mod tests { assert_eq!(body["input"][1]["content"][0]["text"], ""); } + #[test] + fn compact_request_uses_the_codex_request_projection() { + let source = json!({ + "model": "gpt-5.6-sol", + "input": [{ + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "hello"}] + }], + "client_metadata": {"origin": "codex"}, + "include": ["reasoning.encrypted_content"], + "store": false, + "stream": true, + "stream_options": {"reasoning_summary_delivery": "sequential_cutoff"}, + "tool_choice": "auto", + "parallel_tool_calls": true, + "reasoning": {"effort": "max", "context": "all_turns"}, + "text": {"verbosity": "medium"}, + "tools": [{ + "type": "function", + "name": "lookup", + "parameters": {"type": "object", "properties": {}} + }], + "service_tier": "priority", + "prompt_cache_key": "session:compact" + }); + let request = from_raw(&source).expect("canonical Responses request"); + + let regular = to_raw(&request, "gpt-5.6-sol", true, false).expect("Responses request body"); + let compact = to_raw(&request, "gpt-5.6-sol", false, true).expect("Compact request body"); + + for field in COMPACT_OMITTED_REQUEST_FIELDS { + assert!( + regular.get(*field).is_some(), + "regular request should contain {field}" + ); + assert!( + compact.get(*field).is_none(), + "Compact request should omit {field}" + ); + } + for field in [ + "model", + "parallel_tool_calls", + "reasoning", + "text", + "tools", + "service_tier", + "prompt_cache_key", + ] { + assert_eq!( + compact[field], regular[field], + "Compact should preserve {field}" + ); + } + assert_eq!(compact["input"], regular["input"]); + } + + #[test] + fn responses_request_preserves_compaction_trigger_input_item() { + let source = json!({ + "model": "gpt-5.6-sol", + "input": [ + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "compact"}] + }, + {"type": "compaction_trigger"} + ], + "stream": true + }); + let request = from_raw(&source).expect("canonical Responses request"); + let body = to_raw(&request, "gpt-5.6-sol", true, false).expect("Responses request body"); + + assert_eq!(body["input"].as_array().map(Vec::len), Some(2)); + assert_eq!(body["input"][1], json!({"type": "compaction_trigger"})); + } + #[test] fn responses_request_uses_empty_marker_for_empty_tool_output() { let request = CanonicalRequest { diff --git a/crates/aether-ai-formats/src/formats/openai/responses/response.rs b/crates/aether-ai-formats/src/formats/openai/responses/response.rs index bc1abe567..12048cda0 100644 --- a/crates/aether-ai-formats/src/formats/openai/responses/response.rs +++ b/crates/aether-ai-formats/src/formats/openai/responses/response.rs @@ -12,13 +12,15 @@ use crate::{ canonical_tool_use_to_openai_responses_item, canonical_usage_to_openai_responses_usage, flush_openai_responses_message_item, is_openai_responses_raw_block, is_openai_thinking_block, namespace_extension_object, openai_responses_extensions, - openai_responses_output_to_canonical_blocks, openai_usage_to_canonical, - CanonicalContentBlock, CanonicalResponse, CanonicalResponseOutput, CanonicalRole, - CanonicalStopReason, OPENAI_RESPONSES_EXTENSION_NAMESPACE, - OPENAI_RESPONSES_LEGACY_EXTENSION_NAMESPACE, + openai_responses_item_extension_object, openai_responses_output_to_canonical, + openai_responses_usage_to_canonical, openai_service_tier_extension, CanonicalContentBlock, + CanonicalResponse, CanonicalResponseOutput, CanonicalRole, CanonicalStopReason, + OPENAI_RESPONSES_EXTENSION_NAMESPACE, OPENAI_RESPONSES_LEGACY_EXTENSION_NAMESPACE, }, }; +const AETHER_RESPONSES_RAW_OUTPUT_KEY: &str = "openai_responses_raw_output"; + pub fn from(body: &Value, _ctx: &FormatContext) -> Option { from_raw(body) } @@ -38,7 +40,7 @@ pub fn from_raw(body_json: &Value) -> Option { { return None; } - let content = openai_responses_output_to_canonical_blocks(body.get("output"))?; + let (content, output_extensions) = openai_responses_output_to_canonical(body.get("output"))?; let has_tool_use = content .iter() .any(|block| matches!(block, CanonicalContentBlock::ToolUse { .. })); @@ -53,12 +55,18 @@ pub fn from_raw(body_json: &Value) -> Option { }; let mut extensions = openai_responses_extensions( body, - &["id", "object", "model", "output", "usage", "status"], + &[ + "id", "object", "model", "output", "usage", "status", "error", + ], ); if let Some(raw_status) = body.get("status").cloned() { canonical_extension_object_mut(&mut extensions, OPENAI_RESPONSES_EXTENSION_NAMESPACE) .insert("raw_status".to_string(), raw_status); } + if let Some(raw_output) = body.get("output").cloned() { + canonical_extension_object_mut(&mut extensions, "aether") + .insert(AETHER_RESPONSES_RAW_OUTPUT_KEY.to_string(), raw_output); + } Some(CanonicalResponse { id: body .get("id") @@ -75,11 +83,11 @@ pub fn from_raw(body_json: &Value) -> Option { role: CanonicalRole::Assistant, content: content.clone(), stop_reason: stop_reason.clone(), - extensions: BTreeMap::new(), + extensions: output_extensions, }], content, stop_reason, - usage: openai_usage_to_canonical(body.get("usage")), + usage: openai_responses_usage_to_canonical(body.get("usage")), extensions, }) } @@ -97,26 +105,35 @@ fn openai_responses_incomplete_stop_reason(body: &Map) -> Canonic } } -pub fn to_raw(canonical: &CanonicalResponse, report_context: &Value, _compact: bool) -> Value { +pub fn to_raw(canonical: &CanonicalResponse, report_context: &Value, compact: bool) -> Value { let mut response = Map::new(); let response_id = canonical.id.replace("chatcmpl", "resp"); response.insert("id".to_string(), Value::String(response_id.clone())); - response.insert("object".to_string(), Value::String("response".to_string())); - response.insert("status".to_string(), Value::String("completed".to_string())); - response.insert("model".to_string(), Value::String(canonical.model.clone())); - if let Some(raw_status) = canonical - .extensions - .get(OPENAI_RESPONSES_EXTENSION_NAMESPACE) - .or_else(|| { - canonical - .extensions - .get(OPENAI_RESPONSES_LEGACY_EXTENSION_NAMESPACE) - }) - .and_then(Value::as_object) - .and_then(|openai| openai.get("raw_status")) - .cloned() - { - response.insert("status".to_string(), raw_status); + response.insert( + "object".to_string(), + Value::String(if compact { + "response.compaction".to_string() + } else { + "response".to_string() + }), + ); + if !compact { + response.insert("status".to_string(), Value::String("completed".to_string())); + response.insert("model".to_string(), Value::String(canonical.model.clone())); + if let Some(raw_status) = canonical + .extensions + .get(OPENAI_RESPONSES_EXTENSION_NAMESPACE) + .or_else(|| { + canonical + .extensions + .get(OPENAI_RESPONSES_LEGACY_EXTENSION_NAMESPACE) + }) + .and_then(Value::as_object) + .and_then(|openai| openai.get("raw_status")) + .cloned() + { + response.insert("status".to_string(), raw_status); + } } let mut output = Vec::new(); @@ -260,6 +277,8 @@ pub fn to_raw(canonical: &CanonicalResponse, report_context: &Value, _compact: b if *is_error { item.insert("is_error".to_string(), Value::Bool(true)); } + let extension_fields = openai_responses_item_extension_object(extensions, &item); + item.extend(extension_fields); output.push(Value::Object(item)); } CanonicalContentBlock::Unknown { @@ -296,6 +315,15 @@ pub fn to_raw(canonical: &CanonicalResponse, report_context: &Value, _compact: b &response_id, &mut message_index, ); + if let Some(raw_output) = canonical + .extensions + .get("aether") + .and_then(Value::as_object) + .and_then(|aether| aether.get(AETHER_RESPONSES_RAW_OUTPUT_KEY)) + .and_then(Value::as_array) + { + output.clone_from(raw_output); + } response.insert("output".to_string(), Value::Array(output)); if let Some(usage) = &canonical.usage { response.insert( @@ -303,6 +331,22 @@ pub fn to_raw(canonical: &CanonicalResponse, report_context: &Value, _compact: b canonical_usage_to_openai_responses_usage(usage), ); } + if compact { + let created_at = canonical + .extensions + .get(OPENAI_RESPONSES_EXTENSION_NAMESPACE) + .or_else(|| { + canonical + .extensions + .get(OPENAI_RESPONSES_LEGACY_EXTENSION_NAMESPACE) + }) + .and_then(Value::as_object) + .and_then(|openai| openai.get("created_at").or_else(|| openai.get("created"))) + .and_then(openai_responses_timestamp_value) + .unwrap_or_else(openai_responses_current_timestamp); + response.insert("created_at".to_string(), Value::from(created_at)); + return Value::Object(response); + } if let Some(request_object) = report_context .get("original_request_body") .and_then(Value::as_object) @@ -327,9 +371,9 @@ pub fn to_raw(canonical: &CanonicalResponse, report_context: &Value, _compact: b response.insert(key.to_string(), value.clone()); } } - if let Some(service_tier) = request_object.get("service_tier").cloned() { - response.insert("service_tier".to_string(), service_tier); - } + } + if let Some(service_tier) = openai_service_tier_extension(&canonical.extensions).cloned() { + response.insert("service_tier".to_string(), service_tier); } let mut extension_fields = namespace_extension_object( &canonical.extensions, @@ -574,6 +618,42 @@ mod tests { assert_eq!(body["conversation"]["id"], "conv_123"); } + #[test] + fn compact_response_builder_emits_the_compaction_resource_shape() { + let mut extensions = BTreeMap::new(); + extensions.insert( + OPENAI_RESPONSES_EXTENSION_NAMESPACE.to_string(), + json!({"created_at": 123}), + ); + let response = CanonicalResponse { + id: "resp_compact".to_string(), + model: "gpt-5.6-sol".to_string(), + content: vec![CanonicalContentBlock::Text { + text: "compacted context".to_string(), + extensions: BTreeMap::new(), + }], + outputs: Vec::new(), + stop_reason: Some(CanonicalStopReason::EndTurn), + usage: None, + extensions, + }; + + let body = to_raw(&response, &json!({}), true); + let keys = body + .as_object() + .expect("Compact response should be an object") + .keys() + .map(String::as_str) + .collect::>(); + + assert_eq!(body["object"], "response.compaction"); + assert_eq!(body["created_at"], 123); + assert_eq!( + keys, + std::collections::BTreeSet::from(["created_at", "id", "object", "output"]) + ); + } + #[test] fn responses_response_parser_preserves_encrypted_reasoning_without_summary() { let body = json!({ diff --git a/crates/aether-ai-formats/src/formats/openai/responses/spec.rs b/crates/aether-ai-formats/src/formats/openai/responses/spec.rs index 4b6ba6fbb..9be663e05 100644 --- a/crates/aether-ai-formats/src/formats/openai/responses/spec.rs +++ b/crates/aether-ai-formats/src/formats/openai/responses/spec.rs @@ -1,5 +1,4 @@ use crate::contracts::{ - OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, OPENAI_RESPONSES_COMPACT_STREAM_SUCCESS_REPORT_KIND, OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND, OPENAI_RESPONSES_COMPACT_SYNC_SUCCESS_REPORT_KIND, OPENAI_RESPONSES_STREAM_PLAN_KIND, OPENAI_RESPONSES_STREAM_SUCCESS_REPORT_KIND, OPENAI_RESPONSES_SYNC_PLAN_KIND, OPENAI_RESPONSES_SYNC_SUCCESS_REPORT_KIND, @@ -43,13 +42,6 @@ pub fn resolve_stream_spec(plan_kind: &str) -> Option compact: false, require_streaming: true, }), - OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND => Some(LocalOpenAiResponsesSpec { - api_format: "openai:responses:compact", - decision_kind: OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, - report_kind: OPENAI_RESPONSES_COMPACT_STREAM_SUCCESS_REPORT_KIND, - compact: true, - require_streaming: true, - }), _ => None, } } @@ -68,11 +60,7 @@ mod tests { } #[test] - fn resolves_openai_responses_compact_stream_spec() { - let spec = resolve_stream_spec("openai_responses_compact_stream").expect("spec"); - assert_eq!(spec.api_format, "openai:responses:compact"); - assert_eq!(spec.report_kind, "openai_responses_compact_stream_success"); - assert!(spec.compact); - assert!(spec.require_streaming); + fn does_not_resolve_openai_responses_compact_as_streaming() { + assert!(resolve_stream_spec("openai_responses_compact_stream").is_none()); } } diff --git a/crates/aether-ai-formats/src/formats/openai/search.rs b/crates/aether-ai-formats/src/formats/openai/search.rs new file mode 100644 index 000000000..1553acbf0 --- /dev/null +++ b/crates/aether-ai-formats/src/formats/openai/search.rs @@ -0,0 +1,219 @@ +use serde_json::Value; + +const REQUEST_FIELDS: &[&str] = &[ + "id", + "model", + "reasoning", + "input", + "commands", + "settings", + "max_output_tokens", +]; +const REASONING_FIELDS: &[&str] = &["effort", "summary", "context"]; +const COMMAND_FIELDS: &[&str] = &[ + "search_query", + "image_query", + "open", + "click", + "find", + "screenshot", + "finance", + "weather", + "sports", + "time", + "response_length", +]; +const COMMAND_ITEM_FIELDS: &[(&str, &[&str])] = &[ + ("search_query", &["q", "recency", "domains"]), + ("image_query", &["q", "recency", "domains"]), + ("open", &["ref_id", "lineno"]), + ("click", &["ref_id", "id"]), + ("find", &["ref_id", "pattern"]), + ("screenshot", &["ref_id", "pageno"]), + ("finance", &["ticker", "type", "market"]), + ("weather", &["location", "start", "duration"]), + ( + "sports", + &[ + "tool", + "fn", + "league", + "team", + "opponent", + "date_from", + "date_to", + "num_games", + "locale", + ], + ), + ("time", &["utc_offset"]), +]; +const SETTINGS_FIELDS: &[&str] = &[ + "user_location", + "search_context_size", + "filters", + "image_settings", + "allowed_callers", + "external_web_access", +]; +const USER_LOCATION_FIELDS: &[&str] = &["type", "country", "region", "city", "timezone"]; +const FILTER_FIELDS: &[&str] = &["allowed_domains", "blocked_domains"]; +const IMAGE_SETTINGS_FIELDS: &[&str] = &["max_results", "caption"]; + +fn retain_object_fields(value: &mut Value, fields: &[&str]) { + if let Some(object) = value.as_object_mut() { + object.retain(|field, _| fields.contains(&field.as_str())); + } +} + +fn retain_array_object_fields( + object: &mut serde_json::Map, + key: &str, + fields: &[&str], +) { + if let Some(items) = object.get_mut(key).and_then(Value::as_array_mut) { + for item in items { + retain_object_fields(item, fields); + } + } +} + +pub fn apply_openai_search_request_projection(body: &mut Value, provider_api_format: &str) { + if !crate::api_format_alias_matches(provider_api_format, "openai:search") { + return; + } + let Some(body_object) = body.as_object_mut() else { + return; + }; + body_object.retain(|field, _| REQUEST_FIELDS.contains(&field.as_str())); + + if let Some(reasoning) = body_object.get_mut("reasoning") { + retain_object_fields(reasoning, REASONING_FIELDS); + } + if let Some(commands) = body_object + .get_mut("commands") + .and_then(Value::as_object_mut) + { + commands.retain(|field, _| COMMAND_FIELDS.contains(&field.as_str())); + for (key, fields) in COMMAND_ITEM_FIELDS { + retain_array_object_fields(commands, key, fields); + } + } + if let Some(settings) = body_object + .get_mut("settings") + .and_then(Value::as_object_mut) + { + settings.retain(|field, _| SETTINGS_FIELDS.contains(&field.as_str())); + if let Some(user_location) = settings.get_mut("user_location") { + retain_object_fields(user_location, USER_LOCATION_FIELDS); + } + if let Some(filters) = settings.get_mut("filters") { + retain_object_fields(filters, FILTER_FIELDS); + } + if let Some(image_settings) = settings.get_mut("image_settings") { + retain_object_fields(image_settings, IMAGE_SETTINGS_FIELDS); + } + } +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::apply_openai_search_request_projection; + + #[test] + fn projects_the_typed_search_request_contract() { + let mut body = json!({ + "id": "session-1", + "model": "gpt-5.6-sol", + "reasoning": { + "effort": "max", + "summary": "auto", + "context": "current_turn", + "mode": "pro" + }, + "input": "find documentation", + "commands": { + "search_query": [{"q": "Aether", "recency": 7, "unknown": true}], + "image_query": [{"q": "Aether UI", "domains": ["example.com"], "unknown": true}], + "open": [{"ref_id": "turn0search0", "lineno": 12, "unknown": true}], + "click": [{"ref_id": "turn0fetch0", "id": 3, "unknown": true}], + "find": [{"ref_id": "turn0fetch0", "pattern": "Aether", "unknown": true}], + "screenshot": [{"ref_id": "turn0fetch0", "pageno": 2, "unknown": true}], + "finance": [{"ticker": "OPENAI", "type": "equity", "market": "USA", "unknown": true}], + "weather": [{"location": "US, CA, San Francisco", "duration": 3, "unknown": true}], + "sports": [{"tool": "sports", "fn": "schedule", "league": "nba", "team": "GSW", "unknown": true}], + "time": [{"utc_offset": "+08:00", "unknown": true}], + "response_length": "short", + "unknown": true + }, + "settings": { + "user_location": {"type": "approximate", "country": "US", "unknown": true}, + "search_context_size": "high", + "filters": {"allowed_domains": ["openai.com"], "unknown": true}, + "image_settings": {"max_results": 3, "unknown": true}, + "allowed_callers": ["direct"], + "external_web_access": "live", + "unknown": true + }, + "max_output_tokens": 1024, + "store": false, + "stream": true, + "service_tier": "priority", + "unknown": true + }); + + apply_openai_search_request_projection(&mut body, "/v1/alpha/search"); + + assert_eq!( + body["reasoning"], + json!({ + "effort": "max", + "summary": "auto", + "context": "current_turn" + }) + ); + assert_eq!( + body["commands"], + json!({ + "search_query": [{"q": "Aether", "recency": 7}], + "image_query": [{"q": "Aether UI", "domains": ["example.com"]}], + "open": [{"ref_id": "turn0search0", "lineno": 12}], + "click": [{"ref_id": "turn0fetch0", "id": 3}], + "find": [{"ref_id": "turn0fetch0", "pattern": "Aether"}], + "screenshot": [{"ref_id": "turn0fetch0", "pageno": 2}], + "finance": [{"ticker": "OPENAI", "type": "equity", "market": "USA"}], + "weather": [{"location": "US, CA, San Francisco", "duration": 3}], + "sports": [{"tool": "sports", "fn": "schedule", "league": "nba", "team": "GSW"}], + "time": [{"utc_offset": "+08:00"}], + "response_length": "short" + }) + ); + assert_eq!( + body["settings"], + json!({ + "user_location": {"type": "approximate", "country": "US"}, + "search_context_size": "high", + "filters": {"allowed_domains": ["openai.com"]}, + "image_settings": {"max_results": 3}, + "allowed_callers": ["direct"], + "external_web_access": "live" + }) + ); + assert!(body.get("store").is_none()); + assert!(body.get("stream").is_none()); + assert!(body.get("service_tier").is_none()); + assert!(body.get("unknown").is_none()); + } + + #[test] + fn leaves_other_formats_unchanged() { + let mut body = json!({"model": "gpt-5.6-sol", "store": true}); + let expected = body.clone(); + + apply_openai_search_request_projection(&mut body, "openai:responses"); + + assert_eq!(body, expected); + } +} diff --git a/crates/aether-ai-formats/src/formats/openai/shared.rs b/crates/aether-ai-formats/src/formats/openai/shared.rs index 32eee7306..f8d216358 100644 --- a/crates/aether-ai-formats/src/formats/openai/shared.rs +++ b/crates/aether-ai-formats/src/formats/openai/shared.rs @@ -12,6 +12,7 @@ macro_rules! define_openai_reasoning_effort { Medium, High, XHigh, + Max, } impl $name { @@ -23,6 +24,7 @@ macro_rules! define_openai_reasoning_effort { "medium" => Some(Self::Medium), "high" => Some(Self::High), "xhigh" => Some(Self::XHigh), + "max" => Some(Self::Max), _ => None, } } @@ -35,6 +37,7 @@ macro_rules! define_openai_reasoning_effort { Self::Medium => "medium", Self::High => "high", Self::XHigh => "xhigh", + Self::Max => "max", } } } @@ -44,6 +47,29 @@ macro_rules! define_openai_reasoning_effort { define_openai_reasoning_effort!(OpenAiChatReasoningEffort); define_openai_reasoning_effort!(OpenAiResponsesReasoningEffort); +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum OpenAiPromptCacheRetention { + InMemory, + TwentyFourHours, +} + +impl OpenAiPromptCacheRetention { + pub fn parse(value: &str) -> Option { + match value.trim() { + "in_memory" => Some(Self::InMemory), + "24h" => Some(Self::TwentyFourHours), + _ => None, + } + } + + pub fn as_str(self) -> &'static str { + match self { + Self::InMemory => "in_memory", + Self::TwentyFourHours => "24h", + } + } +} + #[deprecated(note = "use OpenAiChatReasoningEffort or OpenAiResponsesReasoningEffort")] pub type OpenAiReasoningEffort = OpenAiChatReasoningEffort; @@ -137,3 +163,36 @@ pub fn extract_openai_reasoning_effort(request: &Map) -> Option gemini::embedding::request::from(body, ctx), FormatId::DoubaoEmbedding => doubao::embedding::request::from(body, ctx), FormatId::AliyunMultimodalEmbedding => aliyun::embedding::request::from(body, ctx), - FormatId::GeminiInteractions => None, + FormatId::OpenAiSearch | FormatId::GeminiInteractions => None, } .ok_or_else(|| FormatError::RequestParseFailed { format: source.as_str().to_string(), @@ -72,7 +72,7 @@ fn emit_request_inner( FormatId::GeminiEmbedding => gemini::embedding::request::to(request, ctx), FormatId::DoubaoEmbedding => doubao::embedding::request::to(request, ctx), FormatId::AliyunMultimodalEmbedding => aliyun::embedding::request::to(request, ctx), - FormatId::GeminiInteractions => None, + FormatId::OpenAiSearch | FormatId::GeminiInteractions => None, } .ok_or_else(|| FormatError::RequestEmitFailed { format: target.as_str().to_string(), @@ -114,7 +114,14 @@ pub fn convert_request_pure_with_context( ) -> Result, FormatError> { let pure_ctx = ctx.without_runtime_request_edits(); let request = parse_request(source_format, body, &pure_ctx)?; - validate_request_conversion(source_format, target_format, body, &request)?; + validate_openai_responses_target_contract(target_format, body)?; + validate_request_conversion( + source_format, + target_format, + body, + &request, + ctx.mapped_model.as_deref(), + )?; let value = emit_request_inner(target_format, &request, &pure_ctx)?; let report = build_request_conversion_report(source_format, target_format, body, &value); Ok(Converted { value, report }) @@ -128,8 +135,15 @@ pub fn convert_request( ) -> Result { let source = parse_format(source_format)?; let target = parse_format(target_format)?; - validate_runtime_request_conversion(source, target, body)?; + validate_openai_responses_target_contract(target_format, body)?; let mut request = parse_request(source_format, body, ctx)?; + validate_runtime_request_conversion( + source, + target, + body, + &request, + ctx.mapped_model.as_deref(), + )?; if let Some(mapped_model) = ctx .mapped_model .as_deref() @@ -144,7 +158,30 @@ fn validate_runtime_request_conversion( source: FormatId, target: FormatId, body: &Value, + _request: &CanonicalRequest, + mapped_model: Option<&str>, ) -> Result<(), FormatError> { + validate_openai_cross_format_store(source, target, body)?; + validate_openai_prompt_cache_contract(source, body, mapped_model)?; + validate_openai_reasoning_effort(source, target, body, mapped_model)?; + validate_openai_responses_cross_format_input(source, target, body)?; + validate_openai_responses_runtime_reasoning(source, target, body)?; + if matches!(source, FormatId::OpenAiChat) + && matches!( + target, + FormatId::OpenAiResponses | FormatId::OpenAiResponsesCompact + ) + && body + .as_object() + .is_some_and(openai_chat_request_has_unsupported_responses_cache_breakpoint) + { + return Err(FormatError::LossyConversionBlocked { + source_format: source.as_str().to_string(), + target_format: target.as_str().to_string(), + field: "messages[].content[].prompt_cache_breakpoint".to_string(), + reason: "OpenAI Responses supports prompt cache breakpoints only on input_text, input_image, and input_file blocks".to_string(), + }); + } if source == FormatId::ClaudeMessages { match target { FormatId::OpenAiChat @@ -177,6 +214,58 @@ fn validate_runtime_request_conversion( Ok(()) } +fn validate_openai_responses_target_contract( + target_format: &str, + body: &Value, +) -> Result<(), FormatError> { + if !matches!( + FormatId::parse(target_format), + Some(FormatId::OpenAiResponses | FormatId::OpenAiResponsesCompact) + ) { + return Ok(()); + } + crate::formats::openai::responses::request::validate_openai_responses_request_contract( + body, + target_format, + ) + .map_err(|violation| FormatError::InvalidTargetField { + format: crate::normalize_api_format_alias(target_format), + field: violation.field.to_string(), + reason: violation.reason.to_string(), + }) +} + +fn validate_openai_responses_runtime_reasoning( + source: FormatId, + target: FormatId, + body: &Value, +) -> Result<(), FormatError> { + if !matches!( + source, + FormatId::OpenAiResponses | FormatId::OpenAiResponsesCompact + ) || matches!( + target, + FormatId::OpenAiResponses | FormatId::OpenAiResponsesCompact + ) { + return Ok(()); + } + let Some(reasoning) = body.get("reasoning").and_then(Value::as_object) else { + return Ok(()); + }; + if let Some(field) = reasoning + .keys() + .find(|field| !matches!(field.as_str(), "effort" | "summary")) + { + return openai_responses_lossy_input( + source, + target, + format!("thinking.openai_responses.{field}"), + "target format cannot preserve this provider-specific Responses reasoning field", + ); + } + Ok(()) +} + pub fn parse_response( source_format: &str, body: &Value, @@ -191,6 +280,7 @@ pub fn parse_response( FormatId::ClaudeMessages => claude_messages::response::from(body, ctx), FormatId::GeminiGenerateContent => gemini_generate_content::response::from(body, ctx), FormatId::OpenAiEmbedding + | FormatId::OpenAiSearch | FormatId::JinaEmbedding | FormatId::OpenAiRerank | FormatId::JinaRerank @@ -229,6 +319,7 @@ fn emit_response_inner( FormatId::ClaudeMessages => claude_messages::response::to(response, ctx), FormatId::GeminiGenerateContent => gemini_generate_content::response::to(response, ctx), FormatId::OpenAiEmbedding + | FormatId::OpenAiSearch | FormatId::JinaEmbedding | FormatId::OpenAiRerank | FormatId::JinaRerank @@ -314,9 +405,12 @@ fn validate_request_conversion( target_format: &str, body: &Value, request: &CanonicalRequest, + mapped_model: Option<&str>, ) -> Result<(), FormatError> { let source = parse_format(source_format)?; let target = parse_format(target_format)?; + validate_openai_prompt_cache_contract(source, body, mapped_model)?; + validate_openai_reasoning_effort(source, target, body, mapped_model)?; if source == target { return Ok(()); } @@ -327,8 +421,9 @@ fn validate_request_conversion( return validate_rerank_request_conversion(source, target, request); } validate_known_standard_request_root_fields(source, target, body)?; + validate_openai_responses_cross_format_input(source, target, body)?; validate_cross_format_generation_target(source, target, request)?; - validate_openai_reasoning_effort(source, target, body)?; + validate_openai_cross_format_store(source, target, body)?; match (source, target) { (FormatId::OpenAiChat, FormatId::OpenAiResponses | FormatId::OpenAiResponsesCompact) => { validate_openai_chat_to_responses(body)?; @@ -347,6 +442,142 @@ fn validate_request_conversion( validate_cross_format_request_extensions(source, target, request) } +fn validate_openai_responses_cross_format_input( + source: FormatId, + target: FormatId, + body: &Value, +) -> Result<(), FormatError> { + if !matches!( + source, + FormatId::OpenAiResponses | FormatId::OpenAiResponsesCompact + ) || matches!( + target, + FormatId::OpenAiResponses | FormatId::OpenAiResponsesCompact + ) { + return Ok(()); + } + if body + .get("multi_agent") + .and_then(Value::as_object) + .and_then(|multi_agent| multi_agent.get("enabled")) + .and_then(Value::as_bool) + == Some(true) + { + return openai_responses_lossy_input( + source, + target, + "multi_agent".to_string(), + "target format cannot represent the OpenAI multi-agent request contract", + ); + } + + let Some(items) = body.get("input").and_then(Value::as_array) else { + return Ok(()); + }; + for (item_index, item) in items.iter().enumerate() { + let Some(item_object) = item.as_object() else { + if !item.is_string() { + return openai_responses_lossy_input( + source, + target, + format!("input[{item_index}]"), + "target format cannot represent this raw Responses input item", + ); + } + continue; + }; + let item_type = item_object + .get("type") + .and_then(Value::as_str) + .unwrap_or("message") + .trim() + .to_ascii_lowercase(); + if !matches!( + item_type.as_str(), + "message" + | "reasoning" + | "function_call" + | "custom_tool_call" + | "function_call_output" + | "custom_tool_call_output" + | "local_shell_call_output" + | "shell_call_output" + | "apply_patch_call_output" + | "computer_call_output" + ) { + return openai_responses_lossy_input( + source, + target, + format!("input[{item_index}]"), + "target format has no lossless mapping for this Responses input item type", + ); + } + if matches!(item_type.as_str(), "function_call" | "function_call_output") + && item_object.contains_key("caller") + { + return openai_responses_lossy_input( + source, + target, + format!("input[{item_index}].caller"), + "target format cannot preserve function call caller provenance", + ); + } + if item_type != "message" { + continue; + } + let Some(content) = item_object.get("content").and_then(Value::as_array) else { + continue; + }; + for (content_index, block) in content.iter().enumerate() { + let block_type = block + .get("type") + .and_then(Value::as_str) + .unwrap_or_default() + .trim() + .to_ascii_lowercase(); + if !matches!( + block_type.as_str(), + "input_text" + | "output_text" + | "text" + | "input_image" + | "output_image" + | "image_url" + | "input_file" + | "file" + | "input_audio" + | "audio" + | "output_audio" + | "reasoning" + | "thinking" + | "refusal" + ) { + return openai_responses_lossy_input( + source, + target, + format!("input[{item_index}].content[{content_index}]"), + "target format has no lossless mapping for this raw Responses content block", + ); + } + } + } + Ok(()) +} + +fn openai_responses_lossy_input( + source: FormatId, + target: FormatId, + field: String, + reason: &str, +) -> Result<(), FormatError> { + Err(FormatError::LossyConversionBlocked { + source_format: source.as_str().to_string(), + target_format: target.as_str().to_string(), + field, + reason: reason.to_string(), + }) +} + fn validate_response_conversion( source_format: &str, target_format: &str, @@ -361,9 +592,279 @@ fn validate_response_conversion( validate_source_response_stop_enums(source, target, body)?; validate_response_content_has_no_unknown_blocks(source, target, response)?; + validate_openai_responses_cross_format_response_extensions(source, target, response)?; validate_canonical_response_stop_reasons(source, target, response) } +fn validate_openai_responses_cross_format_response_extensions( + source: FormatId, + target: FormatId, + response: &CanonicalResponse, +) -> Result<(), FormatError> { + if !matches!( + source, + FormatId::OpenAiResponses | FormatId::OpenAiResponsesCompact + ) || matches!( + target, + FormatId::OpenAiResponses | FormatId::OpenAiResponsesCompact + ) { + return Ok(()); + } + + validate_openai_responses_response_extension_namespace( + source, + target, + "response", + &response.extensions, + |key, value| { + matches!(key, "raw_status" | "incomplete_details") + || (key == "output_text" && value.is_string()) + || (matches!(key, "created_at" | "completed_at") + && (value.is_number() || value.is_null())) + || (key == "error" && value.is_null()) + || (target == FormatId::OpenAiChat && key == "service_tier") + }, + )?; + + for output in &response.outputs { + validate_openai_responses_message_item_provenance(source, target, &output.extensions)?; + validate_openai_responses_response_extension_namespace( + source, + target, + "output[]", + &output.extensions, + |key, _| key == "message_items", + )?; + } + + if let Some(usage) = &response.usage { + validate_openai_responses_usage_extension_fields(source, target, &usage.extensions)?; + validate_openai_responses_response_extension_namespace( + source, + target, + "usage", + &usage.extensions, + |key, _| matches!(key, "input_tokens_details" | "output_tokens_details"), + )?; + } + + for block in &response.content { + if let CanonicalContentBlock::Thinking { + encrypted_content: Some(encrypted_content), + .. + } = block + { + if !encrypted_content.is_empty() { + return openai_responses_lossy_response( + source, + target, + "output[].encrypted_content", + "target format cannot preserve encrypted OpenAI reasoning content", + ); + } + } + let (location, extensions): (&str, &std::collections::BTreeMap) = match block + { + CanonicalContentBlock::Text { extensions, .. } => ("output[].content[]", extensions), + CanonicalContentBlock::Thinking { extensions, .. } => { + ("output[].reasoning", extensions) + } + CanonicalContentBlock::Image { extensions, .. } + | CanonicalContentBlock::File { extensions, .. } + | CanonicalContentBlock::Audio { extensions, .. } => ("output[].content[]", extensions), + CanonicalContentBlock::ToolUse { extensions, .. } => { + ("output[].function_call", extensions) + } + CanonicalContentBlock::ToolResult { extensions, .. } => { + ("output[].function_call_output", extensions) + } + CanonicalContentBlock::Unknown { .. } => continue, + }; + validate_openai_responses_response_extension_namespace( + source, + target, + location, + extensions, + |key, value| match block { + CanonicalContentBlock::Text { .. } => { + key == "annotations" + && (target == FormatId::OpenAiChat + || value.as_array().is_some_and(Vec::is_empty)) + } + CanonicalContentBlock::Thinking { .. } => key == "item_type", + CanonicalContentBlock::ToolUse { .. } => { + matches!(key, "item_id" | "item_type" | "status") + } + CanonicalContentBlock::ToolResult { .. } => key == "item_type", + CanonicalContentBlock::Image { .. } + | CanonicalContentBlock::File { .. } + | CanonicalContentBlock::Audio { .. } => key == "item_type", + CanonicalContentBlock::Unknown { .. } => false, + }, + )?; + } + Ok(()) +} + +fn validate_openai_responses_message_item_provenance( + source: FormatId, + target: FormatId, + extensions: &std::collections::BTreeMap, +) -> Result<(), FormatError> { + for (namespace, value) in extensions { + if !matches!(namespace.as_str(), "openai_responses" | "openai_cli") { + continue; + } + let Some(message_items) = value + .as_object() + .and_then(|object| object.get("message_items")) + else { + continue; + }; + let Some(message_items) = message_items.as_array() else { + return openai_responses_lossy_response( + source, + target, + "output[].message_items", + "Responses message item provenance is not an array", + ); + }; + for item in message_items { + let Some(item) = item.as_object() else { + return openai_responses_lossy_response( + source, + target, + "output[].message_items[]", + "Responses message item provenance entry is not an object", + ); + }; + let output_index = item + .get("output_index") + .and_then(Value::as_u64) + .map(|index| index.to_string()) + .unwrap_or_else(|| "?".to_string()); + let Some(fields) = item.get("fields").and_then(Value::as_object) else { + return openai_responses_lossy_response( + source, + target, + &format!("output[{output_index}]"), + "Responses message item provenance fields are not an object", + ); + }; + if let Some(field) = fields.keys().next() { + return openai_responses_lossy_response( + source, + target, + &format!("output[{output_index}].{field}"), + "target format cannot preserve this Responses message item field", + ); + } + } + } + Ok(()) +} + +fn validate_openai_responses_usage_extension_fields( + source: FormatId, + target: FormatId, + extensions: &std::collections::BTreeMap, +) -> Result<(), FormatError> { + for (namespace, value) in extensions { + if !matches!(namespace.as_str(), "openai_responses" | "openai_cli") { + continue; + } + let Some(usage) = value.as_object() else { + continue; + }; + for (details_key, mapped_fields) in [ + ( + "input_tokens_details", + &["cached_tokens", "cache_write_tokens"] as &[&str], + ), + ("output_tokens_details", &["reasoning_tokens"] as &[&str]), + ] { + let Some(details) = usage.get(details_key) else { + continue; + }; + let Some(details) = details.as_object() else { + return openai_responses_lossy_response( + source, + target, + &format!("usage.{details_key}"), + "Responses usage details are not an object", + ); + }; + for (field, value) in details { + let mapped = mapped_fields.contains(&field.as_str()) && value.as_u64().is_some(); + let cache_write_supported = + field != "cache_write_tokens" || target != FormatId::GeminiGenerateContent; + if mapped && cache_write_supported { + continue; + } + return openai_responses_lossy_response( + source, + target, + &format!("usage.{details_key}.{field}"), + "target format cannot preserve this Responses usage detail field", + ); + } + } + } + Ok(()) +} + +fn validate_openai_responses_response_extension_namespace( + source: FormatId, + target: FormatId, + location: &str, + extensions: &std::collections::BTreeMap, + is_safe: impl Fn(&str, &Value) -> bool, +) -> Result<(), FormatError> { + for (namespace, value) in extensions { + if namespace == "aether" { + continue; + } + let Some(object) = value.as_object() else { + return openai_responses_lossy_response( + source, + target, + location, + "Responses extension namespace is not an object", + ); + }; + for (key, value) in object { + let internal_openai_field = namespace == "openai" && key == "omit_reasoning_parts"; + if internal_openai_field + || (matches!(namespace.as_str(), "openai_responses" | "openai_cli") + && is_safe(key, value)) + { + continue; + } + return openai_responses_lossy_response( + source, + target, + &extension_field_path(location, namespace, key), + "target format cannot preserve this Responses provider field", + ); + } + } + Ok(()) +} + +fn openai_responses_lossy_response( + source: FormatId, + target: FormatId, + field: &str, + reason: &str, +) -> Result<(), FormatError> { + Err(FormatError::LossyConversionBlocked { + source_format: source.as_str().to_string(), + target_format: target.as_str().to_string(), + field: field.to_string(), + reason: reason.to_string(), + }) +} + fn validate_response_content_has_no_unknown_blocks( source: FormatId, target: FormatId, @@ -376,7 +877,14 @@ fn validate_response_content_has_no_unknown_blocks( .flat_map(|output| output.content.iter()), ) { if let CanonicalContentBlock::Unknown { raw_type, .. } = block { - if raw_type == "refusal" { + if raw_type == "refusal" + && matches!( + target, + FormatId::OpenAiChat + | FormatId::OpenAiResponses + | FormatId::OpenAiResponsesCompact + ) + { continue; } return Err(FormatError::LossyConversionBlocked { @@ -435,6 +943,7 @@ fn standard_request_root_field_is_audited(source: FormatId, key: &str) -> bool { | "prediction" | "presence_penalty" | "prompt_cache_key" + | "prompt_cache_options" | "prompt_cache_retention" | "reasoning_effort" | "response_format" @@ -466,10 +975,12 @@ fn standard_request_root_field_is_audited(source: FormatId, key: &str) -> bool { | "max_tool_calls" | "metadata" | "model" + | "multi_agent" | "parallel_tool_calls" | "previous_response_id" | "prompt" | "prompt_cache_key" + | "prompt_cache_options" | "prompt_cache_retention" | "reasoning" | "safety_identifier" @@ -546,6 +1057,7 @@ fn standard_request_root_field_is_audited(source: FormatId, key: &str) -> bool { | "tools" ), FormatId::OpenAiEmbedding + | FormatId::OpenAiSearch | FormatId::OpenAiRerank | FormatId::GeminiEmbedding | FormatId::JinaEmbedding @@ -838,6 +1350,23 @@ fn request_extension_key_is_cross_format_safe( namespace: &str, key: &str, ) -> bool { + if location == "messages[].content[]" + && key == "prompt_cache_breakpoint" + && matches!( + (source, target, namespace), + ( + FormatId::OpenAiChat, + FormatId::OpenAiResponses | FormatId::OpenAiResponsesCompact, + "openai", + ) | ( + FormatId::OpenAiResponses | FormatId::OpenAiResponsesCompact, + FormatId::OpenAiChat, + "openai_responses" | "openai_cli", + ) + ) + { + return true; + } if location == "tools[]" { return tool_extension_key_is_cross_format_safe(source, target, namespace, key); } @@ -858,7 +1387,9 @@ fn request_extension_key_is_cross_format_safe( | "service_tier" | "safety_identifier" | "prompt_cache_key" + | "prompt_cache_options" | "prompt_cache_retention" + | "user" | "verbosity", ) | ( FormatId::OpenAiChat, @@ -879,7 +1410,9 @@ fn request_extension_key_is_cross_format_safe( | "service_tier" | "safety_identifier" | "prompt_cache_key" + | "prompt_cache_options" | "prompt_cache_retention" + | "user" | "verbosity", ) | ( FormatId::ClaudeMessages, @@ -1005,6 +1538,7 @@ fn validate_source_response_stop_enums( FormatId::GeminiGenerateContent => validate_gemini_response_finish_reasons(body, target), FormatId::GeminiInteractions => Ok(()), FormatId::OpenAiEmbedding + | FormatId::OpenAiSearch | FormatId::OpenAiRerank | FormatId::GeminiEmbedding | FormatId::JinaEmbedding @@ -1584,6 +2118,89 @@ fn validate_openai_reasoning_effort( source: FormatId, target: FormatId, body: &Value, + mapped_model: Option<&str>, +) -> Result<(), FormatError> { + let same_openai_format = source == target + && matches!( + source, + FormatId::OpenAiChat | FormatId::OpenAiResponses | FormatId::OpenAiResponsesCompact + ); + let openai_cross_format = matches!( + (source, target), + ( + FormatId::OpenAiChat, + FormatId::OpenAiResponses | FormatId::OpenAiResponsesCompact + ) | ( + FormatId::OpenAiResponses | FormatId::OpenAiResponsesCompact, + FormatId::OpenAiChat + ) + ); + if !same_openai_format && !openai_cross_format { + return Ok(()); + } + crate::validate_openai_reasoning_request( + source.as_str(), + target.as_str(), + mapped_model.unwrap_or_default(), + body, + ) + .map_err(|violation| match violation.kind { + crate::OpenAiReasoningViolationKind::InvalidEnum => FormatError::InvalidEnumValue { + format: source.as_str().to_string(), + field: violation.field, + value: violation.value.unwrap_or_default(), + }, + crate::OpenAiReasoningViolationKind::InvalidType + | crate::OpenAiReasoningViolationKind::UnsupportedForModel => { + FormatError::InvalidTargetField { + format: source.as_str().to_string(), + field: violation.field, + reason: violation.reason, + } + } + }) +} + +fn validate_openai_prompt_cache_contract( + source: FormatId, + body: &Value, + mapped_model: Option<&str>, +) -> Result<(), FormatError> { + if !matches!( + source, + FormatId::OpenAiChat | FormatId::OpenAiResponses | FormatId::OpenAiResponsesCompact + ) { + return Ok(()); + } + let provider_model = mapped_model + .map(str::trim) + .filter(|value| !value.is_empty()) + .or_else(|| body.get("model").and_then(Value::as_str)) + .unwrap_or_default(); + crate::validate_openai_prompt_cache_request(source.as_str(), provider_model, body).map_err( + |violation| match violation.kind { + crate::OpenAiPromptCacheViolationKind::InvalidEnum => FormatError::InvalidEnumValue { + format: source.as_str().to_string(), + field: violation.field, + value: violation.value.unwrap_or_default(), + }, + crate::OpenAiPromptCacheViolationKind::InvalidType + | crate::OpenAiPromptCacheViolationKind::UnsupportedForModel + | crate::OpenAiPromptCacheViolationKind::UnsupportedContentBlock => { + FormatError::InvalidTargetField { + format: source.as_str().to_string(), + field: violation.field, + reason: violation.reason, + } + } + }, + ) +} + +fn validate_openai_cross_format_store( + source: FormatId, + target: FormatId, + body: &Value, ) -> Result<(), FormatError> { if !matches!( (source, target), @@ -1597,67 +2214,37 @@ fn validate_openai_reasoning_effort( ) { return Ok(()); } - let Some(object) = body.as_object() else { + let Some(value) = body.get("store") else { return Ok(()); }; - match source { - FormatId::OpenAiChat => validate_openai_reasoning_effort_value( - source.as_str(), - "reasoning_effort", - object.get("reasoning_effort"), - openai_chat_reasoning_effort_is_valid, - ), - FormatId::OpenAiResponses | FormatId::OpenAiResponsesCompact => { - let effort = object - .get("reasoning") - .and_then(Value::as_object) - .and_then(|reasoning| reasoning.get("effort")); - validate_openai_reasoning_effort_value( - source.as_str(), - "reasoning.effort", - effort, - openai_responses_reasoning_effort_is_valid, - ) - } - _ => Ok(()), + if target == FormatId::OpenAiResponsesCompact { + return if value.is_boolean() { + Ok(()) + } else { + Err(FormatError::InvalidTargetField { + format: source.as_str().to_string(), + field: "store".to_string(), + reason: "store must be a boolean".to_string(), + }) + }; } -} - -fn validate_openai_reasoning_effort_value( - format: &str, - field: &str, - value: Option<&Value>, - is_valid: fn(&str) -> bool, -) -> Result<(), FormatError> { - let Some(value) = value else { - return Ok(()); - }; - let Some(raw) = value.as_str() else { - return Err(FormatError::InvalidTargetField { - format: format.to_string(), - field: field.to_string(), - reason: "reasoning effort must be a string".to_string(), - }); - }; - if is_valid(raw) { - Ok(()) - } else { - Err(FormatError::InvalidEnumValue { - format: format.to_string(), - field: field.to_string(), - value: raw.to_string(), - }) + match value.as_bool() { + Some(false) => Ok(()), + Some(true) => Err(FormatError::LossyConversionBlocked { + source_format: source.as_str().to_string(), + target_format: target.as_str().to_string(), + field: "store".to_string(), + reason: "Chat and Responses assign different persistence behavior to store=true" + .to_string(), + }), + None => Err(FormatError::InvalidTargetField { + format: source.as_str().to_string(), + field: "store".to_string(), + reason: "store must be a boolean".to_string(), + }), } } -fn openai_chat_reasoning_effort_is_valid(value: &str) -> bool { - crate::formats::openai::shared::OpenAiChatReasoningEffort::parse(value).is_some() -} - -fn openai_responses_reasoning_effort_is_valid(value: &str) -> bool { - crate::formats::openai::shared::OpenAiResponsesReasoningEffort::parse(value).is_some() -} - fn validate_openai_chat_to_responses(body: &Value) -> Result<(), FormatError> { let Some(object) = body.as_object() else { return Ok(()); @@ -1670,7 +2257,6 @@ fn validate_openai_chat_to_responses(body: &Value) -> Result<(), FormatError> { "seed", "logprobs", "stream_options", - "user", ] { if object.contains_key(field) { return Err(FormatError::LossyConversionBlocked { @@ -1682,9 +2268,66 @@ fn validate_openai_chat_to_responses(body: &Value) -> Result<(), FormatError> { }); } } + if openai_chat_request_has_unsupported_responses_cache_breakpoint(object) { + return Err(FormatError::LossyConversionBlocked { + source_format: FormatId::OpenAiChat.as_str().to_string(), + target_format: FormatId::OpenAiResponses.as_str().to_string(), + field: "messages[].content[].prompt_cache_breakpoint".to_string(), + reason: "OpenAI Responses supports prompt cache breakpoints only on input_text, input_image, and input_file blocks".to_string(), + }); + } Ok(()) } +fn openai_chat_request_has_unsupported_responses_cache_breakpoint( + request: &Map, +) -> bool { + request + .get("messages") + .and_then(Value::as_array) + .into_iter() + .flatten() + .filter_map(Value::as_object) + .any(|message| { + let role = message + .get("role") + .and_then(Value::as_str) + .map(str::trim) + .map(str::to_ascii_lowercase); + message + .get("content") + .and_then(Value::as_array) + .into_iter() + .flatten() + .filter_map(Value::as_object) + .any(|part| { + part.contains_key("prompt_cache_breakpoint") + && !openai_chat_part_maps_to_cacheable_responses_input( + role.as_deref(), + part, + ) + }) + }) +} + +fn openai_chat_part_maps_to_cacheable_responses_input( + role: Option<&str>, + part: &Map, +) -> bool { + if !matches!(role, Some("user" | "system" | "developer")) { + return false; + } + matches!( + part.get("type") + .and_then(Value::as_str) + .unwrap_or("text") + .trim() + .to_ascii_lowercase() + .as_str(), + "text" | "input_text" | "image_url" | "input_image" | "file" | "input_file" + ) +} + fn validate_openai_responses_to_chat( body: &Value, request: &CanonicalRequest, @@ -1700,6 +2343,7 @@ fn validate_openai_responses_to_chat( "conversation", "background", "max_tool_calls", + "multi_agent", ] { if object.contains_key(field) { return Err(FormatError::LossyConversionBlocked { @@ -2740,7 +3384,7 @@ mod tests { } #[test] - fn pure_claude_to_openai_chat_clamps_max_output_effort_to_high() { + fn pure_claude_to_openai_chat_preserves_max_output_effort() { let body = json!({ "model": "claude-sonnet", "messages": [{"role": "user", "content": "hello"}], @@ -2754,7 +3398,7 @@ mod tests { .expect("pure conversion should succeed") .value; - assert_eq!(converted["reasoning_effort"], "xhigh"); + assert_eq!(converted["reasoning_effort"], "max"); } #[test] @@ -3240,6 +3884,666 @@ mod tests { assert_eq!(converted["tools"][1]["server_label"], "docs"); } + #[test] + fn pure_openai_responses_same_format_preserves_ptc_contract() { + let body = json!({ + "model": "gpt-5.6-sol", + "input": [ + { + "type": "program", + "id": "program_1", + "code": "result = lookup({\"id\": 7})" + }, + { + "type": "program_output", + "program_id": "program_1", + "output": "ok" + }, + { + "type": "agent_message", + "agent_id": "agent_1", + "content": [{"type": "input_text", "text": "delegated result"}] + }, + { + "type": "function_call", + "call_id": "call_1", + "name": "lookup", + "arguments": "{\"id\":7}", + "caller": {"type": "program", "id": "program_1"} + }, + { + "type": "function_call_output", + "call_id": "call_1", + "output": "found", + "caller": {"type": "program", "id": "program_1"} + } + ], + "tools": [ + { + "type": "function", + "name": "lookup", + "description": "Look up a record", + "parameters": { + "type": "object", + "properties": {"id": {"type": "integer"}}, + "required": ["id"] + }, + "strict": true, + "output_schema": { + "type": "object", + "properties": {"name": {"type": "string"}} + }, + "allowed_callers": ["programmatic"] + }, + {"type": "programmatic_tool_calling"} + ] + }); + + let converted = convert_request_pure("openai:responses", "openai:responses", &body) + .expect("same-format Responses conversion should preserve the PTC wire contract") + .value; + + assert_eq!(converted["input"], body["input"]); + assert_eq!(converted["tools"], body["tools"]); + assert_eq!( + converted["tools"][0]["allowed_callers"], + json!(["programmatic"]) + ); + assert_eq!(converted["tools"][1]["type"], "programmatic_tool_calling"); + } + + #[test] + fn pure_openai_responses_same_format_preserves_message_items() { + let body = json!({ + "model": "gpt-5.6-sol", + "input": [ + { + "type": "message", + "id": "msg_developer", + "status": "completed", + "role": "developer", + "content": [{ + "type": "input_text", + "text": "Use policy", + "prompt_cache_breakpoint": {"mode": "explicit"} + }] + }, + { + "type": "message", + "id": "msg_user", + "status": "completed", + "role": "user", + "content": [{"type": "input_text", "text": "hello"}] + }, + { + "type": "message", + "id": "msg_assistant", + "status": "completed", + "role": "assistant", + "phase": "commentary", + "content": [ + {"type": "output_text", "text": ""}, + {"type": "refusal", "refusal": "cannot comply"}, + {"type": "future_content", "payload": {"value": 1}} + ] + } + ] + }); + + let converted = convert_request_pure("openai:responses", "openai:responses", &body) + .expect("same-format Responses conversion should preserve message input items") + .value; + + assert_eq!(converted["input"], body["input"]); + } + + #[test] + fn pure_openai_responses_refusal_history_fails_closed_across_formats() { + let body = json!({ + "model": "gpt-5.6-sol", + "input": [{ + "type": "message", + "role": "assistant", + "content": [{"type": "refusal", "refusal": "cannot comply"}] + }] + }); + + for target in ["openai:chat", "claude:messages", "gemini:generate_content"] { + let error = convert_request_pure("openai:responses", target, &body) + .expect_err("target request format cannot represent refusal history"); + assert!(matches!( + error, + super::FormatError::LossyConversionBlocked { ref field, .. } + if field == "messages[].content[].type" + )); + } + } + + #[test] + fn pure_openai_responses_cross_format_rejects_opaque_input_item() { + let body = json!({ + "model": "gpt-5.6-sol", + "input": [{ + "type": "program", + "call_id": "program-call-1", + "fingerprint": "fp-program-1" + }] + }); + + let error = convert_request_pure("openai:responses", "openai:chat", &body) + .expect_err("provider-scoped opaque item must fail closed across formats"); + + assert!(matches!( + error, + super::FormatError::LossyConversionBlocked { ref field, .. } + if field == "input[0]" + )); + } + + #[test] + fn runtime_openai_responses_cross_format_rejects_unknown_content_block() { + let body = json!({ + "model": "gpt-5.6-sol", + "input": [{ + "type": "message", + "role": "user", + "content": [ + {"type": "input_text", "text": "hello"}, + {"type": "future_content_block", "payload": {"keep": true}} + ] + }] + }); + + let error = convert_request( + "openai:responses", + "claude:messages", + &body, + &FormatContext::default(), + ) + .expect_err("runtime conversion must fail closed for raw provider content blocks"); + + assert!(matches!( + error, + super::FormatError::LossyConversionBlocked { ref field, .. } + if field == "input[0].content[1]" + )); + } + + #[test] + fn openai_responses_rejects_official_multi_agent_incompatibilities() { + let cases = [ + ( + "openai:responses:compact", + json!({ + "model": "gpt-5.6-sol", + "input": "hello", + "multi_agent": {"enabled": true} + }), + "multi_agent", + ), + ( + "openai:responses", + json!({ + "model": "gpt-5.6-sol", + "input": "hello", + "multi_agent": {"enabled": true}, + "reasoning": {"effort": "high", "summary": "auto"} + }), + "reasoning.summary", + ), + ( + "openai:responses", + json!({ + "model": "gpt-5.6-sol", + "input": "hello", + "multi_agent": {"enabled": true}, + "max_tool_calls": 8 + }), + "max_tool_calls", + ), + ]; + + for (format, body, expected_field) in cases { + let error = convert_request_pure(format, format, &body) + .expect_err("invalid multi-agent combination must be rejected"); + assert!(matches!( + error, + super::FormatError::InvalidTargetField { ref field, .. } + if field == expected_field + )); + } + } + + #[test] + fn openai_responses_multi_agent_disabled_does_not_trigger_incompatibilities() { + let body = json!({ + "model": "gpt-5.6-sol", + "input": "hello", + "multi_agent": {"enabled": false}, + "reasoning": {"effort": "high", "summary": "auto"}, + "max_tool_calls": 8 + }); + + let converted = convert_request_pure("openai:responses", "openai:responses", &body) + .expect("disabled multi-agent config must not activate incompatibility rules"); + + assert_eq!(converted.value["multi_agent"]["enabled"], false); + } + + #[test] + fn openai_prompt_cache_contract_converts_responses_to_chat_without_loss() { + let body = json!({ + "model": "gpt-5.6-sol", + "store": false, + "user": "user_123", + "prompt_cache_options": {"mode": "explicit", "ttl": "30m"}, + "input": [{ + "type": "message", + "role": "user", + "content": [ + { + "type": "input_text", + "text": "stable text", + "prompt_cache_breakpoint": {"mode": "explicit"} + }, + { + "type": "input_image", + "image_url": "https://example.com/stable.png", + "prompt_cache_breakpoint": {"mode": "explicit"} + }, + { + "type": "input_file", + "file_id": "file_123", + "prompt_cache_breakpoint": {"mode": "explicit"} + } + ] + }] + }); + + let converted = convert_request_pure("openai:responses", "openai:chat", &body) + .expect("official Responses cache contract should map to Chat") + .value; + + assert_eq!( + converted["prompt_cache_options"], + body["prompt_cache_options"] + ); + assert_eq!(converted["store"], false); + assert_eq!(converted["user"], "user_123"); + let parts = converted["messages"][0]["content"] + .as_array() + .expect("chat content should remain structured"); + assert_eq!(parts.len(), 3); + assert!(parts + .iter() + .all(|part| { part["prompt_cache_breakpoint"] == json!({"mode": "explicit"}) })); + } + + #[test] + fn openai_prompt_cache_contract_converts_chat_to_responses_without_loss() { + let body = json!({ + "model": "gpt-5.6-sol", + "store": false, + "user": "user_123", + "prompt_cache_options": {"mode": "explicit", "ttl": "30m"}, + "messages": [{ + "role": "user", + "content": [ + { + "type": "text", + "text": "stable text", + "prompt_cache_breakpoint": {"mode": "explicit"} + }, + { + "type": "image_url", + "image_url": {"url": "https://example.com/stable.png"}, + "prompt_cache_breakpoint": {"mode": "explicit"} + }, + { + "type": "file", + "file": {"file_id": "file_123"}, + "prompt_cache_breakpoint": {"mode": "explicit"} + } + ] + }] + }); + + let converted = convert_request_pure("openai:chat", "openai:responses", &body) + .expect("official Chat cache contract should map to Responses") + .value; + + assert_eq!( + converted["prompt_cache_options"], + body["prompt_cache_options"] + ); + assert_eq!(converted["store"], false); + assert_eq!(converted["user"], "user_123"); + let parts = converted["input"][0]["content"] + .as_array() + .expect("Responses content should remain structured"); + assert_eq!(parts.len(), 3); + assert!(parts + .iter() + .all(|part| { part["prompt_cache_breakpoint"] == json!({"mode": "explicit"}) })); + } + + #[test] + fn openai_prompt_cache_contract_is_preserved_for_compact() { + let body = json!({ + "model": "gpt-5.6-sol", + "prompt_cache_key": "tenant:compact", + "prompt_cache_options": {"mode": "explicit", "ttl": "30m"}, + "messages": [{ + "role": "user", + "content": [{ + "type": "text", + "text": "stable text", + "prompt_cache_breakpoint": {"mode": "explicit"} + }] + }] + }); + + let converted = convert_request_pure("openai:chat", "openai:responses:compact", &body) + .expect("Compact supports the GPT-5.6 prompt cache contract") + .value; + + assert_eq!(converted["prompt_cache_key"], "tenant:compact"); + assert_eq!( + converted["prompt_cache_options"], + json!({"mode": "explicit", "ttl": "30m"}) + ); + assert_eq!( + converted["input"][0]["content"][0]["prompt_cache_breakpoint"], + json!({"mode": "explicit"}) + ); + } + + #[test] + fn openai_prompt_cache_contract_uses_gpt_5_6_options_and_rejects_invalid_values() { + let retention = json!({ + "model": "gpt-5.6-sol", + "messages": [{"role": "user", "content": "hello"}], + "prompt_cache_retention": "24h" + }); + let error = convert_request_pure("openai:chat", "openai:responses", &retention) + .expect_err("GPT-5.6 uses prompt_cache_options.ttl"); + assert!(matches!( + error, + super::FormatError::InvalidTargetField { ref field, .. } + if field == "prompt_cache_retention" + )); + + let invalid_options = json!({ + "model": "gpt-5.6-sol", + "input": "hello", + "prompt_cache_options": {"mode": "explicit", "ttl": "1h"} + }); + let error = convert_request_pure("openai:responses", "openai:chat", &invalid_options) + .expect_err("cross-format conversion must validate cache option enums"); + assert!(matches!( + error, + super::FormatError::InvalidEnumValue { ref field, .. } + if field == "prompt_cache_options.ttl" + )); + } + + #[test] + fn openai_prompt_cache_contract_preserves_earlier_model_retention() { + let body = json!({ + "model": "gpt-5.5", + "input": "hello", + "prompt_cache_retention": "24h" + }); + let converted = convert_request_pure("openai:responses", "openai:chat", &body) + .expect("GPT-5.5 24h retention should map between OpenAI APIs") + .value; + + assert_eq!(converted["prompt_cache_retention"], "24h"); + } + + #[test] + fn runtime_openai_prompt_cache_contract_uses_mapped_provider_model() { + let options = json!({ + "model": "gpt-5.5", + "messages": [{"role": "user", "content": "hello"}], + "prompt_cache_options": {"mode": "explicit", "ttl": "30m"} + }); + let converted = convert_request( + "openai:chat", + "openai:responses", + &options, + &FormatContext::default().with_mapped_model("gpt-5.6-sol"), + ) + .expect("mapped GPT-5.6 provider model should accept prompt_cache_options"); + assert_eq!( + converted["prompt_cache_options"], + options["prompt_cache_options"] + ); + + let error = convert_request( + "openai:chat", + "openai:responses", + &options, + &FormatContext::default().with_mapped_model("gpt-5.5"), + ) + .expect_err("mapped earlier provider model should reject prompt_cache_options"); + assert!(matches!( + error, + super::FormatError::InvalidTargetField { ref field, .. } + if field == "prompt_cache_options" + )); + + let retention = json!({ + "model": "gpt-5.6-sol", + "input": "hello", + "prompt_cache_retention": "24h" + }); + let converted = convert_request( + "openai:responses", + "openai:chat", + &retention, + &FormatContext::default().with_mapped_model("gpt-5.5"), + ) + .expect("mapped earlier provider model should accept its retention contract"); + assert_eq!(converted["prompt_cache_retention"], "24h"); + + let error = convert_request( + "openai:responses", + "openai:chat", + &retention, + &FormatContext::default().with_mapped_model("gpt-5.6-luna"), + ) + .expect_err("mapped GPT-5.6 provider model uses prompt_cache_options.ttl"); + assert!(matches!( + error, + super::FormatError::InvalidTargetField { ref field, .. } + if field == "prompt_cache_retention" + )); + } + + #[test] + fn openai_cross_format_store_only_maps_false() { + for (source, target, body) in [ + ( + "openai:chat", + "openai:responses", + json!({ + "model": "gpt-5.6-sol", + "messages": [{"role": "user", "content": "hello"}], + "store": true + }), + ), + ( + "openai:responses", + "openai:chat", + json!({ + "model": "gpt-5.6-sol", + "input": "hello", + "store": true + }), + ), + ] { + let error = convert_request_pure(source, target, &body) + .expect_err("store=true has different behavior across the OpenAI APIs"); + assert!(matches!( + error, + super::FormatError::LossyConversionBlocked { ref field, .. } + if field == "store" + )); + } + } + + #[test] + fn openai_chat_store_is_validated_then_omitted_for_compact() { + for store in [true, false] { + let converted = convert_request_pure( + "openai:chat", + "openai:responses:compact", + &json!({ + "model": "gpt-5.6-sol", + "messages": [{"role": "user", "content": "hello"}], + "store": store + }), + ) + .expect("Compact should validate and omit the transport-ineligible store field"); + assert!(converted.value.get("store").is_none()); + } + + let error = convert_request_pure( + "openai:chat", + "openai:responses:compact", + &json!({ + "model": "gpt-5.6-sol", + "messages": [{"role": "user", "content": "hello"}], + "store": "true" + }), + ) + .expect_err("Compact should reject a non-boolean store field"); + assert!(matches!( + error, + super::FormatError::InvalidTargetField { ref field, .. } + if field == "store" + )); + } + + #[test] + fn openai_chat_audio_cache_breakpoint_fails_closed_for_responses() { + let body = json!({ + "model": "gpt-5.6-sol", + "messages": [{ + "role": "user", + "content": [{ + "type": "input_audio", + "input_audio": {"data": "ZmFrZQ==", "format": "mp3"}, + "prompt_cache_breakpoint": {"mode": "explicit"} + }] + }] + }); + + let error = convert_request_pure("openai:chat", "openai:responses", &body) + .expect_err("Responses does not support cache breakpoints on audio blocks"); + + assert!(matches!( + error, + super::FormatError::LossyConversionBlocked { ref field, .. } + if field == "messages[].content[].prompt_cache_breakpoint" + )); + } + + #[test] + fn openai_chat_assistant_text_cache_breakpoint_fails_closed_for_responses() { + let body = json!({ + "model": "gpt-5.6-sol", + "messages": [{ + "role": "assistant", + "content": [{ + "type": "text", + "text": "previous output", + "prompt_cache_breakpoint": {"mode": "explicit"} + }] + }] + }); + + let error = convert_request_pure("openai:chat", "openai:responses", &body) + .expect_err("assistant text maps to output_text, which cannot carry a breakpoint"); + + assert!(matches!( + error, + super::FormatError::LossyConversionBlocked { ref field, .. } + if field == "messages[].content[].prompt_cache_breakpoint" + )); + } + + #[test] + fn runtime_responses_cross_format_blocks_advanced_reasoning_and_caller_fields() { + for (body, expected_field) in [ + ( + json!({ + "model": "gpt-5.6-sol", + "input": "hello", + "reasoning": {"effort": "high", "mode": "pro"} + }), + "thinking.openai_responses.mode", + ), + ( + json!({ + "model": "gpt-5.6-sol", + "input": "hello", + "reasoning": {"effort": "high", "context": "all_turns"} + }), + "thinking.openai_responses.context", + ), + ( + json!({ + "model": "gpt-5.6-sol", + "input": "hello", + "multi_agent": {"enabled": true} + }), + "multi_agent", + ), + ( + json!({ + "model": "gpt-5.6-sol", + "input": [{ + "type": "function_call", + "call_id": "call_1", + "name": "lookup", + "arguments": "{}", + "caller": {"type": "multi_agent", "id": "agent_1"} + }] + }), + "input[0].caller", + ), + ( + json!({ + "model": "gpt-5.6-sol", + "input": [{ + "type": "web_search_call", + "id": "search_1", + "action": {"query": "latest docs"} + }] + }), + "input[0]", + ), + ] { + let error = convert_request( + "openai:responses", + "openai:chat", + &body, + &FormatContext::default(), + ) + .expect_err("runtime conversion must fail closed for provider-only fields"); + assert!( + matches!( + error, + super::FormatError::LossyConversionBlocked { ref field, .. } + if field == expected_field + ), + "unexpected conversion error for {expected_field}: {error:?}" + ); + } + } + #[test] fn pure_openai_chat_to_claude_blocks_target_unsupported_generation_field() { let body = json!({ @@ -3336,40 +4640,233 @@ mod tests { } #[test] - fn pure_openai_cross_format_rejects_invalid_reasoning_effort_enum() { - let body = json!({ - "model": "gpt-source", + fn pure_openai_cross_format_preserves_reasoning_effort_contract() { + for effort in ["none", "low", "medium", "high", "xhigh", "max"] { + let chat = json!({ + "model": "gpt-5.6-sol", + "messages": [{"role": "user", "content": "hello"}], + "reasoning_effort": effort + }); + let responses = convert_request_pure("openai:chat", "openai:responses", &chat) + .expect("Chat effort should convert to Responses") + .value; + assert_eq!(responses["reasoning"]["effort"], effort); + + let responses = json!({ + "model": "gpt-5.6-sol", + "input": [{"role": "user", "content": "hello"}], + "reasoning": {"effort": effort} + }); + let chat = convert_request_pure("openai:responses", "openai:chat", &responses) + .expect("Responses effort should convert to Chat") + .value; + assert_eq!(chat["reasoning_effort"], effort); + } + } + + #[test] + fn runtime_openai_cross_format_preserves_gpt_5_6_reasoning_efforts() { + for target in ["openai:responses", "openai:responses:compact"] { + for effort in ["none", "low", "medium", "high", "xhigh", "max"] { + let chat = json!({ + "model": "deployment-alias", + "messages": [{"role": "user", "content": "hello"}], + "reasoning_effort": effort + }); + let responses = convert_request( + "openai:chat", + target, + &chat, + &FormatContext::default().with_mapped_model("gpt-5.6-sol"), + ) + .expect("runtime Chat effort should convert to Responses"); + assert_eq!(responses["reasoning"]["effort"], effort, "target: {target}"); + } + } + } + + #[test] + fn runtime_reasoning_effort_uses_concrete_mapped_model_as_authoritative_capability() { + let alias_to_gpt_5_6 = json!({ + "model": "deployment-alias", "messages": [{"role": "user", "content": "hello"}], "reasoning_effort": "max" }); + let converted = convert_request( + "openai:chat", + "openai:responses", + &alias_to_gpt_5_6, + &FormatContext::default().with_mapped_model("gpt-5.6-sol"), + ) + .expect("a concrete GPT-5.6 mapped target should accept max"); + assert_eq!(converted["model"], "gpt-5.6-sol"); + assert_eq!(converted["reasoning"]["effort"], "max"); - let error = convert_request_pure("openai:chat", "openai:responses", &body) - .expect_err("invalid OpenAI enum should fail closed"); - + let gpt_5_6_to_gpt_5_4 = json!({ + "model": "gpt-5.6-sol", + "messages": [{"role": "user", "content": "hello"}], + "reasoning_effort": "max" + }); + let error = convert_request( + "openai:chat", + "openai:responses", + &gpt_5_6_to_gpt_5_4, + &FormatContext::default().with_mapped_model("gpt-5.4"), + ) + .expect_err("a concrete GPT-5.4 mapped target must reject max"); assert!(matches!( error, - super::FormatError::InvalidEnumValue { ref field, ref value, .. } - if field == "reasoning_effort" && value == "max" + super::FormatError::InvalidTargetField { ref field, .. } + if field == "reasoning_effort" )); } #[test] - fn pure_openai_responses_to_chat_rejects_invalid_reasoning_effort_enum() { - let body = json!({ - "model": "gpt-source", + fn pure_conversion_reasoning_effort_uses_the_concrete_mapped_model() { + let alias_to_gpt_5_6 = json!({ + "model": "deployment-alias", + "messages": [{"role": "user", "content": "hello"}], + "reasoning_effort": "max" + }); + let converted = convert_request_pure_with_context( + "openai:chat", + "openai:responses", + &alias_to_gpt_5_6, + &FormatContext::default().with_mapped_model("gpt-5.6-sol"), + ) + .expect("a concrete GPT-5.6 target should accept max"); + assert_eq!(converted.value["reasoning"]["effort"], "max"); + + let gpt_5_6_to_gpt_5_4 = json!({ + "model": "gpt-5.6-sol", + "messages": [{"role": "user", "content": "hello"}], + "reasoning_effort": "max" + }); + let error = convert_request_pure_with_context( + "openai:chat", + "openai:responses", + &gpt_5_6_to_gpt_5_4, + &FormatContext::default().with_mapped_model("gpt-5.4"), + ) + .expect_err("a concrete GPT-5.4 target must reject max"); + assert!(matches!( + error, + super::FormatError::InvalidTargetField { ref field, .. } + if field == "reasoning_effort" + )); + } + + #[test] + fn runtime_openai_cross_format_enforces_known_efforts_and_preserves_custom_efforts() { + let alias_minimal = json!({ + "model": "deployment-alias", + "messages": [{"role": "user", "content": "hello"}], + "reasoning_effort": "minimal" + }); + for target in ["openai:responses", "openai:responses:compact"] { + let error = convert_request( + "openai:chat", + target, + &alias_minimal, + &FormatContext::default().with_mapped_model("gpt-5.6-terra"), + ) + .expect_err("mapped GPT-5.6 deployments must reject minimal effort"); + assert!(matches!( + error, + super::FormatError::InvalidTargetField { ref field, .. } + if field == "reasoning_effort" + )); + } + + let custom = json!({ + "model": "gpt-5.6-sol", "input": [{"role": "user", "content": "hello"}], - "reasoning": { - "effort": "max" - } + "reasoning": {"effort": "future"} + }); + for source in ["openai:responses", "openai:responses:compact"] { + let converted = convert_request( + source, + "openai:chat", + &custom, + &FormatContext::default().with_mapped_model("gpt-5.6-sol"), + ) + .expect("model-advertised effort should survive runtime conversion"); + assert_eq!(converted["reasoning_effort"], "future"); + } + + let ultra = json!({ + "model": "gpt-5.6-sol", + "input": [{"role": "user", "content": "hello"}], + "reasoning": {"effort": "ultra"} + }); + for source in ["openai:responses", "openai:responses:compact"] { + let error = convert_request( + source, + "openai:chat", + &ultra, + &FormatContext::default().with_mapped_model("gpt-5.6-sol"), + ) + .expect_err("Codex local ultra preset should not enter the OpenAI wire contract"); + assert!(matches!( + error, + super::FormatError::InvalidEnumValue { ref field, ref value, .. } + if field == "reasoning.effort" && value == "ultra" + )); + } + } + + #[test] + fn pure_openai_cross_format_preserves_custom_reasoning_effort() { + let body = json!({ + "model": "gpt-5.6-sol", + "messages": [{"role": "user", "content": "hello"}], + "reasoning_effort": "future" }); - let error = convert_request_pure("openai:responses", "openai:chat", &body) - .expect_err("invalid Responses reasoning enum should fail closed"); + let converted = convert_request_pure("openai:chat", "openai:responses", &body) + .expect("custom effort should remain wire-compatible"); + assert_eq!(converted.value["reasoning"]["effort"], "future"); + let ultra = json!({ + "model": "gpt-5.6-sol", + "messages": [{"role": "user", "content": "hello"}], + "reasoning_effort": "ultra" + }); + let error = convert_request_pure("openai:chat", "openai:responses", &ultra) + .expect_err("Codex local ultra preset should not enter the OpenAI wire contract"); assert!(matches!( error, super::FormatError::InvalidEnumValue { ref field, ref value, .. } - if field == "reasoning.effort" && value == "max" + if field == "reasoning_effort" && value == "ultra" + )); + } + + #[test] + fn pure_openai_cross_format_rejects_gpt_5_6_minimal_effort() { + let chat = json!({ + "model": "gpt-5.6-sol", + "messages": [{"role": "user", "content": "hello"}], + "reasoning_effort": "minimal" + }); + let error = convert_request_pure("openai:chat", "openai:responses", &chat) + .expect_err("GPT-5.6 does not publish minimal as a supported effort"); + assert!(matches!( + error, + super::FormatError::InvalidTargetField { ref field, .. } + if field == "reasoning_effort" + )); + + let responses = json!({ + "model": "gpt-5.6-terra", + "input": [{"role": "user", "content": "hello"}], + "reasoning": {"effort": "minimal"} + }); + let error = convert_request_pure("openai:responses", "openai:chat", &responses) + .expect_err("GPT-5.6 does not publish minimal as a supported effort"); + assert!(matches!( + error, + super::FormatError::InvalidTargetField { ref field, .. } + if field == "reasoning.effort" )); } @@ -3588,6 +5085,7 @@ mod tests { { "type": "message", "id": "msg_123", + "status": "completed", "role": "assistant", "content": [{"type": "output_text", "text": "done"}] } @@ -3598,10 +5096,49 @@ mod tests { .expect("same-format Responses response should preserve raw output items") .value; - assert_eq!(converted["output"][0]["type"], "file_search_call"); - assert_eq!(converted["output"][0]["results"][0]["file_id"], "file_123"); - assert_eq!(converted["output"][1]["type"], "code_interpreter_call"); - assert_eq!(converted["output"][2]["content"][0]["text"], "done"); + assert_eq!(converted["output"], body["output"]); + } + + #[test] + fn pure_openai_responses_response_preserves_ptc_and_blocks_lossy_targets() { + let body = json!({ + "id": "resp_ptc", + "object": "response", + "model": "gpt-5.6-sol", + "status": "completed", + "output": [ + { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "status": "completed", + "name": "lookup", + "arguments": "{\"id\":7}", + "caller": {"type": "program", "id": "program_1"} + }, + { + "type": "function_call_output", + "call_id": "call_1", + "output": "found", + "caller": {"type": "program", "id": "program_1"} + } + ] + }); + + let same_format = convert_response_pure("openai:responses", "openai:responses", &body) + .expect("same-format response should preserve PTC output items") + .value; + assert_eq!(same_format["output"], body["output"]); + + for target in ["openai:chat", "claude:messages", "gemini:generate_content"] { + let error = convert_response_pure("openai:responses", target, &body) + .expect_err("target response format cannot preserve PTC caller provenance"); + assert!(matches!( + error, + super::FormatError::LossyConversionBlocked { ref field, .. } + if field.ends_with(".caller") + )); + } } #[test] @@ -3629,6 +5166,111 @@ mod tests { )); } + #[test] + fn pure_openai_responses_response_refusal_only_maps_to_openai_chat() { + let body = json!({ + "id": "resp_refusal", + "object": "response", + "model": "gpt-5.6-sol", + "status": "completed", + "output": [{ + "type": "message", + "id": "msg_refusal", + "status": "completed", + "role": "assistant", + "content": [{"type": "refusal", "refusal": "cannot comply"}] + }] + }); + + let chat = convert_response_pure("openai:responses", "openai:chat", &body) + .expect("OpenAI Chat has a refusal response field") + .value; + assert_eq!(chat["choices"][0]["message"]["refusal"], "cannot comply"); + + for target in ["claude:messages", "gemini:generate_content"] { + let error = convert_response_pure("openai:responses", target, &body) + .expect_err("target response format cannot preserve refusal content"); + assert!(matches!( + error, + super::FormatError::LossyConversionBlocked { ref field, .. } + if field == "output[].type" + )); + } + } + + #[test] + fn pure_openai_responses_response_blocks_unknown_message_item_fields_cross_format() { + let body = json!({ + "id": "resp_message_extension", + "object": "response", + "model": "gpt-5.6-sol", + "status": "completed", + "output": [{ + "type": "message", + "id": "msg_extension", + "status": "completed", + "role": "assistant", + "future_message_option": {"mode": "verified"}, + "content": [{"type": "output_text", "text": "done"}] + }] + }); + + let same_format = convert_response_pure("openai:responses", "openai:responses", &body) + .expect("same-format response should preserve message item fields") + .value; + assert_eq!(same_format["output"], body["output"]); + + for target in ["openai:chat", "claude:messages", "gemini:generate_content"] { + let error = convert_response_pure("openai:responses", target, &body) + .expect_err("target response format cannot preserve message item extensions"); + assert!(matches!( + error, + super::FormatError::LossyConversionBlocked { ref field, .. } + if field == "output[0].future_message_option" + )); + } + } + + #[test] + fn pure_openai_responses_response_preserves_usage_and_audits_unknown_details() { + let body = json!({ + "id": "resp_usage_extension", + "object": "response", + "model": "gpt-5.6-sol", + "status": "completed", + "output": [], + "usage": { + "input_tokens": 120, + "output_tokens": 30, + "total_tokens": 150, + "input_tokens_details": { + "cached_tokens": 40, + "future_cache_tokens": 3 + }, + "output_tokens_details": { + "reasoning_tokens": 20, + "future_reasoning_tokens": 2 + }, + "future_usage_option": "metered" + } + }); + + let same_format = convert_response_pure("openai:responses", "openai:responses", &body) + .expect("same-format response should preserve usage extensions") + .value; + assert_eq!(same_format["usage"], body["usage"]); + + for target in ["openai:chat", "claude:messages", "gemini:generate_content"] { + let error = convert_response_pure("openai:responses", target, &body) + .expect_err("target response format cannot preserve unknown usage details"); + assert!(matches!( + error, + super::FormatError::LossyConversionBlocked { ref field, .. } + if field == "usage.input_tokens_details.future_cache_tokens" + )); + } + } + #[test] fn pure_claude_response_same_format_preserves_unknown_stop_reason() { let body = json!({ 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 53ace4d28..a00f053af 100644 --- a/crates/aether-ai-formats/src/formats/shared/error_body.rs +++ b/crates/aether-ai-formats/src/formats/shared/error_body.rs @@ -38,7 +38,11 @@ pub fn build_core_error_body_for_client_format( error_object.insert("message".to_string(), Value::String(message.to_string())); match aether_ai_formats::normalize_api_format_alias(client_api_format).as_str() { - "openai:chat" | "openai:responses" | "openai:responses:compact" | "openai:embedding" => { + "openai:chat" + | "openai:responses" + | "openai:responses:compact" + | "openai:search" + | "openai:embedding" => { error_object.insert( "type".to_string(), Value::String(map_local_sync_error_kind_to_openai_type(kind).to_string()), @@ -157,6 +161,21 @@ mod tests { assert_eq!(body["error"]["code"], "invalid_request"); } + #[test] + fn builds_openai_search_core_error_body() { + let body = build_core_error_body_for_client_format( + "openai:search", + "search unavailable", + Some("upstream_unavailable"), + LocalCoreSyncErrorKind::ServerError, + ) + .expect("body should build"); + + assert_eq!(body["error"]["message"], "search unavailable"); + assert_eq!(body["error"]["type"], "server_error"); + assert_eq!(body["error"]["code"], "upstream_unavailable"); + } + #[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/image_bridge.rs b/crates/aether-ai-formats/src/formats/shared/image_bridge.rs index 623c297be..89e4b11e1 100644 --- a/crates/aether-ai-formats/src/formats/shared/image_bridge.rs +++ b/crates/aether-ai-formats/src/formats/shared/image_bridge.rs @@ -14,6 +14,7 @@ pub struct OpenAiImageRequestForGemini { pub struct GeminiImageRequestForOpenAi { pub requested_model: String, pub mapped_model: String, + pub operation: crate::formats::openai::image::request::OpenAiImageOperation, pub body_json: Value, pub summary_json: Value, } @@ -136,52 +137,53 @@ pub fn build_openai_image_request_body_from_gemini_image_request( .filter(|value| !value.is_empty()) .collect::>() .join("\n\n"); - if !prompt.is_empty() { - content.insert( - 0, - json!({ - "type": "input_text", - "text": prompt, - }), - ); - } - if content.is_empty() { + if prompt.is_empty() { return None; } - let action = if content.iter().any(|value| { + let operation = if content.iter().any(|value| { value .get("type") .and_then(Value::as_str) .is_some_and(|kind| kind == "input_image") }) { - "edit" + crate::formats::openai::image::request::OpenAiImageOperation::Edit } else { - "generate" + crate::formats::openai::image::request::OpenAiImageOperation::Generate }; - let body_json = json!({ - "model": mapped_model, - "input": [{ - "role": "user", - "content": content, - }], - "tools": [{ - "type": "image_generation", - "action": action, - }], - "tool_choice": { - "type": "image_generation" - }, - "stream": false, - }); + let mut body = Map::new(); + body.insert("model".to_string(), Value::String(mapped_model.to_string())); + body.insert("prompt".to_string(), Value::String(prompt)); + if operation == crate::formats::openai::image::request::OpenAiImageOperation::Edit { + let images = content + .iter() + .filter(|value| value.get("type").and_then(Value::as_str) == Some("input_image")) + .map(|value| { + let image_url = value + .get("image_url") + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty())?; + Some(json!({ "image_url": image_url })) + }) + .collect::>>()?; + if images.is_empty() { + return None; + } + crate::formats::openai::image::request::insert_standard_openai_image_inputs( + &mut body, images, + ); + } + let body_json = Value::Object(body); let summary_json = json!({ - "operation": action, + "operation": operation.as_str(), "response_format": "b64_json", }); Some(GeminiImageRequestForOpenAi { requested_model, mapped_model: mapped_model.to_string(), + operation, body_json, summary_json, }) @@ -1003,11 +1005,50 @@ mod tests { assert_eq!(converted.requested_model, "gemini-image"); assert_eq!(converted.body_json["model"], "gpt-image-2"); - assert_eq!(converted.body_json["tools"][0]["action"], "edit"); + assert_eq!(converted.operation.as_str(), "edit"); assert_eq!( - converted.body_json["input"][0]["content"][1]["image_url"], + converted.body_json["image"]["image_url"], "data:image/png;base64,aGVsbG8=" ); + assert!(converted.body_json.get("input").is_none()); + assert!(converted.body_json.get("tools").is_none()); + assert!(converted.body_json.get("stream").is_none()); + } + + #[test] + fn converts_multiple_gemini_image_inputs_to_standard_openai_image_array() { + let body = json!({ + "generationConfig": {"responseModalities": ["TEXT", "IMAGE"]}, + "contents": [{ + "role": "user", + "parts": [ + {"text": "Combine these references"}, + {"inlineData": {"mimeType": "image/png", "data": "aGVsbG8="}}, + {"fileData": {"mimeType": "image/jpeg", "fileUri": "https://example.test/reference.jpg"}} + ] + }] + }); + + let converted = build_openai_image_request_body_from_gemini_image_request( + &body, + "/v1beta/models/gemini-image:generateContent", + "gpt-image-2", + ) + .expect("conversion should succeed"); + + assert_eq!( + converted.body_json["image"].as_array().map(Vec::len), + Some(2) + ); + assert_eq!( + converted.body_json["image"][0]["image_url"], + "data:image/png;base64,aGVsbG8=" + ); + assert_eq!( + converted.body_json["image"][1]["image_url"], + "https://example.test/reference.jpg" + ); + assert!(converted.body_json.get("images").is_none()); } #[test] diff --git a/crates/aether-ai-formats/src/formats/shared/model_directives.rs b/crates/aether-ai-formats/src/formats/shared/model_directives.rs index 4576c3b36..b97501c62 100644 --- a/crates/aether-ai-formats/src/formats/shared/model_directives.rs +++ b/crates/aether-ai-formats/src/formats/shared/model_directives.rs @@ -1,17 +1,48 @@ use serde_json::{json, Value}; +pub const MODEL_DIRECTIVE_API_FORMATS: [&str; 6] = [ + "openai:chat", + "openai:responses", + "openai:responses:compact", + "openai:search", + "claude:messages", + "gemini:generate_content", +]; +pub const OPENAI_MODEL_DIRECTIVE_SUFFIXES: [&str; 9] = [ + "none", "minimal", "low", "medium", "high", "xhigh", "max", "ultra", "fast", +]; +pub const CROSS_PROVIDER_MODEL_DIRECTIVE_SUFFIXES: [&str; 5] = + ["low", "medium", "high", "xhigh", "max"]; + #[derive(Debug, Clone, PartialEq, Eq)] pub struct ModelDirective { pub base_model: String, pub overrides: Vec, } +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ModelDirectiveSuffixResolution { + pub base_model: String, + pub suffixes: Vec, +} + #[derive(Debug, Clone, PartialEq, Eq)] pub enum ModelOverride { ReasoningEffort(ReasoningEffort), + CodexReasoningPreset(CodexReasoningPreset), ServiceTier(ServiceTier), } +impl ModelOverride { + pub fn suffix(&self) -> &'static str { + match self { + Self::ReasoningEffort(effort) => effort.as_str(), + Self::CodexReasoningPreset(preset) => preset.as_str(), + Self::ServiceTier(tier) => tier.as_directive_suffix(), + } + } +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum ReasoningEffort { None, @@ -24,6 +55,16 @@ pub enum ReasoningEffort { } impl ReasoningEffort { + pub const ALL: [Self; 7] = [ + Self::None, + Self::Minimal, + Self::Low, + Self::Medium, + Self::High, + Self::XHigh, + Self::Max, + ]; + pub fn parse(value: &str) -> Option { match value.trim().to_ascii_lowercase().as_str() { "none" => Some(Self::None), @@ -37,7 +78,7 @@ impl ReasoningEffort { } } - pub fn as_openai_chat_value(self) -> &'static str { + pub fn as_str(self) -> &'static str { match self { Self::None => "none", Self::Minimal => "minimal", @@ -49,16 +90,16 @@ impl ReasoningEffort { } } + pub fn as_openai_chat_value(self) -> &'static str { + self.as_str() + } + pub fn as_openai_responses_value(self) -> &'static str { - match self { - Self::None => "none", - Self::Minimal => "minimal", - Self::Low => "low", - Self::Medium => "medium", - Self::High => "high", - Self::XHigh => "xhigh", - Self::Max => "max", - } + self.as_str() + } + + pub fn as_openai_model_directive_value(self) -> &'static str { + self.as_str() } pub fn as_claude_output_value(self) -> &'static str { @@ -92,6 +133,19 @@ impl ReasoningEffort { } } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CodexReasoningPreset { + Ultra, +} + +impl CodexReasoningPreset { + pub fn as_str(self) -> &'static str { + match self { + Self::Ultra => "ultra", + } + } +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum ServiceTier { Priority, @@ -105,6 +159,12 @@ impl ServiceTier { } } + pub fn as_directive_suffix(self) -> &'static str { + match self { + Self::Priority => "fast", + } + } + pub fn as_openai_value(self) -> &'static str { match self { Self::Priority => "priority", @@ -113,28 +173,109 @@ impl ServiceTier { } pub fn parse_model_directive(model: &str) -> Option { - let (base_model, overrides) = parse_model_directive_parts(model)?; + let resolution = parse_model_directive_with_suffixes( + model, + OPENAI_MODEL_DIRECTIVE_SUFFIXES.iter().copied(), + )?; + let overrides = resolution + .suffixes + .iter() + .map(|suffix| parse_model_override_for_model(suffix, &resolution.base_model)) + .collect::>>()?; Some(ModelDirective { - base_model, + base_model: resolution.base_model, overrides, }) } -fn parse_model_directive_parts(model: &str) -> Option<(String, Vec)> { +pub fn parse_model_directive_with_suffixes<'a>( + model: &str, + suffixes: impl IntoIterator, +) -> Option { + let mut configured_suffixes = Vec::::new(); + for suffix in suffixes { + let suffix = suffix.trim(); + if suffix.is_empty() || suffix.starts_with('-') || suffix.ends_with('-') { + continue; + } + if let Some(existing) = configured_suffixes + .iter() + .find(|existing| existing.eq_ignore_ascii_case(suffix)) + { + if existing != suffix { + return None; + } + continue; + } + configured_suffixes.push(suffix.to_string()); + } + configured_suffixes + .sort_by(|left, right| right.len().cmp(&left.len()).then_with(|| left.cmp(right))); + let mut base_model = model.trim(); - let mut overrides = ModelOverrideAccumulator::default(); - while let Some((candidate_base, suffix)) = base_model.rsplit_once('-') { - let Some(override_item) = parse_model_override(suffix) else { - break; - }; - overrides.insert(override_item)?; + let mut matched_suffixes = Vec::::new(); + let mut matched_reasoning_effort = false; + let mut matched_service_tier = false; + while let Some((candidate_base, suffix)) = configured_suffixes.iter().find_map(|suffix| { + strip_model_directive_suffix(base_model, suffix) + .map(|candidate_base| (candidate_base, suffix)) + }) { + if matched_suffixes + .iter() + .any(|matched| matched.eq_ignore_ascii_case(suffix)) + { + return None; + } + if model_directive_suffix_is_reasoning(suffix) { + if matched_reasoning_effort { + return None; + } + matched_reasoning_effort = true; + } else if ServiceTier::parse(suffix).is_some() { + if matched_service_tier { + return None; + } + matched_service_tier = true; + } + matched_suffixes.push(suffix.clone()); base_model = candidate_base.trim(); } - if base_model.is_empty() { + + if base_model.is_empty() || matched_suffixes.is_empty() { return None; } - let overrides = overrides.into_overrides()?; - Some((base_model.to_string(), overrides)) + matched_suffixes.sort_by(|left, right| { + model_directive_suffix_rank(left) + .cmp(&model_directive_suffix_rank(right)) + .then_with(|| left.cmp(right)) + }); + Some(ModelDirectiveSuffixResolution { + base_model: base_model.to_string(), + suffixes: matched_suffixes, + }) +} + +fn strip_model_directive_suffix<'a>(model: &'a str, suffix: &str) -> Option<&'a str> { + let suffix_start = model.len().checked_sub(suffix.len())?; + let separator = suffix_start.checked_sub(1)?; + if !model.is_char_boundary(suffix_start) + || !model.is_char_boundary(separator) + || model.as_bytes().get(separator) != Some(&b'-') + || !model[suffix_start..].eq_ignore_ascii_case(suffix) + { + return None; + } + Some(&model[..separator]) +} + +fn model_directive_suffix_rank(suffix: &str) -> u8 { + if model_directive_suffix_is_reasoning(suffix) { + 0 + } else if ServiceTier::parse(suffix).is_some() { + 1 + } else { + 2 + } } fn parse_model_override(suffix: &str) -> Option { @@ -143,39 +284,35 @@ fn parse_model_override(suffix: &str) -> Option { .or_else(|| ServiceTier::parse(suffix).map(ModelOverride::ServiceTier)) } -#[derive(Default)] -struct ModelOverrideAccumulator { - reasoning_effort: Option, - service_tier: Option, +fn parse_model_override_for_model(suffix: &str, model: &str) -> Option { + if suffix.eq_ignore_ascii_case("ultra") && codex_ultra_preset_supported_for_model(model) { + return Some(ModelOverride::CodexReasoningPreset( + CodexReasoningPreset::Ultra, + )); + } + parse_model_override(suffix) } -impl ModelOverrideAccumulator { - fn insert(&mut self, override_item: ModelOverride) -> Option<()> { - match override_item { - ModelOverride::ReasoningEffort(value) => { - if self.reasoning_effort.replace(value).is_some() { - return None; - } - } - ModelOverride::ServiceTier(value) => { - if self.service_tier.replace(value).is_some() { - return None; - } - } - } - Some(()) - } +pub fn model_directive_suffix_has_builtin_mapping(suffix: &str) -> bool { + parse_model_override(suffix).is_some() || suffix.eq_ignore_ascii_case("ultra") +} - fn into_overrides(self) -> Option> { - let mut overrides = Vec::new(); - if let Some(reasoning_effort) = self.reasoning_effort { - overrides.push(ModelOverride::ReasoningEffort(reasoning_effort)); - } - if let Some(service_tier) = self.service_tier { - overrides.push(ModelOverride::ServiceTier(service_tier)); - } - (!overrides.is_empty()).then_some(overrides) - } +pub fn model_directive_builtin_suffix_supported_for_source_model( + suffix: &str, + source_model: &str, +) -> bool { + parse_model_override_for_model(suffix, source_model).is_some() +} + +fn model_directive_suffix_is_reasoning(suffix: &str) -> bool { + ReasoningEffort::parse(suffix).is_some() || suffix.eq_ignore_ascii_case("ultra") +} + +fn codex_ultra_preset_supported_for_model(model: &str) -> bool { + crate::formats::openai::responses::codex::resolve_codex_responses_model_capabilities( + model, model, None, + ) + .supports_reasoning_effort("ultra") } pub fn model_directive_base_model(model: &str) -> Option { @@ -241,9 +378,17 @@ pub fn apply_model_directive_overrides_from_model( &mut patched_body, provider_api_format, provider_model, + &directive.base_model, *effort, )?; } + ModelOverride::CodexReasoningPreset(preset) => { + apply_codex_reasoning_preset_override( + &mut patched_body, + provider_api_format, + *preset, + )?; + } ModelOverride::ServiceTier(tier) => { apply_service_tier_override(&mut patched_body, provider_api_format, *tier)?; } @@ -261,6 +406,65 @@ pub fn apply_model_directive_mapping_patch( Some(()) } +pub fn default_model_directive_mapping_patch( + provider_api_format: &str, + provider_model: &str, + source_model: &str, + suffix: &str, +) -> Option { + let override_item = parse_model_override_for_model(suffix, source_model)?; + let mut patch = json!({}); + match override_item { + ModelOverride::ReasoningEffort(effort) => apply_reasoning_effort_override( + &mut patch, + provider_api_format, + provider_model, + source_model, + effort, + )?, + ModelOverride::CodexReasoningPreset(preset) => { + apply_codex_reasoning_preset_override(&mut patch, provider_api_format, preset)? + } + ModelOverride::ServiceTier(tier) => { + apply_service_tier_override(&mut patch, provider_api_format, tier)? + } + } + Some(patch) +} + +pub fn default_model_directive_suffixes(provider_api_format: &str) -> &'static [&'static str] { + match crate::normalize_api_format_alias(provider_api_format).as_str() { + "openai:chat" | "openai:responses" | "openai:responses:compact" | "openai:search" => { + &OPENAI_MODEL_DIRECTIVE_SUFFIXES + } + "claude:messages" | "gemini:generate_content" => &CROSS_PROVIDER_MODEL_DIRECTIVE_SUFFIXES, + _ => &[], + } +} + +pub fn default_model_directives_config() -> Value { + let api_formats = MODEL_DIRECTIVE_API_FORMATS + .into_iter() + .map(|api_format| { + ( + api_format.to_string(), + json!({ + "enabled": true, + "suffixes": default_model_directive_suffixes(api_format), + "mappings": {}, + }), + ) + }) + .collect::>(); + + json!({ + "reasoning_effort": { + "enabled": true, + "api_formats": api_formats, + }, + }) +} + fn deep_merge_json(target: &mut Value, patch: &Value) { match (target, patch) { (Value::Object(target_object), Value::Object(patch_object)) => { @@ -283,15 +487,24 @@ fn apply_reasoning_effort_override( provider_request_body: &mut Value, provider_api_format: &str, provider_model: &str, + source_model: &str, effort: ReasoningEffort, ) -> Option<()> { + if !reasoning_effort_supported_for_model( + provider_api_format, + provider_model, + source_model, + effort, + ) { + return None; + } match crate::normalize_api_format_alias(provider_api_format).as_str() { "openai:chat" => set_object_string( provider_request_body, "reasoning_effort", - effort.as_openai_chat_value(), + effort.as_openai_model_directive_value(), ), - "openai:responses" | "openai:responses:compact" => { + "openai:responses" | "openai:responses:compact" | "openai:search" => { set_openai_responses_reasoning_effort(provider_request_body, effort) } "claude:messages" => { @@ -304,6 +517,29 @@ fn apply_reasoning_effort_override( } } +fn apply_codex_reasoning_preset_override( + provider_request_body: &mut Value, + provider_api_format: &str, + preset: CodexReasoningPreset, +) -> Option<()> { + match crate::normalize_api_format_alias(provider_api_format).as_str() { + "openai:chat" => { + set_object_string(provider_request_body, "reasoning_effort", preset.as_str()) + } + "openai:responses" | "openai:responses:compact" | "openai:search" => { + let object = provider_request_body.as_object_mut()?; + let reasoning = object + .entry("reasoning".to_string()) + .or_insert_with(|| json!({})); + reasoning + .as_object_mut()? + .insert("effort".to_string(), json!(preset.as_str())); + Some(()) + } + _ => None, + } +} + fn apply_service_tier_override( provider_request_body: &mut Value, provider_api_format: &str, @@ -315,6 +551,7 @@ fn apply_service_tier_override( "service_tier", tier.as_openai_value(), ), + "openai:search" => Some(()), _ => None, } } @@ -335,7 +572,7 @@ fn set_openai_responses_reasoning_effort(body: &mut Value, effort: ReasoningEffo } reasoning.as_object_mut()?.insert( "effort".to_string(), - Value::String(effort.as_openai_responses_value().to_string()), + Value::String(effort.as_openai_model_directive_value().to_string()), ); Some(()) } @@ -457,6 +694,158 @@ pub fn gemini_model_uses_thinking_level(model: &str) -> bool { .any(|part| part.starts_with("gemini-3")) } +pub fn openai_model_supports_max_reasoning_effort(model: &str) -> bool { + is_openai_gpt_5_6_family(model) +} + +pub fn openai_model_supports_prompt_cache_options(model: &str) -> bool { + let normalized = model.trim().to_ascii_lowercase().replace('_', "-"); + let model = normalized.rsplit('/').next().unwrap_or_default(); + openai_gpt_model_version(model) + .is_some_and(|(major, minor)| major > 5 || (major == 5 && minor >= 6)) +} + +pub fn reasoning_effort_supported_for_model( + provider_api_format: &str, + provider_model: &str, + source_model: &str, + effort: ReasoningEffort, +) -> bool { + match crate::normalize_api_format_alias(provider_api_format).as_str() { + "openai:chat" | "openai:responses" | "openai:responses:compact" | "openai:search" => { + match resolved_openai_model_identity(provider_model, source_model).0 { + OpenAiModelIdentity::Gpt56 => effort != ReasoningEffort::Minimal, + OpenAiModelIdentity::ConcreteOther => effort != ReasoningEffort::Max, + OpenAiModelIdentity::Opaque => true, + } + } + "claude:messages" | "gemini:generate_content" => matches!( + effort, + ReasoningEffort::Low + | ReasoningEffort::Medium + | ReasoningEffort::High + | ReasoningEffort::XHigh + | ReasoningEffort::Max + ), + _ => false, + } +} + +pub(crate) fn openai_model_resolves_to_gpt_5_6(provider_model: &str, source_model: &str) -> bool { + matches!( + resolved_openai_model_identity(provider_model, source_model).0, + OpenAiModelIdentity::Gpt56 | OpenAiModelIdentity::Opaque + ) +} + +pub(crate) fn openai_model_capability_identity(provider_model: &str, source_model: &str) -> String { + resolved_openai_model_identity(provider_model, source_model).1 +} + +pub(crate) fn openai_model_capability_is_opaque(provider_model: &str, source_model: &str) -> bool { + resolved_openai_model_identity(provider_model, source_model).0 == OpenAiModelIdentity::Opaque +} + +fn resolved_openai_model_identity( + provider_model: &str, + source_model: &str, +) -> (OpenAiModelIdentity, String) { + let provider_model = normalize_model_directive_model(provider_model); + let provider_identity = classify_openai_model_identity(&provider_model); + if provider_identity != OpenAiModelIdentity::Opaque { + return (provider_identity, provider_model); + } + let source_model = normalize_model_directive_model(source_model); + (classify_openai_model_identity(&source_model), source_model) +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum OpenAiModelIdentity { + Gpt56, + ConcreteOther, + Opaque, +} + +fn classify_openai_model_identity(model: &str) -> OpenAiModelIdentity { + if is_openai_gpt_5_6_family(model) { + return OpenAiModelIdentity::Gpt56; + } + let normalized = model.trim().to_ascii_lowercase().replace('_', "-"); + let model = normalized.rsplit('/').next().unwrap_or_default(); + if openai_gpt_model_identity_is_concrete(model) + || openai_o_series_model_identity_is_concrete(model) + || model.starts_with("chatgpt-") + || model.starts_with("codex-") + { + OpenAiModelIdentity::ConcreteOther + } else { + OpenAiModelIdentity::Opaque + } +} + +fn openai_gpt_model_identity_is_concrete(model: &str) -> bool { + let Some(rest) = model.strip_prefix("gpt-") else { + return false; + }; + let mut segments = rest.split('-'); + let Some(version) = segments.next() else { + return false; + }; + if openai_gpt_model_version(model).is_none() { + return false; + } + let version_is_omni = version.ends_with('o'); + if version.contains('.') || version_is_omni { + return true; + } + match segments.next() { + None => true, + Some(variant) => { + variant.chars().all(|character| character.is_ascii_digit()) + || matches!( + variant, + "audio" + | "chat" + | "codex" + | "mini" + | "nano" + | "pro" + | "realtime" + | "search" + | "turbo" + | "vision" + ) + } + } +} + +fn openai_gpt_model_version(model: &str) -> Option<(u64, u64)> { + let version = model.strip_prefix("gpt-")?.split('-').next()?; + let version = version.strip_suffix('o').unwrap_or(version); + let mut parts = version.split('.'); + let major = parts.next()?.parse().ok()?; + let minor = parts.next().map(str::parse).transpose().ok()?.unwrap_or(0); + if parts.next().is_some() { + return None; + } + Some((major, minor)) +} + +fn openai_o_series_model_identity_is_concrete(model: &str) -> bool { + let Some(rest) = model.strip_prefix('o') else { + return false; + }; + rest.split('-').next().is_some_and(|version| { + !version.is_empty() && version.chars().all(|character| character.is_ascii_digit()) + }) +} + +fn is_openai_gpt_5_6_family(model: &str) -> bool { + let normalized = model.trim().to_ascii_lowercase().replace('_', "-"); + let model = normalized.rsplit('/').next().unwrap_or_default(); + openai_gpt_model_version(model) == Some((5, 6)) +} + pub fn extract_gemini_model_from_path(path: &str) -> Option { let marker = "/models/"; let start = path.find(marker)? + marker.len(); @@ -471,12 +860,87 @@ mod tests { use serde_json::json; use super::{ - apply_model_directive_overrides_from_model, parse_model_directive, ModelDirective, - ModelOverride, ReasoningEffort, ServiceTier, + apply_model_directive_overrides_from_model, default_model_directive_suffixes, + default_model_directives_config, parse_model_directive, + parse_model_directive_with_suffixes, CodexReasoningPreset, ModelDirective, + ModelDirectiveSuffixResolution, ModelOverride, ReasoningEffort, ServiceTier, + MODEL_DIRECTIVE_API_FORMATS, }; + #[test] + fn policy_suffix_parser_prefers_the_longest_configured_suffix() { + assert_eq!( + parse_model_directive_with_suffixes( + "deployment-alias-VendorFuture", + ["Future", "VendorFuture"], + ), + Some(ModelDirectiveSuffixResolution { + base_model: "deployment-alias".to_string(), + suffixes: vec!["VendorFuture".to_string()], + }) + ); + assert_eq!( + parse_model_directive_with_suffixes( + "deployment-alias-high-VendorFuture", + ["high", "VendorFuture"], + ), + Some(ModelDirectiveSuffixResolution { + base_model: "deployment-alias".to_string(), + suffixes: vec!["high".to_string(), "VendorFuture".to_string()], + }) + ); + } + + #[test] + fn policy_suffix_parser_rejects_duplicate_or_ambiguous_suffixes() { + assert_eq!( + parse_model_directive_with_suffixes("deployment-low-high", ["low", "high"]), + None + ); + assert_eq!( + parse_model_directive_with_suffixes( + "deployment-VendorFuture", + ["VendorFuture", "vendorfuture"], + ), + None + ); + assert_eq!( + parse_model_directive_with_suffixes( + "deployment-VendorFuture-VendorFuture", + ["VendorFuture"], + ), + None + ); + } + + #[test] + fn policy_suffix_parser_does_not_strip_unconfigured_suffixes() { + assert_eq!( + parse_model_directive_with_suffixes("deployment-alias-VendorFuture", ["high"]), + None + ); + } + #[test] fn parses_supported_reasoning_effort_suffixes() { + let expected = [ + ("none", ReasoningEffort::None), + ("minimal", ReasoningEffort::Minimal), + ("low", ReasoningEffort::Low), + ("medium", ReasoningEffort::Medium), + ("high", ReasoningEffort::High), + ("xhigh", ReasoningEffort::XHigh), + ("max", ReasoningEffort::Max), + ]; + for (suffix, effort) in expected { + assert_eq!( + parse_model_directive(&format!("gpt-5.6-sol-{suffix}")), + Some(ModelDirective { + base_model: "gpt-5.6-sol".to_string(), + overrides: vec![ModelOverride::ReasoningEffort(effort)], + }) + ); + } assert_eq!( parse_model_directive("gpt-5.4-xhigh"), Some(ModelDirective { @@ -493,6 +957,21 @@ mod tests { ); } + #[test] + fn default_config_is_generated_from_the_shared_directive_contract() { + let config = default_model_directives_config(); + for api_format in MODEL_DIRECTIVE_API_FORMATS { + assert_eq!( + config["reasoning_effort"]["api_formats"][api_format]["mappings"], + json!({}) + ); + assert_eq!( + config["reasoning_effort"]["api_formats"][api_format]["suffixes"], + json!(default_model_directive_suffixes(api_format)) + ); + } + } + #[test] fn parses_supported_service_tier_suffixes() { assert_eq!( @@ -519,13 +998,44 @@ mod tests { #[test] fn ignores_unknown_or_incomplete_suffixes() { - assert_eq!(parse_model_directive("gpt-5.4-ultra"), None); + assert_eq!(parse_model_directive("gpt-5.4-turbo"), None); assert_eq!(parse_model_directive("gpt-5.4"), None); assert_eq!(parse_model_directive("-high"), None); assert_eq!(parse_model_directive("gpt-5.4-high-json"), None); assert_eq!(parse_model_directive("gpt-5.4-low-high"), None); } + #[test] + fn parses_gpt_5_6_ultra_as_an_internal_reasoning_preset() { + assert_eq!( + parse_model_directive("gpt-5.6-sol-ultra"), + Some(ModelDirective { + base_model: "gpt-5.6-sol".to_string(), + overrides: vec![ModelOverride::CodexReasoningPreset( + CodexReasoningPreset::Ultra, + )], + }) + ); + + for model in [ + "gpt-5.6-ultra", + "gpt-5.6-luna-ultra", + "gpt-5.4-ultra", + "gemini-ultra", + ] { + assert_eq!(parse_model_directive(model), None); + } + + let mut unsupported = json!({"model": "gpt-5.4"}); + assert!(apply_model_directive_overrides_from_model( + &mut unsupported, + "openai:responses", + "gpt-5.4", + "gpt-5.4-ultra", + ) + .is_none()); + } + #[test] fn applies_reasoning_effort_to_provider_body_shapes() { let mut openai_chat = json!({"model": "gpt-5-upstream", "reasoning_effort": "low"}); @@ -546,7 +1056,7 @@ mod tests { &mut responses, "openai:responses", "gpt-5-upstream", - "gpt-5.4-max", + "gpt-5.6-max", ) .expect("directive should apply"); assert_eq!(responses["reasoning"]["effort"], "max"); @@ -560,7 +1070,7 @@ mod tests { &mut compact, "openai:responses:compact", "gpt-5-upstream", - "gpt-5.4-max", + "gpt-5.6-max", ) .expect("directive should apply"); assert_eq!(compact["reasoning"]["effort"], "max"); @@ -570,7 +1080,7 @@ mod tests { &mut openai_chat_max, "openai:chat", "gpt-5-upstream", - "gpt-5.4-max", + "gpt-5.6-max", ) .expect("directive should apply"); assert_eq!(openai_chat_max["reasoning_effort"], "max"); @@ -599,6 +1109,142 @@ mod tests { ); } + #[test] + fn max_suffix_is_capability_aware_for_openai_models() { + for model in ["gpt-5.6", "gpt-5.6-sol", "gpt-5.6-terra", "gpt-5.6-luna"] { + let mut body = json!({"model": model}); + apply_model_directive_overrides_from_model( + &mut body, + "openai:responses", + model, + &format!("{model}-max"), + ) + .expect("max directive should apply"); + assert_eq!(body["reasoning"]["effort"], "max", "model: {model}"); + } + + let mut unsupported_model = json!({"model": "gpt-5.4"}); + let original = unsupported_model.clone(); + assert!(apply_model_directive_overrides_from_model( + &mut unsupported_model, + "openai:responses", + "gpt-5.4", + "gpt-5.4-max", + ) + .is_none()); + assert_eq!(unsupported_model, original); + + let mut mapped_deployment = json!({"model": "azure-production"}); + apply_model_directive_overrides_from_model( + &mut mapped_deployment, + "openai:responses", + "azure-production", + "gpt-5.6-sol-max", + ) + .expect("source model capability should survive provider model mapping"); + assert_eq!(mapped_deployment["reasoning"]["effort"], "max"); + + let mut explicit_unsupported_target = json!({"model": "gpt-5.4"}); + let original = explicit_unsupported_target.clone(); + assert!(apply_model_directive_overrides_from_model( + &mut explicit_unsupported_target, + "openai:responses", + "gpt-5.4", + "gpt-5.6-sol-max", + ) + .is_none()); + assert_eq!(explicit_unsupported_target, original); + + let mut unknown_future = json!({"model": "gpt-6.0"}); + let original = unknown_future.clone(); + assert!(apply_model_directive_overrides_from_model( + &mut unknown_future, + "openai:responses", + "gpt-6.0", + "gpt-6.0-max", + ) + .is_none()); + assert_eq!(unknown_future, original); + } + + #[test] + fn openai_only_efforts_do_not_leak_into_cross_provider_mappings() { + for effort in ["none", "minimal"] { + for (api_format, provider_model) in [ + ("claude:messages", "claude-sonnet-4-6"), + ("gemini:generate_content", "gemini-3-pro"), + ] { + let mut body = json!({"model": provider_model}); + let original = body.clone(); + assert!(apply_model_directive_overrides_from_model( + &mut body, + api_format, + provider_model, + &format!("gpt-5.6-sol-{effort}"), + ) + .is_none()); + assert_eq!(body, original); + } + } + } + + #[test] + fn gpt_5_6_rejects_minimal_but_accepts_published_efforts() { + for effort in ["none", "low", "medium", "high", "xhigh", "max"] { + let mut body = json!({"model": "azure-production"}); + apply_model_directive_overrides_from_model( + &mut body, + "openai:responses", + "azure-production", + &format!("gpt-5.6-sol-{effort}"), + ) + .expect("published GPT-5.6 effort should apply"); + assert_eq!(body["reasoning"]["effort"], effort); + } + + let mut body = json!({"model": "azure-production"}); + let original = body.clone(); + assert!(apply_model_directive_overrides_from_model( + &mut body, + "openai:responses", + "azure-production", + "gpt-5.6-sol-minimal", + ) + .is_none()); + assert_eq!(body, original); + + for family_variant in ["gpt-5.6-preview", "gpt-5.6-sol-2026-07-01"] { + assert!(super::openai_model_supports_max_reasoning_effort( + family_variant + )); + let mut body = json!({"model": family_variant}); + apply_model_directive_overrides_from_model( + &mut body, + "openai:responses", + family_variant, + &format!("{family_variant}-max"), + ) + .expect("GPT-5.6 family variants should accept max"); + assert_eq!(body["reasoning"]["effort"], "max"); + } + } + + #[test] + fn prompt_cache_options_capability_requires_gpt_5_6_or_later() { + for model in [ + "gpt-5.6", + "gpt-5.6-sol", + "gpt-5.7", + "gpt-6", + "openai/gpt-6.1-pro", + ] { + assert!(super::openai_model_supports_prompt_cache_options(model)); + } + for model in ["gpt-5.5", "gpt-4o", "o3", "azure-production", "gpt-5.6.1"] { + assert!(!super::openai_model_supports_prompt_cache_options(model)); + } + } + #[test] fn applies_fast_suffix_to_openai_service_tier() { let mut openai_chat = json!({"model": "gpt-5-upstream"}); diff --git a/crates/aether-ai-formats/src/formats/shared/passthrough.rs b/crates/aether-ai-formats/src/formats/shared/passthrough.rs index a0fef4ed8..4c1fd03ce 100644 --- a/crates/aether-ai-formats/src/formats/shared/passthrough.rs +++ b/crates/aether-ai-formats/src/formats/shared/passthrough.rs @@ -5,7 +5,8 @@ use crate::contracts::{ 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_RERANK_SYNC_PLAN_KIND, OPENAI_SEARCH_SYNC_PLAN_KIND, + OPENAI_SEARCH_SYNC_SUCCESS_REPORT_KIND, }; #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -81,6 +82,13 @@ pub fn resolve_sync_spec(plan_kind: &str) -> Option family: LocalSameFormatProviderFamily::Standard, require_streaming: false, }), + OPENAI_SEARCH_SYNC_PLAN_KIND => Some(LocalSameFormatProviderSpec { + api_format: "openai:search", + decision_kind: OPENAI_SEARCH_SYNC_PLAN_KIND, + report_kind: OPENAI_SEARCH_SYNC_SUCCESS_REPORT_KIND, + family: LocalSameFormatProviderFamily::Standard, + require_streaming: false, + }), _ => None, } } @@ -185,4 +193,12 @@ mod tests { assert_eq!(spec.report_kind, "openai_rerank_sync_success"); assert!(!spec.require_streaming); } + + #[test] + fn resolves_openai_search_sync_spec() { + let spec = resolve_sync_spec("openai_search_sync").expect("spec"); + assert_eq!(spec.api_format, "openai:search"); + assert_eq!(spec.report_kind, "openai_search_sync_success"); + assert!(!spec.require_streaming); + } } diff --git a/crates/aether-ai-formats/src/formats/shared/request.rs b/crates/aether-ai-formats/src/formats/shared/request.rs index 9cf4f0131..e4c81f6a8 100644 --- a/crates/aether-ai-formats/src/formats/shared/request.rs +++ b/crates/aether-ai-formats/src/formats/shared/request.rs @@ -46,6 +46,18 @@ pub fn force_upstream_streaming_for_provider( && aether_ai_formats::is_openai_responses_format(provider_api_format) } +pub fn forbid_upstream_streaming_for_provider( + provider_type: &str, + provider_api_format: &str, +) -> bool { + aether_ai_formats::api_format_alias_matches(provider_api_format, "openai:search") + || aether_ai_formats::is_openai_responses_compact_format(provider_api_format) + || (provider_type.trim().eq_ignore_ascii_case("codex") + && provider_api_format + .trim() + .eq_ignore_ascii_case("openai:image")) +} + pub(crate) fn parse_upstream_stream_policy( value: Option<&serde_json::Value>, ) -> UpstreamStreamPolicy { @@ -160,13 +172,31 @@ pub fn resolve_upstream_is_stream_from_endpoint_config( ) } +pub fn resolve_upstream_is_stream_for_provider( + endpoint_config: Option<&serde_json::Value>, + provider_type: &str, + provider_api_format: &str, + client_is_stream: bool, + hard_requires_streaming: bool, +) -> bool { + if forbid_upstream_streaming_for_provider(provider_type, provider_api_format) { + return false; + } + resolve_upstream_is_stream_from_endpoint_config( + endpoint_config, + client_is_stream, + hard_requires_streaming + || force_upstream_streaming_for_provider(provider_type, provider_api_format), + ) +} + #[cfg(test)] mod tests { use super::{ endpoint_config_forces_upstream_stream_policy, enforce_request_body_stream_field, - force_upstream_streaming_for_provider, parse_direct_request_body, - parse_upstream_stream_policy, resolve_upstream_is_stream, - resolve_upstream_is_stream_from_endpoint_config, + forbid_upstream_streaming_for_provider, force_upstream_streaming_for_provider, + parse_direct_request_body, parse_upstream_stream_policy, resolve_upstream_is_stream, + resolve_upstream_is_stream_for_provider, resolve_upstream_is_stream_from_endpoint_config, upstream_stream_policy_from_endpoint_config, UpstreamStreamPolicy, }; use serde_json::json; @@ -228,6 +258,42 @@ mod tests { )); } + #[test] + fn forbids_streaming_for_sync_only_openai_formats() { + assert!(forbid_upstream_streaming_for_provider( + "codex", + "openai:search" + )); + assert!(forbid_upstream_streaming_for_provider( + "custom", + "/v1/alpha/search" + )); + assert!(forbid_upstream_streaming_for_provider( + "codex", + "openai:responses:compact" + )); + assert!(forbid_upstream_streaming_for_provider( + "openai", + "openai:responses:compact" + )); + assert!(forbid_upstream_streaming_for_provider( + "custom", + "openai:responses:compact" + )); + assert!(forbid_upstream_streaming_for_provider( + "codex", + "openai:image" + )); + assert!(!forbid_upstream_streaming_for_provider( + "codex", + "openai:responses" + )); + assert!(!forbid_upstream_streaming_for_provider( + "openai", + "openai:image" + )); + } + #[test] fn parses_python_compatible_upstream_stream_policy_values() { assert_eq!( @@ -400,4 +466,37 @@ mod tests { None, false, false, )); } + + #[test] + fn provider_policy_gives_non_stream_contracts_precedence() { + let force_stream = json!({"upstream_stream_policy": "force_stream"}); + assert!(!resolve_upstream_is_stream_for_provider( + Some(&force_stream), + "codex", + "openai:search", + true, + true, + )); + assert!(!resolve_upstream_is_stream_for_provider( + Some(&force_stream), + "codex", + "openai:responses:compact", + true, + true, + )); + assert!(!resolve_upstream_is_stream_for_provider( + Some(&force_stream), + "codex", + "openai:image", + true, + true, + )); + assert!(resolve_upstream_is_stream_for_provider( + Some(&json!({"upstream_stream_policy": "force_non_stream"})), + "codex", + "openai:responses", + false, + false, + )); + } } diff --git a/crates/aether-ai-formats/src/formats/shared/routing.rs b/crates/aether-ai-formats/src/formats/shared/routing.rs index ac5b02874..fab37faa0 100644 --- a/crates/aether-ai-formats/src/formats/shared/routing.rs +++ b/crates/aether-ai-formats/src/formats/shared/routing.rs @@ -13,9 +13,9 @@ use crate::contracts::{ 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_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, + 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; @@ -95,14 +95,6 @@ pub fn resolve_execution_runtime_stream_plan_kind( return Some(OPENAI_RESPONSES_STREAM_PLAN_KIND); } - if route_family == Some("openai") - && is_openai_responses_compact_route_kind(route_kind) - && *method == Method::POST - && path == "/v1/responses/compact" - { - return Some(OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND); - } - if route_family == Some("openai") && route_kind == Some("image") && *method == Method::POST @@ -224,6 +216,14 @@ pub fn resolve_execution_runtime_sync_plan_kind( return Some(OPENAI_RERANK_SYNC_PLAN_KIND); } + if route_family == Some("openai") + && route_kind == Some("search") + && *method == Method::POST + && path == "/v1/alpha/search" + { + return Some(OPENAI_SEARCH_SYNC_PLAN_KIND); + } + if route_family == Some("openai") && route_kind == Some("image") && *method == Method::POST @@ -402,11 +402,13 @@ pub fn is_matching_stream_request( path: &str, body_json: &serde_json::Value, ) -> bool { + if plan_kind == OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND { + return false; + } match plan_kind { OPENAI_CHAT_STREAM_PLAN_KIND | CLAUDE_CHAT_STREAM_PLAN_KIND | OPENAI_RESPONSES_STREAM_PLAN_KIND - | OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND | CLAUDE_CLI_STREAM_PLAN_KIND | OPENAI_IMAGE_STREAM_PLAN_KIND | GEMINI_INTERACTIONS_STREAM_PLAN_KIND => body_json @@ -439,6 +441,7 @@ pub fn supports_sync_execution_decision_kind(plan_kind: &str) -> bool { OPENAI_CHAT_SYNC_PLAN_KIND | OPENAI_EMBEDDING_SYNC_PLAN_KIND | OPENAI_RERANK_SYNC_PLAN_KIND + | OPENAI_SEARCH_SYNC_PLAN_KIND | OPENAI_IMAGE_SYNC_PLAN_KIND | OPENAI_RESPONSES_SYNC_PLAN_KIND | OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND @@ -468,7 +471,6 @@ pub fn supports_stream_execution_decision_kind(plan_kind: &str) -> bool { | CLAUDE_CHAT_STREAM_PLAN_KIND | GEMINI_CHAT_STREAM_PLAN_KIND | OPENAI_RESPONSES_STREAM_PLAN_KIND - | OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND | OPENAI_IMAGE_STREAM_PLAN_KIND | CLAUDE_CLI_STREAM_PLAN_KIND | GEMINI_CLI_STREAM_PLAN_KIND @@ -499,6 +501,7 @@ mod tests { 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, }; #[test] @@ -592,7 +595,7 @@ mod tests { &Method::POST, "/v1/responses/compact", ), - Some(OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND) + None ); assert_eq!( resolve_execution_runtime_sync_plan_kind( @@ -608,11 +611,51 @@ mod tests { assert!(supports_sync_execution_decision_kind( OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND )); - assert!(supports_stream_execution_decision_kind( + assert!(!supports_stream_execution_decision_kind( OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND )); } + #[test] + fn resolves_openai_search_as_sync_only() { + assert_eq!( + resolve_execution_runtime_sync_plan_kind( + Some("ai_public"), + Some("openai"), + Some("search"), + None, + &Method::POST, + "/v1/alpha/search", + ), + Some(OPENAI_SEARCH_SYNC_PLAN_KIND) + ); + assert_eq!( + resolve_execution_runtime_stream_plan_kind( + Some("ai_public"), + Some("openai"), + Some("search"), + None, + &Method::POST, + "/v1/alpha/search", + ), + None + ); + assert!(supports_sync_execution_decision_kind( + OPENAI_SEARCH_SYNC_PLAN_KIND + )); + assert_eq!( + resolve_execution_runtime_sync_plan_kind( + Some("ai_public"), + Some("openai"), + Some("search"), + None, + &Method::POST, + "/backend-api/codex/alpha/search", + ), + None + ); + } + #[test] fn resolves_claude_messages_plan_kinds_by_request_auth_channel() { assert_eq!( diff --git a/crates/aether-ai-formats/src/formats/shared/standard_matrix.rs b/crates/aether-ai-formats/src/formats/shared/standard_matrix.rs index 65b3c4e24..5e79591cf 100644 --- a/crates/aether-ai-formats/src/formats/shared/standard_matrix.rs +++ b/crates/aether-ai-formats/src/formats/shared/standard_matrix.rs @@ -1292,6 +1292,36 @@ mod tests { } } + #[test] + fn codex_store_body_rule_preserves_response_input_item_ids() { + let request = json!({ + "model": "gpt-5.4", + "input": [{ + "id": "msg-1", + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "hello"}] + }] + }); + let body_rules = json!([{"action":"set","path":"store","value":true}]); + + let converted = build_standard_request_body( + &request, + "openai:responses", + "gpt-5.4", + "codex", + "openai:responses", + "/v1/responses", + true, + Some(&body_rules), + Some("key-1"), + ) + .expect("Codex Responses request should apply the store body rule"); + + assert_eq!(converted["store"], true); + assert_eq!(converted["input"][0]["id"], "msg-1"); + } + #[test] fn standard_openai_responses_strips_content_cache_control_after_body_rules() { let request = json!({ @@ -1329,7 +1359,7 @@ mod tests { } #[test] - fn standard_codex_responses_derives_prompt_cache_key_before_stripping_cache_control() { + fn standard_codex_responses_strip_cache_control_without_synthesizing_a_cache_key() { fn claude_request(user_text: &str) -> Value { json!({ "model": "claude-sonnet", @@ -1373,13 +1403,8 @@ mod tests { ) .expect("claude to codex responses request should build"); - assert!(converted_a["prompt_cache_key"] - .as_str() - .is_some_and(|value| !value.trim().is_empty())); - assert_eq!( - converted_a["prompt_cache_key"], - converted_b["prompt_cache_key"] - ); + assert!(converted_a.get("prompt_cache_key").is_none()); + assert!(converted_b.get("prompt_cache_key").is_none()); assert!(!converted_a.to_string().contains("cache_control")); assert!(!converted_b.to_string().contains("cache_control")); } diff --git a/crates/aether-ai-formats/src/formats/shared/standard_normalize.rs b/crates/aether-ai-formats/src/formats/shared/standard_normalize.rs index 8192d0928..470336daa 100644 --- a/crates/aether-ai-formats/src/formats/shared/standard_normalize.rs +++ b/crates/aether-ai-formats/src/formats/shared/standard_normalize.rs @@ -760,7 +760,7 @@ mod tests { #[test] fn local_openai_responses_request_body_applies_reasoning_effort_suffix() { let body_json = json!({ - "model": "gpt-5.4-max", + "model": "gpt-5.6-sol-max", "input": "hello", "reasoning": {"effort": "low", "summary": "auto"} }); diff --git a/crates/aether-ai-formats/src/formats/shared/stream_core/common.rs b/crates/aether-ai-formats/src/formats/shared/stream_core/common.rs index 33ed38989..5c08b6f16 100644 --- a/crates/aether-ai-formats/src/formats/shared/stream_core/common.rs +++ b/crates/aether-ai-formats/src/formats/shared/stream_core/common.rs @@ -152,6 +152,7 @@ pub fn canonical_usage_from_openai_usage(value: Option<&Value>) -> Option, client: Option, + propagated_actual_service_tier: Option, terminated: bool, } @@ -40,10 +41,21 @@ impl StreamingStandardFormatMatrix { self.terminated = true; return self.emit_error(error_body); } - let Some(provider) = self.provider.as_mut() else { + let (provider, client, propagated_actual_service_tier) = ( + &mut self.provider, + &mut self.client, + &mut self.propagated_actual_service_tier, + ); + let Some(provider) = provider.as_mut() else { return Ok(Vec::new()); }; let frames = provider.push_line(report_context, line)?; + if provider.actual_service_tier() != propagated_actual_service_tier.as_deref() { + *propagated_actual_service_tier = provider.actual_service_tier().map(ToOwned::to_owned); + if let Some(client) = client.as_mut() { + client.set_actual_service_tier(propagated_actual_service_tier.as_deref()); + } + } self.emit_frames(frames) } @@ -52,10 +64,21 @@ impl StreamingStandardFormatMatrix { return Ok(Vec::new()); } self.ensure_initialized(report_context); - let Some(provider) = self.provider.as_mut() else { + let (provider, client, propagated_actual_service_tier) = ( + &mut self.provider, + &mut self.client, + &mut self.propagated_actual_service_tier, + ); + let Some(provider) = provider.as_mut() else { return Ok(Vec::new()); }; let frames = provider.finish(report_context)?; + if provider.actual_service_tier() != propagated_actual_service_tier.as_deref() { + *propagated_actual_service_tier = provider.actual_service_tier().map(ToOwned::to_owned); + if let Some(client) = client.as_mut() { + client.set_actual_service_tier(propagated_actual_service_tier.as_deref()); + } + } let mut out = self.emit_frames(frames)?; if let Some(client) = self.client.as_mut() { out.extend(client.finish()?); @@ -100,6 +123,14 @@ impl StreamingStandardFormatMatrix { out.extend(client.emit_unknown_event(payload)?); break; } + if let CanonicalStreamEvent::OpenAiResponsesOutputItem { raw_event, .. } = &frame.event + { + if !matches!(client, ClientStreamEmitter::OpenAIResponses(_)) { + self.terminated = true; + out.extend(client.emit_unknown_event(raw_event)?); + break; + } + } out.extend(client.emit(frame)?); } Ok(out) @@ -132,7 +163,13 @@ impl StreamingStandardTerminalObserver { match provider { TerminalStreamParser::Standard(provider) => { let frames = provider.push_line(report_context, line)?; + let actual_service_tier = provider.actual_service_tier().map(ToOwned::to_owned); self.observe_frames(frames); + if let Some(actual_service_tier) = actual_service_tier { + self.latest_summary + .get_or_insert_with(ExecutionStreamTerminalSummary::default) + .provider_actual_service_tier = Some(actual_service_tier); + } } TerminalStreamParser::OpenAIImage(provider) => { if let Some(summary) = provider.push_line(report_context, line)? { @@ -272,6 +309,7 @@ impl ProviderStreamParser { FormatId::ClaudeMessages => Self::Claude(ClaudeProviderState::default()), FormatId::GeminiGenerateContent => Self::Gemini(GeminiProviderState::default()), FormatId::OpenAiEmbedding + | FormatId::OpenAiSearch | FormatId::OpenAiRerank | FormatId::GeminiEmbedding | FormatId::GeminiInteractions @@ -306,11 +344,19 @@ impl ProviderStreamParser { ProviderStreamParser::Gemini(state) => state.finish(report_context), } } + + fn actual_service_tier(&self) -> Option<&str> { + match self { + ProviderStreamParser::OpenAIChat(state) => state.actual_service_tier(), + ProviderStreamParser::OpenAIResponses(state) => state.actual_service_tier(), + ProviderStreamParser::Claude(_) | ProviderStreamParser::Gemini(_) => None, + } + } } enum ClientStreamEmitter { OpenAIChat(OpenAIChatClientEmitter), - OpenAIResponses(OpenAIResponsesClientEmitter), + OpenAIResponses(Box), Claude(ClaudeClientEmitter), Gemini(GeminiClientEmitter), } @@ -359,11 +405,12 @@ impl ClientStreamEmitter { Some(match FormatId::parse(client_api_format)? { FormatId::OpenAiChat => Self::OpenAIChat(OpenAIChatClientEmitter::default()), FormatId::OpenAiResponses | FormatId::OpenAiResponsesCompact => { - Self::OpenAIResponses(OpenAIResponsesClientEmitter::default()) + Self::OpenAIResponses(Box::default()) } FormatId::ClaudeMessages => Self::Claude(ClaudeClientEmitter::default()), FormatId::GeminiGenerateContent => Self::Gemini(GeminiClientEmitter::default()), FormatId::OpenAiEmbedding + | FormatId::OpenAiSearch | FormatId::OpenAiRerank | FormatId::GeminiEmbedding | FormatId::GeminiInteractions @@ -383,6 +430,14 @@ impl ClientStreamEmitter { } } + fn set_actual_service_tier(&mut self, value: Option<&str>) { + match self { + ClientStreamEmitter::OpenAIChat(state) => state.set_actual_service_tier(value), + ClientStreamEmitter::OpenAIResponses(state) => state.set_actual_service_tier(value), + ClientStreamEmitter::Claude(_) | ClientStreamEmitter::Gemini(_) => {} + } + } + fn finish(&mut self) -> Result, AiSurfaceFinalizeError> { match self { ClientStreamEmitter::OpenAIChat(state) => state.finish(), @@ -478,6 +533,7 @@ fn parse_provider_error( parse_gemini_error(payload) } FormatId::OpenAiEmbedding + | FormatId::OpenAiSearch | FormatId::OpenAiRerank | FormatId::GeminiEmbedding | FormatId::JinaEmbedding @@ -577,6 +633,15 @@ mod tests { format!("data: {}\n", value).into_bytes() } + fn json_data_events(bytes: &[u8]) -> Vec { + String::from_utf8_lossy(bytes) + .lines() + .filter_map(|line| line.strip_prefix("data: ")) + .filter(|payload| *payload != "[DONE]") + .filter_map(|payload| serde_json::from_str(payload).ok()) + .collect() + } + #[test] fn transforms_provider_errors_to_openai_chat_error_bodies() { let cases = [ @@ -1093,6 +1158,66 @@ mod tests { } } + #[test] + fn responses_compaction_output_is_lossless_within_family_and_rejected_cross_format() { + let compaction_event = json!({ + "type": "response.output_item.done", + "item": { + "type": "compaction", + "encrypted_content": "ENCRYPTED_CONTEXT_COMPACTION_SUMMARY" + } + }); + + let mut responses_matrix = StreamingStandardFormatMatrix::default(); + let responses_context = report_context("openai:responses", "openai:responses"); + let mut responses_output = responses_matrix + .transform_line(&responses_context, data_line(compaction_event.clone())) + .expect("same-family compaction output should convert"); + responses_output.extend( + responses_matrix + .transform_line( + &responses_context, + data_line(json!({ + "type": "response.completed", + "response": { + "id": "resp-compact", + "model": "gpt-5.6-sol", + "usage": { + "input_tokens": 0, + "output_tokens": 0, + "total_tokens": 0 + } + } + })), + ) + .expect("terminal response should convert"), + ); + let responses_sse = String::from_utf8(responses_output).expect("valid Responses SSE"); + assert!(responses_sse.contains("event: response.output_item.done\n")); + assert!(responses_sse.contains("\"type\":\"compaction\"")); + assert!(!responses_sse.contains("\"output_index\"")); + assert!(responses_sse.contains("event: response.completed\n")); + + for client_api_format in ["openai:chat", "claude:messages", "gemini:generate_content"] { + let mut matrix = StreamingStandardFormatMatrix::default(); + let context = report_context("openai:responses", client_api_format); + let output = matrix + .transform_line(&context, data_line(compaction_event.clone())) + .expect("cross-format rejection should be encoded for the client"); + let sse = String::from_utf8(output).expect("valid error SSE"); + assert!( + sse.contains("Unsupported provider stream event cannot be converted losslessly") + && sse.contains("compaction"), + "{client_api_format}: {sse}" + ); + if client_api_format == "gemini:generate_content" { + assert!(sse.contains("\"status\":\"INTERNAL\""), "{sse}"); + } else { + assert!(sse.contains("unsupported_stream_event"), "{sse}"); + } + } + } + #[test] fn transforms_openai_responses_known_sidecar_events_without_unsupported_errors() { let report_context = report_context("openai:responses", "claude:messages"); @@ -1795,6 +1920,147 @@ mod tests { assert_eq!(summary.unknown_event_count, 0); } + #[test] + fn terminal_observer_preserves_actual_service_tier_without_response_capture() { + let chat_context = report_context("openai:chat", "openai:chat"); + let mut chat_observer = StreamingStandardTerminalObserver::default(); + chat_observer + .push_line( + &chat_context, + data_line(json!({ + "id": "chatcmpl_tier_1", + "object": "chat.completion.chunk", + "model": "gpt-5.6", + "service_tier": "Default", + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12}, + })), + ) + .expect("Chat terminal tier should be observed"); + assert_eq!( + chat_observer + .latest_summary() + .and_then(|summary| summary.provider_actual_service_tier.as_deref()), + Some("default") + ); + let chat_summary = chat_observer + .latest_summary() + .expect("Chat summary should exist"); + assert!(chat_summary.observed_finish); + assert_eq!( + chat_summary + .standardized_usage + .as_ref() + .map(|usage| (usage.input_tokens, usage.output_tokens)), + Some((10, 2)) + ); + + let responses_context = report_context("openai:responses", "openai:responses"); + let mut responses_observer = StreamingStandardTerminalObserver::default(); + responses_observer + .push_line( + &responses_context, + data_line(json!({ + "type": "response.completed", + "response": { + "id": "resp_tier_1", + "model": "gpt-5.6", + "status": "completed", + "service_tier": "Flex", + "output": [], + "usage": {"input_tokens": 10, "output_tokens": 2, "total_tokens": 12}, + }, + "sequence_number": 1, + })), + ) + .expect("Responses terminal tier should be observed"); + assert_eq!( + responses_observer + .latest_summary() + .and_then(|summary| summary.provider_actual_service_tier.as_deref()), + Some("flex") + ); + let responses_summary = responses_observer + .latest_summary() + .expect("Responses summary should exist"); + assert!(responses_summary.observed_finish); + assert_eq!( + responses_summary + .standardized_usage + .as_ref() + .map(|usage| (usage.input_tokens, usage.output_tokens)), + Some((10, 2)) + ); + } + + #[test] + fn openai_chat_client_chunks_carry_provider_actual_service_tier() { + let context = report_context("openai:chat", "openai:chat"); + let mut matrix = StreamingStandardFormatMatrix::default(); + let output = matrix + .transform_line( + &context, + data_line(json!({ + "id": "chatcmpl_tier_stream", + "object": "chat.completion.chunk", + "model": "gpt-5.6", + "service_tier": "Default", + "choices": [{ + "index": 0, + "delta": {"role": "assistant", "content": "done"}, + "finish_reason": "stop" + }], + "usage": {"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12} + })), + ) + .expect("Chat stream should transform"); + let events = json_data_events(&output); + + assert!(!events.is_empty()); + assert!(events + .iter() + .all(|event| event.get("service_tier") == Some(&json!("default")))); + } + + #[test] + fn responses_actual_service_tier_reaches_transformed_chat_chunks() { + let context = report_context("openai:responses", "openai:chat"); + let mut matrix = StreamingStandardFormatMatrix::default(); + let output = matrix + .transform_line( + &context, + data_line(json!({ + "type": "response.completed", + "response": { + "id": "resp_tier_stream", + "object": "response", + "model": "gpt-5.6", + "status": "completed", + "service_tier": "Flex", + "output": [{ + "type": "message", + "id": "msg_tier_stream", + "role": "assistant", + "status": "completed", + "content": [{ + "type": "output_text", + "text": "done", + "annotations": [] + }] + }], + "usage": {"input_tokens": 10, "output_tokens": 2, "total_tokens": 12} + } + })), + ) + .expect("Responses stream should transform"); + let events = json_data_events(&output); + + assert!(!events.is_empty()); + assert!(events + .iter() + .all(|event| event.get("service_tier") == Some(&json!("flex")))); + } + #[test] fn terminal_observer_tracks_openai_image_stream_usage() { let mut report_context = report_context("openai:image", "openai:chat"); 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 0285bf28e..55349c4b8 100644 --- a/crates/aether-ai-formats/src/formats/shared/sync_products.rs +++ b/crates/aether-ai-formats/src/formats/shared/sync_products.rs @@ -699,6 +699,13 @@ fn maybe_build_openai_responses_same_family_stream_sync_body( return Ok(None); }; let body_bytes = base64::engine::general_purpose::STANDARD.decode(body_base64)?; + if let Some(terminal_body) = terminal_openai_responses_stream_response(&body_bytes) { + return Ok(Some(client_body_with_report_context_model( + terminal_body, + report_context, + &client_api_format, + ))); + } Ok( try_aggregate_openai_responses_stream_sync_response(&body_bytes)?.map(|body| { client_body_with_report_context_model(body, report_context, &client_api_format) @@ -706,6 +713,32 @@ fn maybe_build_openai_responses_same_family_stream_sync_body( ) } +fn terminal_openai_responses_stream_response(body: &[u8]) -> Option { + parse_stream_json_events(body)? + .into_iter() + .rev() + .find_map(|event| { + let event = event.as_object()?; + let event_type = event.get("type").and_then(Value::as_str)?; + if !matches!( + event_type, + "response.completed" | "response.done" | "response.incomplete" | "response.failed" + ) { + return None; + } + let response = event.get("response").and_then(Value::as_object).cloned()?; + if matches!(event_type, "response.completed" | "response.done") + && response + .get("output") + .and_then(Value::as_array) + .is_none_or(|output| output.is_empty()) + { + return None; + } + Some(Value::Object(response)) + }) +} + fn maybe_build_openai_cross_format_provider_body_from_normalized_payload( body_json: Option<&Value>, body_base64: Option<&str>, @@ -781,10 +814,20 @@ pub fn maybe_build_standard_cross_format_sync_product( let client_api_format = client_api_format.trim().to_ascii_lowercase(); if provider_api_format == "openai:image" && client_api_format == "gemini:generate_content" { - let client_body_json = crate::formats::shared::image_bridge::build_gemini_image_response_from_openai_responses_image_response( - &provider_body_json, - Some(report_context), - )?; + let client_body_json = match ( + provider_body_json.get("data").and_then(Value::as_array), + provider_body_json.get("output").and_then(Value::as_array), + ) { + (Some(_), None) => crate::formats::shared::image_bridge::build_gemini_image_response_from_openai_image_response( + &provider_body_json, + Some(report_context), + )?, + (None, Some(_)) => crate::formats::shared::image_bridge::build_gemini_image_response_from_openai_responses_image_response( + &provider_body_json, + Some(report_context), + )?, + _ => return None, + }; return Some(StandardCrossFormatSyncProduct { client_body_json, provider_body_json, @@ -1875,7 +1918,10 @@ pub fn aggregate_openai_responses_stream_sync_response(body: &[u8]) -> Option = BTreeMap::new(); let mut tool_states: BTreeMap = BTreeMap::new(); let mut image_items: BTreeMap = BTreeMap::new(); + let mut opaque_items: BTreeMap = BTreeMap::new(); let mut item_output_indexes = BTreeMap::::new(); + let mut generic_output_indexes = BTreeMap::::new(); + let mut next_output_index = 0_usize; for event in events { let event_object = event.as_object()?; @@ -2057,8 +2103,20 @@ pub fn aggregate_openai_responses_stream_sync_response(body: &[u8]) -> Option merge_openai_responses_message_item( message_states.entry(output_index).or_default(), @@ -2082,6 +2140,11 @@ pub fn aggregate_openai_responses_stream_sync_response(body: &[u8]) -> Option { image_items.insert(output_index, Value::Object(item.clone())); } + _ if event_object.get("type").and_then(Value::as_str) + == Some("response.output_item.done") => + { + opaque_items.insert(output_index, Value::Object(item.clone())); + } _ => {} } } @@ -2239,7 +2302,9 @@ pub fn aggregate_openai_responses_stream_sync_response(body: &[u8]) -> Option { image_items.insert(output_index, Value::Object(item.clone())); } - _ => {} + _ => { + opaque_items.insert(output_index, Value::Object(item.clone())); + } } } } @@ -2282,6 +2347,7 @@ pub fn aggregate_openai_responses_stream_sync_response(body: &[u8]) -> Option>(); output_indexes.sort_unstable(); @@ -2308,6 +2374,9 @@ pub fn aggregate_openai_responses_stream_sync_response(body: &[u8]) -> Option) -> String { + let item_type = item.get("type").and_then(Value::as_str).unwrap_or_default(); + if let Some(item_id) = item.get("id").and_then(Value::as_str) { + return format!("{item_type}:id:{item_id}"); + } + if let Some(encrypted_content) = item.get("encrypted_content").and_then(Value::as_str) { + return format!("{item_type}:encrypted_content:{encrypted_content}"); + } + format!( + "{item_type}:{}", + serde_json::to_string(item).unwrap_or_default() + ) +} + fn openai_responses_event_output_index(event: &Map) -> Option { event .get("output_index") @@ -3158,6 +3241,9 @@ fn try_aggregate_gemini_stream_sync_response( parts.push(gemini_sync_part_from_canonical_content_part(part)); } } + CanonicalStreamEvent::OpenAiResponsesOutputItem { raw_event, .. } => { + return Err(unsupported_stream_event_finalize_error(&raw_event)) + } CanonicalStreamEvent::ToolCallStart { index, call_id, @@ -3554,6 +3640,87 @@ mod tests { use base64::Engine as _; use serde_json::json; + #[test] + fn converts_openai_images_sync_body_to_gemini_image_body() { + let provider_body_json = json!({ + "created": 1776839946, + "model": "gpt-image-2", + "data": [{ + "revised_prompt": "revised prompt", + "b64_json": "aGVsbG8=" + }], + "usage": { + "input_tokens": 3, + "output_tokens": 4, + "total_tokens": 7 + } + }); + let report_context = json!({ + "provider_api_format": "openai:image", + "client_api_format": "gemini:generate_content", + }); + + let product = maybe_build_standard_cross_format_sync_product( + "gemini_chat_sync_finalize", + "openai:image", + "gemini:generate_content", + &report_context, + provider_body_json, + ) + .expect("OpenAI Images sync body should convert"); + + assert_eq!(product.client_body_json["modelVersion"], "gpt-image-2"); + assert_eq!( + product.client_body_json["candidates"][0]["content"]["parts"][0]["text"], + "revised prompt" + ); + assert_eq!( + product.client_body_json["candidates"][0]["content"]["parts"][1]["inlineData"]["data"], + "aGVsbG8=" + ); + assert_eq!( + product.client_body_json["usageMetadata"]["totalTokenCount"], + 7 + ); + } + + #[test] + fn converts_openai_responses_image_body_to_gemini_image_body() { + let provider_body_json = json!({ + "model": "gpt-image-2", + "output": [{ + "type": "image_generation_call", + "status": "completed", + "result": "aGVsbG8=", + "output_format": "png" + }] + }); + let report_context = json!({ + "provider_api_format": "openai:image", + "client_api_format": "gemini:generate_content", + }); + + let product = maybe_build_standard_cross_format_sync_product( + "gemini_chat_sync_finalize", + "openai:image", + "gemini:generate_content", + &report_context, + provider_body_json, + ) + .expect("OpenAI Responses image body should convert"); + + assert_eq!(product.client_body_json["modelVersion"], "gpt-image-2"); + assert_eq!( + product.client_body_json["candidates"][0]["content"]["parts"][0]["inlineData"] + ["mimeType"], + "image/png" + ); + assert_eq!( + product.client_body_json["candidates"][0]["content"]["parts"][0]["inlineData"]["data"], + "aGVsbG8=" + ); + } + #[test] fn aggregates_openai_chat_stream_tool_usage_and_finish_into_sync_body() { let body = concat!( @@ -4203,7 +4370,7 @@ mod tests { "event: response.output_text.delta\n", "data: {\"type\":\"response.output_text.delta\",\"output_index\":0,\"content_index\":0,\"delta\":\"Hello\"}\n\n", "event: response.completed\n", - "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_123\",\"object\":\"response\",\"model\":\"gpt-5\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":1,\"output_tokens\":2,\"total_tokens\":3}}}\n\n", + "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_123\",\"object\":\"response\",\"model\":\"gpt-5\",\"status\":\"completed\",\"output\":[{\"type\":\"message\",\"role\":\"assistant\",\"status\":\"completed\",\"content\":[{\"type\":\"output_text\",\"text\":\"Hello\"}]}],\"usage\":{\"input_tokens\":1,\"output_tokens\":2,\"total_tokens\":3}}}\n\n", ); let report_context = json!({ "provider_api_format": "openai:responses:compact", @@ -4224,9 +4391,241 @@ mod tests { assert_eq!(body_json.get("id"), Some(&json!("resp_123"))); assert_eq!(body_json.get("status"), Some(&json!("completed"))); assert_eq!(body_json["output"][0]["content"][0]["text"], json!("Hello")); - assert_eq!(body_json["output_text"], "Hello"); - assert!(body_json["created_at"].as_i64().is_some()); - assert!(body_json["completed_at"].as_i64().is_some()); + } + + #[test] + fn same_family_completed_metadata_uses_materialized_stream_output() { + for terminal_response in [ + json!({ + "id": "resp_codex_terminal_metadata", + "object": "response", + "model": "gpt-5.6-sol", + "status": "completed", + "output": [], + "usage": {"input_tokens": 1, "output_tokens": 2, "total_tokens": 3} + }), + json!({ + "id": "resp_codex_terminal_metadata", + "object": "response", + "model": "gpt-5.6-sol", + "status": "completed", + "usage": {"input_tokens": 1, "output_tokens": 2, "total_tokens": 3} + }), + ] { + let body = format!( + "event: response.created\ndata: {}\n\n\ + event: response.output_text.delta\ndata: {}\n\n\ + event: response.completed\ndata: {}\n\n", + json!({ + "type": "response.created", + "response": { + "id": "resp_codex_terminal_metadata", + "object": "response", + "model": "gpt-5.6-sol", + "status": "in_progress", + "output": [] + } + }), + json!({ + "type": "response.output_text.delta", + "output_index": 0, + "content_index": 0, + "delta": "Hello from Codex" + }), + json!({"type": "response.completed", "response": terminal_response}) + ); + let report_context = json!({ + "provider_api_format": "openai:responses", + "client_api_format": "openai:responses", + "needs_conversion": false, + }); + + let body_json = + maybe_build_openai_responses_same_family_sync_body_from_normalized_payload( + "openai_responses_sync_finalize", + 200, + Some(&report_context), + None, + Some(&base64::engine::general_purpose::STANDARD.encode(body)), + ) + .expect("Responses stream should aggregate") + .expect("sync response should exist"); + + assert_eq!( + body_json["output"][0]["content"][0]["text"], + "Hello from Codex" + ); + } + } + + #[test] + fn same_family_responses_stream_uses_terminal_response_as_authoritative_snapshot() { + let terminal_response = json!({ + "id": "resp_gpt56_terminal", + "object": "response", + "model": "gpt-5.6-sol", + "status": "completed", + "output": [ + { + "type": "program", + "call_id": "program-call-1", + "fingerprint": "fp-program-1" + }, + { + "type": "program_output", + "call_id": "program-call-1", + "result": "hello", + "status": "completed" + }, + { + "type": "multi_agent_call", + "call_id": "agent-call-1", + "action": "delegate", + "arguments": {"query": "release notes"}, + "agent": "researcher" + }, + { + "type": "agent_message", + "author": "researcher", + "recipient": "assistant", + "encrypted_content": "encrypted-agent-message" + } + ], + "usage": { + "input_tokens": 12, + "output_tokens": 4, + "total_tokens": 16, + "input_tokens_details": {"cache_write_tokens": 3} + } + }); + let body = format!( + "event: response.program.delta\ndata: {{\"type\":\"response.program.delta\",\"delta\":\"print\"}}\n\n\ + event: response.multi_agent_call.in_progress\ndata: {{\"type\":\"response.multi_agent_call.in_progress\",\"item_id\":\"agent-call-1\"}}\n\n\ + event: response.completed\ndata: {}\n\n", + json!({"type": "response.completed", "response": terminal_response}) + ); + let report_context = json!({ + "provider_api_format": "openai:responses", + "client_api_format": "openai:responses", + "needs_conversion": false, + }); + + let body_json = maybe_build_openai_responses_same_family_sync_body_from_normalized_payload( + "openai_responses_sync_finalize", + 200, + Some(&report_context), + None, + Some(&base64::engine::general_purpose::STANDARD.encode(body)), + ) + .expect("same-family terminal snapshot should bypass unknown-event aggregation") + .expect("terminal response should become the sync response"); + + assert_eq!(body_json, terminal_response); + } + + #[test] + fn cross_format_responses_stream_still_rejects_unknown_gpt_5_6_events() { + let body = concat!( + "event: response.program.delta\n", + "data: {\"type\":\"response.program.delta\",\"delta\":\"print\"}\n\n", + "event: response.completed\n", + "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_cross_unknown\",\"object\":\"response\",\"model\":\"gpt-5.6-sol\",\"status\":\"completed\",\"output\":[{\"type\":\"program\",\"id\":\"program-1\"}]}}\n\n", + ); + let report_context = json!({ + "provider_api_format": "openai:responses", + "client_api_format": "openai:chat", + "needs_conversion": true, + }); + + let result = maybe_build_openai_chat_cross_format_sync_product_from_normalized_payload( + "openai_chat_sync_finalize", + 200, + Some(&report_context), + None, + Some(&base64::engine::general_purpose::STANDARD.encode(body)), + ); + + assert!(result.is_err()); + } + + #[test] + fn same_family_responses_stream_accepts_authoritative_incomplete_snapshot() { + let terminal_response = json!({ + "id": "resp_gpt56_incomplete", + "object": "response", + "model": "gpt-5.6-sol", + "status": "incomplete", + "output": [{ + "type": "agent_message", + "author": "researcher", + "recipient": "assistant", + "encrypted_content": "encrypted-partial-message" + }], + "incomplete_details": {"reason": "max_output_tokens"} + }); + let body = format!( + "event: response.agent_message.delta\ndata: {{\"type\":\"response.agent_message.delta\",\"delta\":\"partial\"}}\n\n\ + event: response.incomplete\ndata: {}\n\n", + json!({"type": "response.incomplete", "response": terminal_response}) + ); + let report_context = json!({ + "provider_api_format": "openai:responses", + "client_api_format": "openai:responses", + "needs_conversion": false, + }); + + let body_json = maybe_build_openai_responses_same_family_sync_body_from_normalized_payload( + "openai_responses_sync_finalize", + 200, + Some(&report_context), + None, + Some(&base64::engine::general_purpose::STANDARD.encode(body)), + ) + .expect("incomplete terminal snapshot should remain a valid Responses result") + .expect("terminal response should become the sync response"); + + assert_eq!(body_json, terminal_response); + } + + #[test] + fn same_family_responses_stream_accepts_authoritative_failed_snapshot() { + let terminal_response = json!({ + "id": "resp_gpt56_failed", + "object": "response", + "model": "gpt-5.6-sol", + "status": "failed", + "output": [{ + "type": "program", + "call_id": "program-call-1", + "fingerprint": "fp-program-1" + }], + "error": {"code": "server_error", "message": "upstream failed"} + }); + let body = format!( + "event: response.failed\ndata: {}\n\n", + json!({ + "type": "response.failed", + "sequence_number": 3, + "response": terminal_response + }) + ); + let report_context = json!({ + "provider_api_format": "openai:responses", + "client_api_format": "openai:responses", + "needs_conversion": false, + }); + + let body_json = maybe_build_openai_responses_same_family_sync_body_from_normalized_payload( + "openai_responses_sync_finalize", + 200, + Some(&report_context), + None, + Some(&base64::engine::general_purpose::STANDARD.encode(body)), + ) + .expect("failed terminal snapshot should remain an authoritative Responses result") + .expect("failed terminal response should become the sync response"); + + assert_eq!(body_json, terminal_response); } #[test] @@ -4237,7 +4636,7 @@ mod tests { "event: response.outtext.delta\n", "data: {\"type\":\"response.outtext.delta\",\"output_index\":0,\"content_index\":0,\"delta\":\"Hello from legacy alias\"}\n\n", "event: response.completed\n", - "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_legacy_123\",\"object\":\"response\",\"model\":\"gpt-5\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":1,\"output_tokens\":4,\"total_tokens\":5}}}\n\n", + "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_legacy_123\",\"object\":\"response\",\"model\":\"gpt-5\",\"status\":\"completed\",\"output\":[{\"type\":\"message\",\"role\":\"assistant\",\"status\":\"completed\",\"content\":[{\"type\":\"output_text\",\"text\":\"Hello from legacy alias\"}]}],\"usage\":{\"input_tokens\":1,\"output_tokens\":4,\"total_tokens\":5}}}\n\n", ); let report_context = json!({ "provider_api_format": "openai:responses", @@ -4296,6 +4695,26 @@ mod tests { assert_eq!(result["output"][0]["output_format"], "png"); } + #[test] + fn reconstructs_openai_responses_compaction_without_output_index_or_terminal_output() { + let body = concat!( + "event: response.output_item.done\n", + "data: {\"type\":\"response.output_item.done\",\"item\":{\"type\":\"compaction\",\"encrypted_content\":\"ENCRYPTED_CONTEXT_COMPACTION_SUMMARY\"}}\n\n", + "event: response.completed\n", + "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp-compact\",\"usage\":{\"input_tokens\":0,\"output_tokens\":0,\"total_tokens\":0}}}\n\n", + ); + + let result = aggregate_openai_responses_stream_sync_response(body.as_bytes()) + .expect("compaction stream should aggregate into a sync response"); + + assert_eq!(result["output"].as_array().map(Vec::len), Some(1)); + assert_eq!(result["output"][0]["type"], "compaction"); + assert_eq!( + result["output"][0]["encrypted_content"], + "ENCRYPTED_CONTEXT_COMPACTION_SUMMARY" + ); + } + #[test] fn reconstructs_openai_responses_multi_part_message_content_order() { let body = concat!( @@ -4947,8 +5366,19 @@ mod tests { .expect("canonical openai responses -> openai chat"); assert_eq!(converted_openai_chat, legacy_openai_chat); + let mut representable_body_json = provider_body_json.clone(); + representable_body_json + .as_object_mut() + .expect("response object") + .remove("service_tier"); + representable_body_json["output"][1]["content"] = json!([{ + "type": "output_text", + "text": "Hello", + "annotations": [] + }]); + let converted_claude = convert_standard_chat_response( - &provider_body_json, + &representable_body_json, "openai:responses", "claude:messages", &report_context, @@ -4965,7 +5395,7 @@ mod tests { assert_eq!(converted_claude["usage"]["output_tokens"], 5); let converted_gemini = convert_standard_chat_response( - &provider_body_json, + &representable_body_json, "openai:responses", "gemini:generate_content", &report_context, 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 3a2eef03f..1533a4a60 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 @@ -18,7 +18,8 @@ use crate::formats::shared::stream_core::common::{ build_openai_chat_usage_chunk_with_cache, }; use crate::formats::shared::stream_core::{ - CanonicalStreamFrame, StreamingStandardFormatMatrix, StreamingStandardTerminalObserver, + CanonicalStreamEvent, CanonicalStreamFrame, StreamingStandardFormatMatrix, + StreamingStandardTerminalObserver, }; use crate::formats::shared::stream_rewrite::maybe_build_ai_surface_stream_rewriter; use crate::formats::shared::AiSurfaceFinalizeError; @@ -70,12 +71,21 @@ pub fn maybe_bridge_standard_sync_json_to_stream( { return Ok(None); } - let bridge_context = build_bridge_report_context( report_context, provider_api_format.as_str(), client_api_format.as_str(), ); + if is_openai_responses_family_api_format(provider_api_format.as_str()) + && is_openai_responses_family_api_format(client_api_format.as_str()) + { + return bridge_openai_responses_same_family_sync_json_to_stream( + provider_body_json, + provider_api_format.as_str(), + &bridge_context, + ); + } + let Some(openai_responses_response) = convert_provider_sync_response_to_openai_responses( provider_body_json, provider_api_format.as_str(), @@ -83,14 +93,33 @@ pub fn maybe_bridge_standard_sync_json_to_stream( ) else { return Ok(None); }; - let terminal_summary = - build_terminal_summary_from_openai_responses_response(&openai_responses_response); + let provider_actual_service_tier = provider_actual_service_tier_from_sync_response( + provider_body_json, + provider_api_format.as_str(), + ); + let terminal_summary = build_terminal_summary_from_openai_responses_response( + &openai_responses_response, + provider_actual_service_tier.clone(), + ); let canonical_frames = build_canonical_frames_from_openai_responses_response( &openai_responses_response, &bridge_context, )?; - let sse_body = - emit_client_stream_from_canonical_frames(canonical_frames, client_api_format.as_str())?; + let sse_body = if is_openai_responses_family_api_format(client_api_format.as_str()) { + emit_openai_responses_stream_with_authoritative_terminal( + canonical_frames, + &openai_responses_response, + openai_responses_terminal_event_type(&openai_responses_response) + .unwrap_or("response.completed"), + provider_actual_service_tier.as_deref(), + )? + } else { + emit_client_stream_from_canonical_frames( + canonical_frames, + client_api_format.as_str(), + provider_actual_service_tier.as_deref(), + )? + }; Ok(Some(SyncToStreamBridgeOutcome { sse_body, @@ -98,6 +127,50 @@ pub fn maybe_bridge_standard_sync_json_to_stream( })) } +fn bridge_openai_responses_same_family_sync_json_to_stream( + response: &Value, + provider_api_format: &str, + report_context: &Value, +) -> Result, AiSurfaceFinalizeError> { + let bridge_response = if openai_responses_terminal_event_type(response).is_some() { + Cow::Borrowed(response) + } else if provider_api_format == "openai:responses:compact" + && response.get("output").is_some_and(Value::is_array) + { + let mut response = response.clone(); + let object = response + .as_object_mut() + .expect("Compact response with output should be an object"); + object.insert("status".to_string(), Value::String("completed".to_string())); + object + .entry("object".to_string()) + .or_insert_with(|| Value::String("response.compaction".to_string())); + Cow::Owned(response) + } else { + return Ok(None); + }; + let canonical_frames = build_canonical_frames_from_openai_responses_response( + bridge_response.as_ref(), + report_context, + )?; + let terminal_event_type = openai_responses_terminal_event_type(bridge_response.as_ref()) + .expect("bridge response should have a terminal status"); + let sse_body = emit_openai_responses_stream_with_authoritative_terminal( + canonical_frames, + response, + terminal_event_type, + provider_actual_service_tier_from_sync_response(response, provider_api_format).as_deref(), + )?; + + Ok(Some(SyncToStreamBridgeOutcome { + sse_body, + terminal_summary: build_terminal_summary_from_openai_responses_response( + response, + provider_actual_service_tier_from_sync_response(response, provider_api_format), + ), + })) +} + fn maybe_bridge_openai_image_sync_json_to_stream( provider_body_json: &Value, report_context: Option<&Value>, @@ -462,6 +535,7 @@ fn openai_image_terminal_summary( .and_then(Value::as_str) .map(ToOwned::to_owned) .or_else(|| image_bridge_model(report_context)), + provider_actual_service_tier: None, observed_finish: true, unknown_event_count: 0, parser_error: None, @@ -616,6 +690,10 @@ fn is_standard_api_format(value: &str) -> bool { ) } +fn is_openai_responses_family_api_format(value: &str) -> bool { + matches!(value, "openai:responses" | "openai:responses:compact") +} + fn maybe_bridge_aether_sse_response_capture_to_stream( provider_body_json: &Value, provider_api_format: &str, @@ -922,10 +1000,12 @@ fn build_canonical_frames_from_openai_responses_response( report_context: &Value, ) -> Result, AiSurfaceFinalizeError> { let mut state = OpenAIResponsesProviderState::default(); + let event_type = openai_responses_terminal_event_type(openai_responses_response) + .unwrap_or("response.completed"); let line = format!( "data: {}\n", serde_json::to_string(&json!({ - "type": "response.completed", + "type": event_type, "response": openai_responses_response, })) .map_err(|err| AiSurfaceFinalizeError::new(err.to_string()))? @@ -941,17 +1021,29 @@ fn build_canonical_frames_from_openai_responses_response( Ok(frames) } +fn openai_responses_terminal_event_type(response: &Value) -> Option<&'static str> { + match response.get("status").and_then(Value::as_str) { + Some("completed") => Some("response.completed"), + Some("incomplete") => Some("response.incomplete"), + Some("failed") => Some("response.failed"), + _ => None, + } +} + fn emit_client_stream_from_canonical_frames( canonical_frames: Vec, client_api_format: &str, + provider_actual_service_tier: Option<&str>, ) -> Result, AiSurfaceFinalizeError> { match client_api_format { "openai:chat" => { let mut emitter = OpenAIChatClientEmitter::default(); + emitter.set_actual_service_tier(provider_actual_service_tier); emit_with_openai_chat_emitter(&mut emitter, canonical_frames) } "openai:responses" | "openai:responses:compact" => { let mut emitter = OpenAIResponsesClientEmitter::default(); + emitter.set_actual_service_tier(provider_actual_service_tier); emit_with_openai_responses_emitter(&mut emitter, canonical_frames) } "claude:messages" => { @@ -1006,6 +1098,39 @@ fn emit_with_openai_responses_emitter( Ok(output) } +fn emit_openai_responses_stream_with_authoritative_terminal( + canonical_frames: Vec, + authoritative_response: &Value, + terminal_event_type: &'static str, + provider_actual_service_tier: Option<&str>, +) -> Result, AiSurfaceFinalizeError> { + let mut emitter = OpenAIResponsesClientEmitter::default(); + emitter.set_actual_service_tier(provider_actual_service_tier); + let mut output = Vec::new(); + for frame in canonical_frames { + if matches!( + &frame.event, + CanonicalStreamEvent::Finish { .. } | CanonicalStreamEvent::UnknownEvent(_) + ) { + continue; + } + output.extend( + emitter + .emit(frame) + .map_err(|err| AiSurfaceFinalizeError::new(err.to_string()))?, + ); + } + output.extend( + emitter + .finish_with_authoritative_response_event( + authoritative_response.clone(), + terminal_event_type, + ) + .map_err(|err| AiSurfaceFinalizeError::new(err.to_string()))?, + ); + Ok(output) +} + fn emit_with_claude_emitter( emitter: &mut ClaudeClientEmitter, canonical_frames: Vec, @@ -1048,6 +1173,7 @@ fn emit_with_gemini_emitter( fn build_terminal_summary_from_openai_responses_response( openai_responses_response: &Value, + provider_actual_service_tier: Option, ) -> Option { let response = openai_responses_response.as_object()?; let response_id = response @@ -1071,12 +1197,34 @@ fn build_terminal_summary_from_openai_responses_response( finish_reason, response_id, model, + provider_actual_service_tier, observed_finish: true, unknown_event_count: 0, parser_error: None, }) } +fn provider_actual_service_tier_from_sync_response( + provider_response: &Value, + provider_api_format: &str, +) -> Option { + if !matches!( + normalize_api_format(provider_api_format).as_str(), + "openai:chat" | "openai:responses" | "openai:responses:compact" + ) { + return None; + } + provider_response + .get("response") + .and_then(Value::as_object) + .and_then(|response| response.get("service_tier")) + .or_else(|| provider_response.get("service_tier")) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty() && value.len() <= 64) + .map(str::to_ascii_lowercase) +} + fn resolve_openai_responses_finish_reason(output: &[Value]) -> String { let has_tool_calls = output.iter().filter_map(Value::as_object).any(|item| { item.get("type") @@ -1114,6 +1262,7 @@ fn standardized_usage_from_openai_usage(value: &Value) -> Option Option Vec { + body.lines() + .filter_map(|line| line.strip_prefix("data: ")) + .filter(|line| *line != "[DONE]") + .map(|line| serde_json::from_str(line).expect("valid SSE JSON")) + .collect() + } + + fn assert_strictly_increasing_sequence_numbers(events: &[Value]) { + let sequence_numbers = events + .iter() + .filter_map(|event| event.get("sequence_number").and_then(Value::as_u64)) + .collect::>(); + assert_eq!(sequence_numbers.len(), events.len()); + assert!( + sequence_numbers.windows(2).all(|pair| pair[0] < pair[1]), + "sequence numbers must be strictly increasing: {sequence_numbers:?}" + ); + } + #[test] fn openai_sync_usage_derives_missing_input_tokens_from_total() { let usage = standardized_usage_from_openai_usage(&json!({ @@ -1176,6 +1345,246 @@ mod tests { assert_eq!(usage.cache_read_tokens, 19_840); } + #[test] + fn bridges_same_family_responses_sync_with_standard_lifecycle_and_authoritative_terminal() { + let response = json!({ + "id": "resp_gpt56_raw_bridge", + "object": "response", + "model": "gpt-5.6-sol", + "status": "completed", + "output": [ + { + "type": "program", + "call_id": "program-call-1", + "fingerprint": "fp-program-1" + }, + { + "type": "program_output", + "call_id": "program-call-1", + "result": "hello", + "status": "completed" + }, + { + "type": "multi_agent_call", + "call_id": "agent-call-1", + "action": "delegate", + "arguments": {"query": "release notes"}, + "agent": "researcher" + }, + { + "type": "agent_message", + "author": "researcher", + "recipient": "assistant", + "encrypted_content": "encrypted-agent-message" + } + ], + "usage": { + "input_tokens": 20, + "input_tokens_details": { + "cached_tokens": 4, + "cache_write_tokens": 6 + }, + "output_tokens": 5, + "total_tokens": 25 + } + }); + + let outcome = maybe_bridge_standard_sync_json_to_stream( + &response, + "openai:responses", + "openai:responses:compact", + None, + ) + .expect("same-family raw bridge should succeed") + .expect("same-family raw bridge should emit terminal SSE"); + + let output = utf8(outcome.sse_body); + assert!(output.contains("event: response.created")); + assert!(output.contains("event: response.in_progress")); + let events = json_sse_events(&output); + assert_strictly_increasing_sequence_numbers(&events); + let event = events + .iter() + .find(|event| event["type"] == "response.completed") + .expect("terminal event data should exist"); + assert_eq!(events.last(), Some(event)); + assert_eq!(event["type"], "response.completed"); + assert_eq!(event["response"], response); + } + + #[test] + fn cross_format_bridge_keeps_provider_actual_service_tier_authoritative() { + let provider_response = json!({ + "id": "chatcmpl_actual_tier", + "object": "chat.completion", + "model": "gpt-5.6-sol", + "service_tier": "Default", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "done"}, + "finish_reason": "stop" + }], + "usage": {"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12} + }); + let report_context = json!({ + "original_request_body": { + "model": "gpt-5.6-sol", + "service_tier": "priority" + } + }); + + let outcome = maybe_bridge_standard_sync_json_to_stream( + &provider_response, + "openai:chat", + "openai:responses", + Some(&report_context), + ) + .expect("cross-format bridge should succeed") + .expect("cross-format bridge should emit terminal SSE"); + + assert_eq!( + outcome + .terminal_summary + .as_ref() + .and_then(|summary| summary.provider_actual_service_tier.as_deref()), + Some("default") + ); + let events = json_sse_events(&utf8(outcome.sse_body)); + let terminal = events + .iter() + .find(|event| event["type"] == "response.completed") + .expect("terminal event should exist"); + assert_eq!(terminal["response"]["service_tier"], "Default"); + let response_events = events + .iter() + .filter(|event| { + matches!( + event.get("type").and_then(Value::as_str), + Some("response.created" | "response.in_progress" | "response.completed") + ) + }) + .collect::>(); + assert!(!response_events.is_empty()); + assert!(response_events.iter().all(|event| { + event + .pointer("/response/service_tier") + .and_then(Value::as_str) + .is_some_and(|tier| tier.eq_ignore_ascii_case("default")) + })); + } + + #[test] + fn sync_bridge_does_not_echo_requested_service_tier_as_provider_actual() { + let provider_response = json!({ + "id": "chatcmpl_without_actual_tier", + "object": "chat.completion", + "model": "gpt-5.6-sol", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "done"}, + "finish_reason": "stop" + }], + "usage": {"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12} + }); + let report_context = json!({ + "original_request_body": { + "model": "gpt-5.6-sol", + "service_tier": "priority" + } + }); + + let outcome = maybe_bridge_standard_sync_json_to_stream( + &provider_response, + "openai:chat", + "openai:responses", + Some(&report_context), + ) + .expect("cross-format bridge should succeed") + .expect("cross-format bridge should emit terminal SSE"); + + assert_eq!( + outcome + .terminal_summary + .as_ref() + .and_then(|summary| summary.provider_actual_service_tier.as_deref()), + None + ); + let events = json_sse_events(&utf8(outcome.sse_body)); + assert!(events.iter().all(|event| { + event + .get("response") + .and_then(|response| response.get("service_tier")) + .is_none() + })); + } + + #[test] + fn bridges_output_only_compact_json_without_reshaping_the_terminal_response() { + let response = json!({ + "output": [{ + "type": "compaction", + "id": "cmp_123", + "encrypted_content": "encrypted-compact-history" + }] + }); + + let outcome = maybe_bridge_standard_sync_json_to_stream( + &response, + "openai:responses:compact", + "openai:responses:compact", + None, + ) + .expect("Compact bridge should succeed") + .expect("Compact output should emit terminal SSE"); + + let output = utf8(outcome.sse_body); + assert!(output.contains("event: response.created")); + assert!(output.contains("event: response.in_progress")); + let events = json_sse_events(&output); + let terminal = events + .iter() + .find(|event| event["type"] == "response.completed") + .expect("Compact terminal event should exist"); + assert_eq!(terminal["response"], response); + assert_eq!(events.last(), Some(terminal)); + } + + #[test] + fn bridges_failed_same_family_responses_sync_with_authoritative_terminal() { + let response = json!({ + "id": "resp_gpt56_failed", + "object": "response", + "model": "gpt-5.6-sol", + "status": "failed", + "error": {"code": "server_error", "message": "upstream failed"}, + "output": [{ + "type": "program", + "call_id": "program-call-1", + "fingerprint": "fp-program-1" + }] + }); + + let outcome = maybe_bridge_standard_sync_json_to_stream( + &response, + "openai:responses", + "openai:responses", + None, + ) + .expect("failed same-family raw bridge should succeed") + .expect("failed response should emit terminal SSE"); + + let output = utf8(outcome.sse_body); + let events = json_sse_events(&output); + assert_strictly_increasing_sequence_numbers(&events); + let event = events + .iter() + .find(|event| event["type"] == "response.failed") + .expect("terminal event data should exist"); + assert_eq!(events.last(), Some(event)); + assert_eq!(event["type"], "response.failed"); + assert_eq!(event["response"], response); + } + #[test] fn bridges_openai_image_sync_json_to_generation_completed_sse() { let report_context = json!({ @@ -1284,6 +1693,7 @@ mod tests { assert!(output.contains("\"finish_reason\":\"stop\"")); assert!(output.contains("\"cached_tokens\":20")); assert!(output.contains("\"cache_write_tokens\":10")); + assert!(!output.contains("\"cached_creation_tokens\"")); assert!(output.contains("data: [DONE]")); assert!(!output.contains("image_generation.completed")); diff --git a/crates/aether-ai-formats/src/lib.rs b/crates/aether-ai-formats/src/lib.rs index f49a6aee3..b5ce22cf9 100644 --- a/crates/aether-ai-formats/src/lib.rs +++ b/crates/aether-ai-formats/src/lib.rs @@ -11,7 +11,10 @@ pub use formats::context::{ FormatError, }; pub use formats::id::{ - api_format_alias_matches, api_format_storage_aliases, api_format_uses_body_stream_field, + api_format_alias_matches, api_format_defaults_to_client_error_failover, + api_format_defaults_to_non_stream, api_format_permission_covers, + api_format_permission_storage_aliases, api_format_storage_aliases, + api_format_uses_body_stream_field, intersect_api_format_allowed_lists, is_openai_responses_compact_format, is_openai_responses_family_format, is_openai_responses_format, normalize_api_format_alias, FormatFamily, FormatId, FormatProfile, }; @@ -22,6 +25,30 @@ pub use formats::matrix::{ sync_chat_response_conversion_kind, sync_cli_response_conversion_kind, RequestConversionKind, SyncChatResponseConversionKind, SyncCliResponseConversionKind, }; +pub use formats::openai::prompt_cache::resolve_openai_prompt_cache_ttl_minutes; +pub use formats::openai::prompt_cache::{ + validate_openai_prompt_cache_request, OpenAiPromptCacheContractViolation, + OpenAiPromptCacheViolationKind, +}; +pub use formats::openai::reasoning::{ + validate_openai_reasoning_request, OpenAiReasoningContractViolation, + OpenAiReasoningViolationKind, +}; +pub use formats::openai::request_contract::{ + finalize_openai_provider_request, + finalize_openai_provider_request_with_codex_model_capabilities, + validate_openai_provider_request_contract, OpenAiProviderRequestContractViolation, + OpenAiProviderRequestFinalization, +}; +pub use formats::openai::responses::codex::{ + build_codex_model_catalog_metadata, bundled_codex_model_cards, effective_codex_model_cards, + parse_codex_auth_identity, resolve_codex_responses_model_capabilities, CodexAuthIdentity, + CodexResponsesModelCapabilities, CODEX_CLIENT_ORIGINATOR, CODEX_CLIENT_USER_AGENT, + CODEX_CLIENT_VERSION, CODEX_MODEL_CATALOG_METADATA_FIELD, CODEX_RESPONSES_LITE_HEADER, +}; +pub use formats::openai::responses::request::{ + validate_openai_responses_request_contract, OpenAiResponsesRequestContractViolation, +}; pub use formats::registry::{ build_stream_transcoder, convert_request, convert_request_pure, convert_request_pure_with_context, convert_response, convert_response_pure, emit_request_pure, @@ -30,12 +57,21 @@ pub use formats::registry::{ pub use formats::shared::model_directives::{ apply_model_directive_mapping_patch, apply_model_directive_overrides_from_model, apply_model_directive_overrides_from_request, claude_model_uses_adaptive_effort, - extract_gemini_model_from_path, gemini_model_uses_thinking_level, model_directive_base_model, - normalize_model_directive_model, parse_model_directive, ModelDirective, ModelOverride, - ReasoningEffort, ServiceTier, + default_model_directive_mapping_patch, default_model_directive_suffixes, + default_model_directives_config, extract_gemini_model_from_path, + gemini_model_uses_thinking_level, model_directive_base_model, + model_directive_builtin_suffix_supported_for_source_model, + model_directive_suffix_has_builtin_mapping, normalize_model_directive_model, + openai_model_supports_prompt_cache_options, parse_model_directive, + parse_model_directive_with_suffixes, reasoning_effort_supported_for_model, ModelDirective, + ModelDirectiveSuffixResolution, ModelOverride, ReasoningEffort, ServiceTier, + CROSS_PROVIDER_MODEL_DIRECTIVE_SUFFIXES, MODEL_DIRECTIVE_API_FORMATS, + OPENAI_MODEL_DIRECTIVE_SUFFIXES, }; pub use formats::shared::request::{ endpoint_config_forces_upstream_stream_policy, enforce_request_body_stream_field, + forbid_upstream_streaming_for_provider, force_upstream_streaming_for_provider, + parse_direct_request_body, resolve_upstream_is_stream_for_provider, resolve_upstream_is_stream_from_endpoint_config, UPSTREAM_IS_STREAM_KEY, }; pub use protocol::canonical::{ diff --git a/crates/aether-ai-formats/src/protocol/canonical.rs b/crates/aether-ai-formats/src/protocol/canonical.rs index 1eaa561bb..f71fade9d 100644 --- a/crates/aether-ai-formats/src/protocol/canonical.rs +++ b/crates/aether-ai-formats/src/protocol/canonical.rs @@ -24,6 +24,8 @@ const OPENAI_CHAT_TOOL_RESULT_SOURCE_MARKER: &str = "openai_chat_tool_result"; const OPENAI_RESPONSES_TOOL_RESULT_SOURCE_MARKER: &str = "openai_responses_tool_result"; const OPENAI_RESPONSES_INPUT_MESSAGE_SOURCE_MARKER: &str = "openai_responses_input_message"; const OPENAI_RESPONSES_RAW_SOURCE_MARKER: &str = "openai_responses_raw"; +const OPENAI_RESPONSES_RAW_CONTENT_SOURCE_MARKER: &str = "openai_responses_raw_content"; +const OPENAI_RESPONSES_CONTENT_MARKER: &str = "openai_responses_content"; const OPENAI_CHAT_TOOL_ERROR_PREFIX: &str = "[tool error]"; #[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] @@ -785,27 +787,45 @@ pub(crate) fn canonical_tool_use_to_openai_responses_item( } else if !name.trim().is_empty() { item.insert("name".to_string(), Value::String(name.to_string())); } + let extension_fields = openai_responses_item_extension_object(extensions, &item); + item.extend(extension_fields); return Value::Object(item); } if is_openai_custom_tool_call(extensions) { let item_id = openai_responses_tool_call_item_id(id, extensions); - return json!({ - "type": "custom_tool_call", - "id": item_id, - "call_id": id, - "status": "completed", - "name": name, - "input": openai_custom_tool_input_text(input), - }); + let mut item = Map::new(); + item.insert( + "type".to_string(), + Value::String("custom_tool_call".to_string()), + ); + item.insert("id".to_string(), Value::String(item_id)); + item.insert("call_id".to_string(), Value::String(id.to_string())); + item.insert("status".to_string(), Value::String("completed".to_string())); + item.insert("name".to_string(), Value::String(name.to_string())); + item.insert( + "input".to_string(), + Value::String(openai_custom_tool_input_text(input)), + ); + let extension_fields = openai_responses_item_extension_object(extensions, &item); + item.extend(extension_fields); + return Value::Object(item); } let item_id = openai_responses_tool_call_item_id(id, extensions); - json!({ - "type": "function_call", - "id": item_id, - "call_id": id, - "name": name, - "arguments": canonicalize_tool_arguments(input), - }) + let mut item = Map::new(); + item.insert( + "type".to_string(), + Value::String("function_call".to_string()), + ); + item.insert("id".to_string(), Value::String(item_id)); + item.insert("call_id".to_string(), Value::String(id.to_string())); + item.insert("name".to_string(), Value::String(name.to_string())); + item.insert( + "arguments".to_string(), + Value::String(canonicalize_tool_arguments(input)), + ); + let extension_fields = openai_responses_item_extension_object(extensions, &item); + item.extend(extension_fields); + Value::Object(item) } pub(crate) fn canonical_tool_use_to_openai_responses_input_item( @@ -831,6 +851,8 @@ pub(crate) fn canonical_tool_use_to_openai_responses_input_item( } else if !name.trim().is_empty() { item.insert("name".to_string(), Value::String(name.to_string())); } + let extension_fields = openai_responses_item_extension_object(extensions, &item); + item.extend(extension_fields); return Value::Object(item); } if is_openai_custom_tool_call(extensions) { @@ -849,6 +871,8 @@ pub(crate) fn canonical_tool_use_to_openai_responses_input_item( "input".to_string(), Value::String(openai_custom_tool_input_text(input)), ); + let extension_fields = openai_responses_item_extension_object(extensions, &item); + item.extend(extension_fields); return Value::Object(item); } let mut item = Map::new(); @@ -865,6 +889,8 @@ pub(crate) fn canonical_tool_use_to_openai_responses_input_item( "arguments".to_string(), Value::String(canonicalize_tool_arguments(input)), ); + let extension_fields = openai_responses_item_extension_object(extensions, &item); + item.extend(extension_fields); Value::Object(item) } @@ -1940,6 +1966,10 @@ pub(crate) fn openai_responses_input_to_canonical_messages( continue; } let Some(item_object) = item.as_object() else { + messages.push(openai_responses_opaque_input_item_message( + item, + String::new(), + )); pending_reasoning = None; continue; }; @@ -1960,36 +1990,16 @@ pub(crate) fn openai_responses_input_to_canonical_messages( .and_then(Value::as_str) .unwrap_or("user"), ); - if matches!(role, CanonicalRole::System | CanonicalRole::Developer) { - let text = openai_content_text(item_object.get("content")); - if !text.trim().is_empty() { - let mut extensions = openai_responses_extensions( - item_object, - &["type", "role", "content"], - ); - mark_openai_responses_input_message(&mut extensions); - messages.push(CanonicalMessage { - role, - content: vec![CanonicalContentBlock::Text { - text, - extensions: BTreeMap::new(), - }], - extensions, - }); - } - pending_reasoning = None; - continue; - } let is_assistant = role == CanonicalRole::Assistant; + let mut extensions = + openai_responses_extensions(item_object, &["type", "role", "content"]); + mark_openai_responses_input_message(&mut extensions); let mut message = CanonicalMessage { role, - content: openai_responses_chat_safe_content_to_blocks( + content: openai_responses_content_to_blocks( item_object.get("content"), )?, - extensions: openai_responses_extensions( - item_object, - &["type", "role", "content"], - ), + extensions, }; if is_assistant { if let Some(reasoning) = pending_reasoning.take() { @@ -2154,6 +2164,7 @@ pub(crate) fn openai_responses_input_to_canonical_messages( pending_reasoning = None; } _ => { + messages.push(openai_responses_opaque_input_item_message(item, item_type)); pending_reasoning = None; } } @@ -2167,6 +2178,18 @@ pub(crate) fn openai_responses_input_to_canonical_messages( } } +fn openai_responses_opaque_input_item_message(item: &Value, raw_type: String) -> CanonicalMessage { + CanonicalMessage { + role: CanonicalRole::Unknown, + content: vec![CanonicalContentBlock::Unknown { + raw_type, + payload: item.clone(), + extensions: openai_responses_raw_extensions(BTreeMap::new()), + }], + extensions: BTreeMap::new(), + } +} + fn append_openai_responses_tool_use( messages: &mut Vec, tool_use: CanonicalContentBlock, @@ -2174,7 +2197,10 @@ fn append_openai_responses_tool_use( ) { let reasoning = pending_reasoning.take(); if let Some(last_message) = messages.last_mut() { - if last_message.role == CanonicalRole::Assistant { + if last_message.role == CanonicalRole::Assistant + && (!is_openai_responses_input_message(&last_message.extensions) + || canonical_assistant_message_has_visible_content(last_message)) + { if let Some(reasoning) = reasoning { prepend_openai_responses_reasoning_block(last_message, reasoning); } @@ -2195,6 +2221,20 @@ fn append_openai_responses_tool_use( }); } +fn canonical_assistant_message_has_visible_content(message: &CanonicalMessage) -> bool { + message.content.iter().any(|block| match block { + CanonicalContentBlock::Text { text, .. } | CanonicalContentBlock::Thinking { text, .. } => { + !text.trim().is_empty() + } + CanonicalContentBlock::Unknown { payload, .. } => !payload.is_null(), + CanonicalContentBlock::Image { .. } + | CanonicalContentBlock::File { .. } + | CanonicalContentBlock::Audio { .. } + | CanonicalContentBlock::ToolUse { .. } + | CanonicalContentBlock::ToolResult { .. } => true, + }) +} + fn prepend_openai_responses_reasoning_block( message: &mut CanonicalMessage, reasoning: CanonicalContentBlock, @@ -2280,17 +2320,6 @@ fn openai_responses_reasoning_text_part(raw: &Value) -> Option { (!text.is_empty()).then(|| text.to_string()) } -fn openai_responses_chat_safe_content_to_blocks( - content: Option<&Value>, -) -> Option> { - Some( - openai_responses_content_to_blocks(content)? - .into_iter() - .filter(|block| !matches!(block, CanonicalContentBlock::Unknown { .. })) - .collect(), - ) -} - pub(crate) fn openai_responses_content_to_blocks( content: Option<&Value>, ) -> Option> { @@ -2311,14 +2340,15 @@ pub(crate) fn openai_responses_content_to_blocks( } } -pub(crate) fn openai_responses_output_to_canonical_blocks( +pub(crate) fn openai_responses_output_to_canonical( output: Option<&Value>, -) -> Option> { +) -> Option<(Vec, BTreeMap)> { let Some(output) = output else { - return Some(Vec::new()); + return Some((Vec::new(), BTreeMap::new())); }; let output_items = output.as_array()?; let mut blocks = Vec::new(); + let mut message_item_provenance = Vec::new(); for (index, item) in output_items.iter().enumerate() { let Some(item_object) = item.as_object() else { blocks.push(CanonicalContentBlock::Unknown { @@ -2336,6 +2366,19 @@ pub(crate) fn openai_responses_output_to_canonical_blocks( .to_ascii_lowercase(); match item_type.as_str() { "message" => { + let message_extensions = openai_responses_extensions( + item_object, + &["type", "id", "status", "role", "content"], + ) + .remove(OPENAI_RESPONSES_EXTENSION_NAMESPACE) + .and_then(|value| value.as_object().cloned()) + .unwrap_or_default(); + if !message_extensions.is_empty() { + message_item_provenance.push(json!({ + "output_index": index, + "fields": message_extensions, + })); + } blocks.extend(openai_responses_content_to_blocks( item_object.get("content"), )?); @@ -2415,7 +2458,7 @@ pub(crate) fn openai_responses_output_to_canonical_blocks( .unwrap_or_else(|| format!("call_auto_{index}")); let mut extensions = openai_responses_extensions( item_object, - &["type", "id", "call_id", "name", "arguments", "status"], + &["type", "id", "call_id", "name", "arguments"], ); remember_openai_responses_tool_call_item_id(&mut extensions, item_object); blocks.push(CanonicalContentBlock::ToolUse { @@ -2562,7 +2605,14 @@ pub(crate) fn openai_responses_output_to_canonical_blocks( }), } } - Some(blocks) + let mut extensions = BTreeMap::new(); + if !message_item_provenance.is_empty() { + extensions.insert( + OPENAI_RESPONSES_EXTENSION_NAMESPACE.to_string(), + json!({ "message_items": message_item_provenance }), + ); + } + Some((blocks, extensions)) } fn openai_responses_hosted_tool_call_to_block( @@ -2790,7 +2840,11 @@ pub(crate) fn openai_responses_part_to_canonical_block( .and_then(Value::as_str) .unwrap_or_default() .to_string(), - extensions: openai_responses_extensions(part_object, &["type", "text"]), + extensions: { + let mut extensions = openai_responses_extensions(part_object, &["type", "text"]); + mark_openai_responses_content_block(&mut extensions); + extensions + }, }), "reasoning" | "thinking" => Some(CanonicalContentBlock::Thinking { text: part_object @@ -2940,12 +2994,12 @@ pub(crate) fn openai_responses_part_to_canonical_block( "refusal" => Some(CanonicalContentBlock::Unknown { raw_type, payload: part.clone(), - extensions: BTreeMap::new(), + extensions: openai_responses_raw_content_extensions(BTreeMap::new()), }), _ => Some(CanonicalContentBlock::Unknown { raw_type, payload: part.clone(), - extensions: BTreeMap::new(), + extensions: openai_responses_raw_content_extensions(BTreeMap::new()), }), } } @@ -3309,6 +3363,23 @@ fn openai_responses_raw_extensions( extensions } +fn openai_responses_raw_content_extensions( + mut extensions: BTreeMap, +) -> BTreeMap { + canonical_extension_object_mut(&mut extensions, AETHER_EXTENSION_NAMESPACE).insert( + "source".to_string(), + Value::String(OPENAI_RESPONSES_RAW_CONTENT_SOURCE_MARKER.to_string()), + ); + extensions +} + +fn mark_openai_responses_content_block(extensions: &mut BTreeMap) { + canonical_extension_object_mut(extensions, AETHER_EXTENSION_NAMESPACE).insert( + OPENAI_RESPONSES_CONTENT_MARKER.to_string(), + Value::Bool(true), + ); +} + fn openai_thinking_extensions(mut extensions: BTreeMap) -> BTreeMap { canonical_extension_object_mut(&mut extensions, AETHER_EXTENSION_NAMESPACE).insert( "source".to_string(), @@ -3377,6 +3448,22 @@ pub(crate) fn is_openai_responses_raw_block(extensions: &BTreeMap == Some(OPENAI_RESPONSES_RAW_SOURCE_MARKER) } +pub(crate) fn is_openai_responses_raw_content_block(extensions: &BTreeMap) -> bool { + extensions + .get(AETHER_EXTENSION_NAMESPACE) + .and_then(|value| value.get("source")) + .and_then(Value::as_str) + == Some(OPENAI_RESPONSES_RAW_CONTENT_SOURCE_MARKER) +} + +pub(crate) fn is_openai_responses_content_block(extensions: &BTreeMap) -> bool { + extensions + .get(AETHER_EXTENSION_NAMESPACE) + .and_then(|value| value.get(OPENAI_RESPONSES_CONTENT_MARKER)) + .and_then(Value::as_bool) + == Some(true) +} + pub(crate) fn is_openai_responses_input_message(extensions: &BTreeMap) -> bool { extensions .get(AETHER_EXTENSION_NAMESPACE) @@ -3620,16 +3707,19 @@ pub(crate) fn canonical_content_block_to_openai_part( block: &CanonicalContentBlock, ) -> Option { match block { - CanonicalContentBlock::Text { text, .. } => Some(json!({ - "type": "text", - "text": text, - })), + CanonicalContentBlock::Text { text, extensions } => { + let mut part = Map::new(); + part.insert("type".to_string(), Value::String("text".to_string())); + part.insert("text".to_string(), Value::String(text.clone())); + insert_openai_prompt_cache_breakpoint(&mut part, extensions); + Some(Value::Object(part)) + } CanonicalContentBlock::Image { data, url, media_type, detail, - .. + extensions, } => { let mut image = Map::new(); image.insert( @@ -3639,10 +3729,11 @@ pub(crate) fn canonical_content_block_to_openai_part( if let Some(detail) = detail { image.insert("detail".to_string(), Value::String(detail.clone())); } - Some(json!({ - "type": "image_url", - "image_url": Value::Object(image), - })) + let mut part = Map::new(); + part.insert("type".to_string(), Value::String("image_url".to_string())); + part.insert("image_url".to_string(), Value::Object(image)); + insert_openai_prompt_cache_breakpoint(&mut part, extensions); + Some(Value::Object(part)) } CanonicalContentBlock::File { data, @@ -3650,7 +3741,7 @@ pub(crate) fn canonical_content_block_to_openai_part( file_url, media_type, filename, - .. + extensions, } => { let mut file = Map::new(); if let Some(value) = file_id { @@ -3669,18 +3760,30 @@ pub(crate) fn canonical_content_block_to_openai_part( if let Some(value) = filename { file.insert("filename".to_string(), Value::String(value.clone())); } - Some(json!({ - "type": "file", - "file": Value::Object(file), - })) + let mut part = Map::new(); + part.insert("type".to_string(), Value::String("file".to_string())); + part.insert("file".to_string(), Value::Object(file)); + insert_openai_prompt_cache_breakpoint(&mut part, extensions); + Some(Value::Object(part)) + } + CanonicalContentBlock::Audio { + data, + format, + extensions, + .. + } => { + let mut part = Map::new(); + part.insert("type".to_string(), Value::String("input_audio".to_string())); + part.insert( + "input_audio".to_string(), + json!({ + "data": data.clone().unwrap_or_default(), + "format": format.clone().unwrap_or_else(|| "mp3".to_string()), + }), + ); + insert_openai_prompt_cache_breakpoint(&mut part, extensions); + Some(Value::Object(part)) } - CanonicalContentBlock::Audio { data, format, .. } => Some(json!({ - "type": "input_audio", - "input_audio": { - "data": data.clone().unwrap_or_default(), - "format": format.clone().unwrap_or_else(|| "mp3".to_string()), - } - })), CanonicalContentBlock::Thinking { text, .. } => Some(json!({ "type": "text", "text": text, @@ -3691,6 +3794,32 @@ pub(crate) fn canonical_content_block_to_openai_part( } } +pub(crate) fn openai_prompt_cache_breakpoint_from_extensions( + extensions: &BTreeMap, +) -> Option { + [ + "openai", + OPENAI_RESPONSES_EXTENSION_NAMESPACE, + OPENAI_RESPONSES_LEGACY_EXTENSION_NAMESPACE, + ] + .into_iter() + .find_map(|namespace| { + extensions + .get(namespace) + .and_then(|value| value.get("prompt_cache_breakpoint")) + .cloned() + }) +} + +fn insert_openai_prompt_cache_breakpoint( + part: &mut Map, + extensions: &BTreeMap, +) { + if let Some(value) = openai_prompt_cache_breakpoint_from_extensions(extensions) { + part.insert("prompt_cache_breakpoint".to_string(), value); + } +} + pub(crate) fn canonical_content_block_to_openai_responses_part( block: &CanonicalContentBlock, ) -> Option { @@ -5819,6 +5948,40 @@ pub(crate) fn openai_usage_to_canonical(value: Option<&Value>) -> Option) -> Option { + let usage = value?.as_object()?; + let mut canonical = openai_usage_to_canonical(value)?; + let provider_fields = usage + .iter() + .filter(|(key, _)| { + !matches!( + key.as_str(), + "input_tokens" | "output_tokens" | "total_tokens" + ) + }) + .map(|(key, value)| { + let value = if key == "input_tokens_details" { + let mut details = value.as_object().cloned().unwrap_or_default(); + details.remove("cached_creation_tokens"); + details.remove("cache_creation_tokens"); + Value::Object(details) + } else { + value.clone() + }; + (key.clone(), value) + }) + .collect::>(); + canonical.extensions = if provider_fields.is_empty() { + BTreeMap::new() + } else { + BTreeMap::from([( + OPENAI_RESPONSES_EXTENSION_NAMESPACE.to_string(), + Value::Object(provider_fields), + )]) + }; + Some(canonical) +} + pub(crate) fn claude_usage_to_canonical(value: Option<&Value>) -> Option { let usage = value?.as_object()?; let input_tokens = usage @@ -5989,9 +6152,29 @@ pub(crate) fn canonical_usage_to_openai_responses_usage(value: &CanonicalUsage) output["input_tokens_details"]["cache_write_tokens"] = Value::from(value.cache_write_tokens); } + if let (Some(output), Some(provider_fields)) = ( + output.as_object_mut(), + openai_responses_extension(&value.extensions).and_then(Value::as_object), + ) { + merge_json_object_missing(output, provider_fields); + } output } +fn merge_json_object_missing(target: &mut Map, source: &Map) { + for (key, source_value) in source { + match (target.get_mut(key), source_value) { + (Some(Value::Object(target_object)), Value::Object(source_object)) => { + merge_json_object_missing(target_object, source_object); + } + (Some(_), _) => {} + (None, _) => { + target.insert(key.clone(), source_value.clone()); + } + } + } +} + pub(crate) fn canonical_usage_to_claude(value: &CanonicalUsage) -> Value { let mut output = json!({ "input_tokens": canonical_usage_uncached_input_tokens(value), @@ -6394,6 +6577,41 @@ pub(crate) fn openai_responses_extension(extensions: &BTreeMap) - .or_else(|| extensions.get(OPENAI_RESPONSES_LEGACY_EXTENSION_NAMESPACE)) } +pub(crate) fn openai_service_tier_extension( + extensions: &BTreeMap, +) -> Option<&Value> { + [ + OPENAI_RESPONSES_EXTENSION_NAMESPACE, + OPENAI_RESPONSES_LEGACY_EXTENSION_NAMESPACE, + "openai", + ] + .into_iter() + .find_map(|namespace| { + extensions + .get(namespace) + .and_then(Value::as_object) + .and_then(|object| object.get("service_tier")) + }) +} + +pub(crate) fn openai_responses_item_extension_object( + extensions: &BTreeMap, + existing: &Map, +) -> Map { + openai_responses_extension(extensions) + .and_then(Value::as_object) + .map(|object| { + object + .iter() + .filter(|(key, _)| { + !existing.contains_key(*key) && !matches!(key.as_str(), "item_id" | "item_type") + }) + .map(|(key, value)| (key.clone(), value.clone())) + .collect() + }) + .unwrap_or_default() +} + pub(crate) fn strip_claude_billing_header(text: &str) -> String { let trimmed = text.trim(); let prefix = "x-anthropic-billing-header:"; @@ -8179,6 +8397,67 @@ mod tests { assert_eq!(rebuilt_gemini["usageMetadata"]["totalTokenCount"], 23); } + #[test] + fn openai_usage_prefers_cache_write_tokens_and_emits_the_official_field() { + let response = json!({ + "id": "resp_cache_write", + "model": "gpt-5.6-sol", + "status": "completed", + "output": [], + "usage": { + "input_tokens": 20, + "output_tokens": 4, + "total_tokens": 24, + "input_tokens_details": { + "cached_tokens": 3, + "cache_write_tokens": 7, + "cached_creation_tokens": 99 + } + } + }); + + let canonical = from_openai_responses_to_canonical_response(&response) + .expect("Responses usage should parse"); + assert_eq!(canonical.usage.as_ref().unwrap().cache_write_tokens, 7); + + let rebuilt = canonical_to_openai_responses_response(&canonical, &json!({})); + assert_eq!( + rebuilt["usage"]["input_tokens_details"]["cache_write_tokens"], + 7 + ); + assert!(rebuilt["usage"]["input_tokens_details"] + .get("cached_creation_tokens") + .is_none()); + + let rebuilt_chat = canonical_to_openai_chat_response(&canonical); + assert_eq!( + rebuilt_chat["usage"]["prompt_tokens_details"]["cache_write_tokens"], + 7 + ); + assert!(rebuilt_chat["usage"]["prompt_tokens_details"] + .get("cached_creation_tokens") + .is_none()); + } + + #[test] + fn openai_usage_accepts_cached_creation_tokens_as_legacy_input_alias() { + let response = json!({ + "id": "resp_cache_write_legacy", + "model": "gpt-5.6-sol", + "status": "completed", + "output": [], + "usage": { + "input_tokens": 10, + "output_tokens": 2, + "input_tokens_details": {"cached_creation_tokens": 5} + } + }); + + let canonical = from_openai_responses_to_canonical_response(&response) + .expect("legacy usage alias should parse"); + assert_eq!(canonical.usage.as_ref().unwrap().cache_write_tokens, 5); + } + #[test] fn gemini_request_adapter_preserves_thinking_tools_media_and_extensions() { let request = json!({ diff --git a/crates/aether-ai-formats/src/protocol/stream.rs b/crates/aether-ai-formats/src/protocol/stream.rs index 97b5e0097..c003de412 100644 --- a/crates/aether-ai-formats/src/protocol/stream.rs +++ b/crates/aether-ai-formats/src/protocol/stream.rs @@ -50,6 +50,11 @@ pub enum CanonicalStreamEvent { index: usize, item: Value, }, + OpenAiResponsesOutputItem { + output_index: Option, + item: Value, + raw_event: Value, + }, ToolCallStart { index: usize, call_id: String, diff --git a/crates/aether-ai-serving/src/attempt_loop.rs b/crates/aether-ai-serving/src/attempt_loop.rs index 512b4941c..f7af08f61 100644 --- a/crates/aether-ai-serving/src/attempt_loop.rs +++ b/crates/aether-ai-serving/src/attempt_loop.rs @@ -51,7 +51,14 @@ where while let Some(attempt) = remaining.next() { last_attempted = Some((attempt.execution_plan().clone(), attempt.report_context())); - if let Some(response) = port.execute_attempt(&attempt).await? { + let response = match port.execute_attempt(&attempt).await { + Ok(response) => response, + 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)); } @@ -94,3 +101,127 @@ impl AiExecutionAttempt for crate::dto::AiStreamAttempt { self.report_context.clone() } } + +#[cfg(test)] +mod tests { + use std::collections::BTreeMap; + use std::sync::Mutex; + + use async_trait::async_trait; + + use super::{run_ai_attempt_loop, AiAttemptLoopPort, AiExecutionAttempt}; + + #[derive(Clone)] + struct TestAttempt { + id: &'static str, + plan: aether_contracts::ExecutionPlan, + } + + impl AiExecutionAttempt for TestAttempt { + fn execution_plan(&self) -> &aether_contracts::ExecutionPlan { + &self.plan + } + + fn report_kind(&self) -> Option { + None + } + + fn report_context(&self) -> Option { + None + } + } + + struct FailingPort { + fail_on: &'static str, + unused: Mutex>, + } + + #[async_trait] + impl AiAttemptLoopPort for FailingPort { + type Response = (); + type Exhaustion = (); + type Error = &'static str; + + async fn execute_attempt( + &self, + attempt: &TestAttempt, + ) -> Result, Self::Error> { + if attempt.id == self.fail_on { + Err("attempt failed") + } else { + Ok(None) + } + } + + 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(()) + } + } + + fn attempt(id: &'static str) -> TestAttempt { + TestAttempt { + id, + plan: aether_contracts::ExecutionPlan { + request_id: format!("request-{id}"), + candidate_id: Some(id.to_string()), + provider_name: Some("provider".to_string()), + provider_id: "provider-1".to_string(), + endpoint_id: "endpoint-1".to_string(), + key_id: "key-1".to_string(), + method: "POST".to_string(), + url: "https://example.test/v1/responses".to_string(), + headers: BTreeMap::new(), + content_type: Some("application/json".to_string()), + content_encoding: None, + body: aether_contracts::RequestBody::from_json(serde_json::json!({})), + stream: false, + client_api_format: "openai:responses".to_string(), + provider_api_format: "openai:responses".to_string(), + model_name: Some("gpt-5.6-sol".to_string()), + proxy: None, + transport_profile: None, + timeouts: None, + }, + } + } + + #[tokio::test] + async fn marks_unattempted_candidates_unused_when_execution_returns_error() { + let port = FailingPort { + fail_on: "candidate-2", + unused: Mutex::new(Vec::new()), + }; + + let error = run_ai_attempt_loop( + &port, + vec![ + attempt("candidate-1"), + attempt("candidate-2"), + attempt("candidate-3"), + ], + ) + .await + .expect_err("second attempt should fail"); + + assert_eq!(error, "attempt failed"); + assert_eq!( + *port.unused.lock().expect("unused attempts should lock"), + vec!["candidate-3"] + ); + } +} diff --git a/crates/aether-ai-serving/src/attempt_plan.rs b/crates/aether-ai-serving/src/attempt_plan.rs index cbb5317d4..63bccbd3e 100644 --- a/crates/aether-ai-serving/src/attempt_plan.rs +++ b/crates/aether-ai-serving/src/attempt_plan.rs @@ -240,6 +240,7 @@ pub fn build_ai_execution_decision_from_plan( request_id: Some(request_id), candidate_id, provider_name, + provider_type: None, provider_id: Some(provider_id), endpoint_id: Some(endpoint_id), key_id: Some(key_id), @@ -792,6 +793,7 @@ mod tests { request_id: Some("req_1".to_string()), candidate_id: Some("candidate_1".to_string()), provider_name: Some("provider".to_string()), + provider_type: None, provider_id: Some("provider_1".to_string()), endpoint_id: Some("endpoint_1".to_string()), key_id: Some("key_1".to_string()), diff --git a/crates/aether-ai-serving/src/decision_payload.rs b/crates/aether-ai-serving/src/decision_payload.rs index a5755c74c..6d5125348 100644 --- a/crates/aether-ai-serving/src/decision_payload.rs +++ b/crates/aether-ai-serving/src/decision_payload.rs @@ -20,6 +20,7 @@ pub struct AiExecutionDecisionResponseParts { pub request_id: String, pub candidate_id: String, pub provider_name: String, + pub provider_type: String, pub provider_id: String, pub endpoint_id: String, pub key_id: String, @@ -65,6 +66,7 @@ pub fn build_ai_execution_decision_response( request_id: Some(parts.request_id), candidate_id: Some(parts.candidate_id), provider_name: Some(parts.provider_name), + provider_type: Some(parts.provider_type), provider_id: Some(parts.provider_id), endpoint_id: Some(parts.endpoint_id), key_id: Some(parts.key_id), @@ -209,6 +211,7 @@ mod tests { request_id: "trace-1".to_string(), candidate_id: "candidate-1".to_string(), provider_name: "OpenAI".to_string(), + provider_type: "openai".to_string(), provider_id: "provider-1".to_string(), endpoint_id: "endpoint-1".to_string(), key_id: "key-1".to_string(), diff --git a/crates/aether-ai-serving/src/dto.rs b/crates/aether-ai-serving/src/dto.rs index 52629adf6..686132cef 100644 --- a/crates/aether-ai-serving/src/dto.rs +++ b/crates/aether-ai-serving/src/dto.rs @@ -83,6 +83,8 @@ pub struct AiExecutionDecision { #[serde(default)] pub provider_name: Option, #[serde(default)] + pub provider_type: Option, + #[serde(default)] pub provider_id: Option, #[serde(default)] pub endpoint_id: Option, @@ -208,6 +210,7 @@ mod tests { request_id: None, candidate_id: None, provider_name: None, + provider_type: None, provider_id: None, endpoint_id: None, key_id: None, diff --git a/crates/aether-billing/src/default_rule.rs b/crates/aether-billing/src/default_rule.rs index ea1ef2f51..0fdb9a8d7 100644 --- a/crates/aether-billing/src/default_rule.rs +++ b/crates/aether-billing/src/default_rule.rs @@ -3,7 +3,7 @@ use std::collections::BTreeMap; use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; -use crate::pricing::BillingModelPricingSnapshot; +use crate::pricing::BillingPricingResolution; #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct VirtualBillingRule { @@ -20,12 +20,12 @@ pub struct DefaultBillingRuleGenerator; impl DefaultBillingRuleGenerator { pub fn generate_for_pricing( - pricing: &BillingModelPricingSnapshot, + global_model_name: &str, + pricing: &BillingPricingResolution, task_type: &str, ) -> Option { - let pricing_config = pricing.effective_tiered_pricing(); - let tiers = pricing - .effective_tiered_pricing() + let pricing_config = pricing.tiered_pricing.as_ref(); + let tiers = pricing_config .and_then(|value| value.get("tiers")) .and_then(Value::as_array) .cloned() @@ -41,10 +41,7 @@ impl DefaultBillingRuleGenerator { || has_image_output_ranges || explicit_image_output_price_default.is_some(); - if tiers.is_empty() - && pricing.effective_price_per_request().is_none() - && !has_image_output_pricing - { + if tiers.is_empty() && pricing.price_per_request.is_none() && !has_image_output_pricing { return None; } @@ -55,7 +52,7 @@ impl DefaultBillingRuleGenerator { tier_value_with_fallback(&first_tier, "cache_creation_price_per_1m", 1.25); let base_cache_read_price = tier_value_with_fallback(&first_tier, "cache_read_price_per_1m", 0.1); - let base_request_price = pricing.effective_price_per_request().unwrap_or(0.0); + let base_request_price = pricing.price_per_request.unwrap_or(0.0); let mut variables = BTreeMap::new(); variables.insert("input_price_per_1m".to_string(), json!(base_input_price)); @@ -237,7 +234,7 @@ impl DefaultBillingRuleGenerator { Some(VirtualBillingRule { id: "__default__".to_string(), - name: format!("Default rule for {}", pricing.global_model_name), + name: format!("Default rule for {global_model_name}"), task_type: normalize_task_type(task_type).to_string(), expression: "input_cost + output_cost + cache_creation_uncategorized_cost + cache_creation_ephemeral_5m_cost + cache_creation_ephemeral_1h_cost + cache_read_cost + image_output_cost + request_cost".to_string(), variables, diff --git a/crates/aether-billing/src/event_enrichment.rs b/crates/aether-billing/src/event_enrichment.rs index e2d1fb751..8cd34c47a 100644 --- a/crates/aether-billing/src/event_enrichment.rs +++ b/crates/aether-billing/src/event_enrichment.rs @@ -1,4 +1,10 @@ use aether_data_contracts::repository::billing::StoredBillingModelContext; +use aether_data_contracts::repository::usage::{ + extract_provider_actual_service_tier_from_response, + extract_provider_cache_ttl_minutes_from_metadata, extract_provider_service_tier_from_body, + normalize_provider_service_tier, resolve_provider_cache_ttl_minutes, + PROVIDER_ACTUAL_SERVICE_TIER_METADATA_KEY, PROVIDER_SERVICE_TIER_METADATA_KEY, +}; use aether_data_contracts::DataLayerError; use aether_usage_runtime::{UsageEvent, UsageEventType}; use async_trait::async_trait; @@ -137,6 +143,7 @@ fn calculate_billing_computation( } else { 1 }; + let processing_tiers = usage_event_processing_tiers(&event.data); let input = BillingUsageInput { task_type: if is_image_usage { "image".to_string() @@ -152,6 +159,8 @@ fn calculate_billing_computation( .endpoint_api_format .clone() .or_else(|| event.data.api_format.clone()), + requested_processing_tier: processing_tiers.requested, + actual_processing_tier: processing_tiers.actual, request_count, input_tokens: event.data.input_tokens.unwrap_or_default() as i64, output_tokens: event.data.output_tokens.unwrap_or_default() as i64, @@ -169,7 +178,8 @@ fn calculate_billing_computation( image_size: usage_event_dimension_string(&event.data, "image_size"), image_quality: usage_event_dimension_string(&event.data, "image_quality"), image_output_format: usage_event_dimension_string(&event.data, "image_output_format"), - cache_ttl_minutes: pricing.provider_api_key_cache_ttl_minutes, + cache_ttl_minutes: usage_event_provider_cache_ttl_minutes(&event.data) + .or(pricing.provider_api_key_cache_ttl_minutes), }; BillingService::new() @@ -179,6 +189,48 @@ fn calculate_billing_computation( }) } +#[derive(Debug, Clone, PartialEq, Eq)] +struct UsageEventProcessingTiers { + requested: Option, + actual: Option, +} + +fn usage_event_processing_tiers( + data: &aether_usage_runtime::UsageEventData, +) -> UsageEventProcessingTiers { + let metadata = data.request_metadata.as_ref().and_then(Value::as_object); + let requested = extract_provider_service_tier_from_body(data.provider_request_body.as_ref()) + .or_else(|| { + metadata + .and_then(|metadata| metadata.get(PROVIDER_SERVICE_TIER_METADATA_KEY)) + .and_then(Value::as_str) + .and_then(normalize_provider_service_tier) + }); + let actual = metadata + .and_then(|metadata| metadata.get(PROVIDER_ACTUAL_SERVICE_TIER_METADATA_KEY)) + .and_then(Value::as_str) + .and_then(normalize_provider_service_tier) + .or_else(|| { + extract_provider_actual_service_tier_from_response(data.response_body.as_ref()) + }); + + UsageEventProcessingTiers { requested, actual } +} + +fn usage_event_provider_cache_ttl_minutes( + data: &aether_usage_runtime::UsageEventData, +) -> Option { + resolve_provider_cache_ttl_minutes( + data.endpoint_api_format + .as_deref() + .or(data.api_format.as_deref()), + data.target_model.as_deref().or(Some(data.model.as_str())), + Some(data.model.as_str()), + data.provider_request_body.as_ref(), + ) + .or_else(|| extract_provider_cache_ttl_minutes_from_metadata(data.request_metadata.as_ref())) +} + fn usage_event_is_image_usage(data: &aether_usage_runtime::UsageEventData) -> bool { data.request_type .as_deref() @@ -262,54 +314,23 @@ fn apply_billing_computation( ) -> Result<(), DataLayerError> { event.data.total_cost_usd = Some(computation.cost_result.cost); event.data.actual_total_cost_usd = Some(computation.actual_total_cost); - merge_billing_snapshot_metadata( - &mut event.data.request_metadata, - pricing, - &computation.cost_result.snapshot, - computation.actual_total_cost, - computation.rate_multiplier, - computation.is_free_tier, - ) + merge_billing_snapshot_metadata(&mut event.data.request_metadata, pricing, &computation) } fn map_pricing_context(context: StoredBillingModelContext) -> BillingModelPricingSnapshot { - BillingModelPricingSnapshot { - provider_id: context.provider_id, - provider_billing_type: context.provider_billing_type, - provider_api_key_id: context.provider_api_key_id, - provider_api_key_rate_multipliers: context.provider_api_key_rate_multipliers, - provider_api_key_cache_ttl_minutes: context.provider_api_key_cache_ttl_minutes, - global_model_id: context.global_model_id, - global_model_name: context.global_model_name, - global_model_config: context.global_model_config, - default_price_per_request: context.default_price_per_request, - default_tiered_pricing: context.default_tiered_pricing, - model_id: context.model_id, - model_provider_model_name: context.model_provider_model_name, - model_config: context.model_config, - model_price_per_request: context.model_price_per_request, - model_tiered_pricing: context.model_tiered_pricing, - } + context.into() } fn merge_billing_snapshot_metadata( request_metadata: &mut Option, pricing: &BillingModelPricingSnapshot, - snapshot: &crate::BillingSnapshot, - actual_total_cost: f64, - rate_multiplier: f64, - is_free_tier: bool, + computation: &BillingComputation, ) -> Result<(), DataLayerError> { + let snapshot = &computation.cost_result.snapshot; let billing_snapshot = serde_json::to_value(snapshot).map_err(|err| { DataLayerError::UnexpectedValue(format!("failed to serialize billing snapshot: {err}")) })?; - let settlement_snapshot = build_settlement_snapshot( - pricing, - snapshot, - actual_total_cost, - rate_multiplier, - is_free_tier, - ); + let settlement_snapshot = build_settlement_snapshot(pricing, computation); let mut metadata = match request_metadata.take() { Some(Value::Object(object)) => object, @@ -325,19 +346,24 @@ fn merge_billing_snapshot_metadata( "billing_dimensions".to_string(), Value::Object(snapshot.resolved_dimensions.clone().into_iter().collect()), ); - metadata.insert("rate_multiplier".to_string(), Value::from(rate_multiplier)); - metadata.insert("is_free_tier".to_string(), Value::from(is_free_tier)); + metadata.insert( + "rate_multiplier".to_string(), + Value::from(computation.rate_multiplier), + ); + metadata.insert( + "is_free_tier".to_string(), + Value::from(computation.is_free_tier), + ); *request_metadata = Some(Value::Object(metadata)); Ok(()) } fn build_settlement_snapshot( pricing: &BillingModelPricingSnapshot, - snapshot: &crate::BillingSnapshot, - actual_total_cost: f64, - rate_multiplier: f64, - is_free_tier: bool, + computation: &BillingComputation, ) -> Value { + let snapshot = &computation.cost_result.snapshot; + let resolution = &computation.pricing_resolution; json!({ "schema_version": SETTLEMENT_SNAPSHOT_SCHEMA_VERSION, "pricing_snapshot": { @@ -348,11 +374,16 @@ fn build_settlement_snapshot( "global_model_name": pricing.global_model_name.clone(), "model_id": pricing.model_id.clone(), "provider_model_name": pricing.model_provider_model_name.clone(), - "pricing_source": pricing.pricing_source(), - "tiered_pricing": pricing.effective_tiered_pricing().cloned(), - "price_per_request": pricing.effective_price_per_request(), - "rate_multiplier": rate_multiplier, - "is_free_tier": is_free_tier, + "requested_processing_tier": resolution.requested_processing_tier, + "actual_processing_tier": resolution.actual_processing_tier, + "billing_processing_tier": resolution.billing_processing_tier, + "pricing_source": resolution.pricing_source(), + "tiered_pricing_source": resolution.tiered_pricing_source.map(|source| source.as_str()), + "price_per_request_source": resolution.price_per_request_source.map(|source| source.as_str()), + "tiered_pricing": resolution.tiered_pricing, + "price_per_request": resolution.price_per_request, + "rate_multiplier": computation.rate_multiplier, + "is_free_tier": computation.is_free_tier, }, "billing_plan_snapshot": { "rule_id": snapshot.rule_id.clone(), @@ -365,7 +396,7 @@ fn build_settlement_snapshot( "resolved_variables": snapshot.resolved_variables.clone(), "cost_breakdown": snapshot.cost_breakdown.clone(), "total_cost": snapshot.total_cost, - "actual_total_cost": actual_total_cost, + "actual_total_cost": computation.actual_total_cost, "status": snapshot.status, "calculated_at": snapshot.calculated_at.clone(), }) @@ -379,7 +410,9 @@ mod tests { use serde_json::json; use serde_json::Value; - use super::{enrich_usage_event_with_billing, BillingModelContextLookup}; + use super::{ + enrich_usage_event_with_billing, usage_event_processing_tiers, BillingModelContextLookup, + }; struct TestLookup { name_context: Option, @@ -409,6 +442,122 @@ mod tests { } } + #[test] + fn processing_tier_facts_keep_request_and_terminal_response_independent() { + let data = UsageEventData { + provider_request_body: Some(json!({"service_tier": "Priority"})), + response_body: Some(json!({"service_tier": "priority"})), + request_metadata: Some(json!({ + "provider_service_tier": "batch", + "provider_actual_service_tier": "Default" + })), + ..UsageEventData::default() + }; + + let tiers = usage_event_processing_tiers(&data); + + assert_eq!(tiers.requested.as_deref(), Some("priority")); + assert_eq!(tiers.actual.as_deref(), Some("default")); + } + + #[tokio::test] + async fn settlement_uses_effective_gpt_5_6_cache_ttl_after_body_capture() { + let lookup = TestLookup { + name_context: Some( + StoredBillingModelContext::new( + "provider-1".to_string(), + Some("pay_as_you_go".to_string()), + Some("key-1".to_string()), + None, + Some(60), + "global-model-1".to_string(), + "gpt-5.6-sol".to_string(), + None, + None, + Some(json!({ + "tiers": [{ + "up_to": null, + "input_price_per_1m": 5.0, + "output_price_per_1m": 30.0, + "cache_creation_price_per_1m": 6.25, + "cache_read_price_per_1m": 0.5, + "cache_ttl_pricing": [{ + "ttl_minutes": 60, + "cache_creation_price_per_1m": 100.0, + "cache_read_price_per_1m": 100.0 + }] + }] + })), + Some("model-1".to_string()), + Some("gpt-5.6-sol".to_string()), + None, + None, + None, + ) + .expect("billing context should build"), + ), + model_id_context: None, + }; + + for (request_id, provider_request_body, request_metadata) in [ + ( + "req-cache-body", + Some(json!({"model": "gpt-5.6-sol"})), + None, + ), + ( + "req-cache-metadata", + None, + Some(json!({"provider_cache_ttl_minutes": 30})), + ), + ] { + let mut event = UsageEvent::new( + UsageEventType::Completed, + request_id, + UsageEventData { + provider_name: "OpenAI".to_string(), + model: "gpt-5.6-sol".to_string(), + target_model: Some("gpt-5.6-sol".to_string()), + provider_id: Some("provider-1".to_string()), + provider_api_key_id: Some("key-1".to_string()), + request_type: Some("chat".to_string()), + api_format: Some("openai:responses".to_string()), + endpoint_api_format: Some("openai:responses".to_string()), + provider_request_body, + request_metadata, + input_tokens: Some(1_000_000), + cache_creation_input_tokens: Some(1_000_000), + status_code: Some(200), + ..UsageEventData::default() + }, + ); + + enrich_usage_event_with_billing(&lookup, &mut event) + .await + .expect("billing should succeed"); + + let snapshot = event + .data + .request_metadata + .as_ref() + .and_then(|value| value.get("billing_snapshot")) + .expect("billing snapshot should exist"); + assert_eq!( + snapshot + .get("resolved_dimensions") + .and_then(|value| value.get("cache_ttl_minutes")), + Some(&json!(30)) + ); + assert_eq!( + snapshot + .get("resolved_variables") + .and_then(|value| value.get("cache_creation_price_per_1m")), + Some(&json!(6.25)) + ); + assert_eq!(event.data.total_cost_usd, Some(6.25)); + } + } + #[tokio::test] async fn enriches_completed_usage_event_with_billing_snapshot() { let lookup = TestLookup { @@ -471,6 +620,181 @@ mod tests { ); } + #[tokio::test] + async fn settlement_uses_actual_processing_tier_catalog_and_source() { + let lookup = TestLookup { + name_context: Some( + StoredBillingModelContext::new( + "provider-1".to_string(), + Some("pay_as_you_go".to_string()), + Some("key-1".to_string()), + None, + Some(60), + "global-model-1".to_string(), + "gpt-5.6".to_string(), + None, + None, + Some(json!({ + "tiers": [{"up_to": null, "input_price_per_1m": 5.0, "output_price_per_1m": 30.0}], + "processing_tiers": { + "flex": {"tiers": [{"up_to": null, "input_price_per_1m": 2.5, "output_price_per_1m": 15.0}]} + } + })), + Some("model-1".to_string()), + Some("gpt-5.6-upstream".to_string()), + None, + None, + Some(json!({ + "processing_tiers": { + "priority": {"tiers": [{"up_to": 272000, "input_price_per_1m": 10.0, "output_price_per_1m": 60.0}]} + } + })), + ) + .expect("billing context should build"), + ), + model_id_context: None, + }; + let mut event = UsageEvent::new( + UsageEventType::Completed, + "req-billing-tier-1", + UsageEventData { + provider_name: "OpenAI".to_string(), + model: "gpt-5.6".to_string(), + provider_id: Some("provider-1".to_string()), + provider_api_key_id: Some("key-1".to_string()), + request_type: Some("chat".to_string()), + api_format: Some("openai:responses".to_string()), + endpoint_api_format: Some("openai:responses".to_string()), + provider_request_body: Some(json!({"service_tier": "priority"})), + response_body: Some(json!({"service_tier": "priority"})), + request_metadata: Some(json!({ + "provider_service_tier": "priority", + "provider_actual_service_tier": "flex" + })), + input_tokens: Some(1_000), + output_tokens: Some(100), + status_code: Some(200), + ..UsageEventData::default() + }, + ); + + enrich_usage_event_with_billing(&lookup, &mut event) + .await + .expect("billing should succeed"); + + let pricing_snapshot = event + .data + .request_metadata + .as_ref() + .and_then(|value| value.pointer("/settlement_snapshot/pricing_snapshot")) + .expect("settlement pricing snapshot should exist"); + assert_eq!(pricing_snapshot["requested_processing_tier"], "priority"); + assert_eq!(pricing_snapshot["actual_processing_tier"], "flex"); + assert_eq!(pricing_snapshot["billing_processing_tier"], "flex"); + assert_eq!(pricing_snapshot["tiered_pricing_source"], "global_default"); + assert_eq!( + pricing_snapshot["tiered_pricing"]["tiers"][0]["input_price_per_1m"], + 2.5 + ); + assert_eq!( + event + .data + .request_metadata + .as_ref() + .and_then(|value| { value.pointer("/billing_dimensions/actual_processing_tier") }), + Some(&json!("flex")) + ); + } + + #[tokio::test] + async fn actual_processing_catalog_controls_image_price_with_independent_fixed_price() { + let lookup = TestLookup { + name_context: Some( + StoredBillingModelContext::new( + "provider-1".to_string(), + Some("pay_as_you_go".to_string()), + Some("key-1".to_string()), + None, + None, + "global-image-1".to_string(), + "gpt-image-2".to_string(), + None, + Some(0.01), + Some(json!({ + "image_output_price_default": 0.1, + "processing_tiers": { + "flex": {"image_output_price_default": 0.2} + } + })), + Some("model-image-1".to_string()), + Some("gpt-image-2".to_string()), + None, + Some(0.02), + Some(json!({ + "processing_tiers": { + "priority": {"image_output_price_default": 0.4} + } + })), + ) + .expect("billing context should build"), + ), + model_id_context: None, + }; + let mut event = UsageEvent::new( + UsageEventType::Completed, + "req-image-processing-tier-1", + UsageEventData { + provider_name: "OpenAI Image".to_string(), + model: "gpt-image-2".to_string(), + provider_id: Some("provider-1".to_string()), + provider_api_key_id: Some("key-1".to_string()), + request_type: Some("image".to_string()), + api_format: Some("openai:image".to_string()), + endpoint_api_format: Some("openai:image".to_string()), + provider_request_body: Some(json!({"service_tier": "priority"})), + request_metadata: Some(json!({ + "provider_actual_service_tier": "flex", + "dimensions": {"image_count": 2} + })), + status_code: Some(200), + ..UsageEventData::default() + }, + ); + + enrich_usage_event_with_billing(&lookup, &mut event) + .await + .expect("billing should succeed"); + + assert_eq!(event.data.total_cost_usd, Some(0.44)); + assert_eq!(event.data.actual_total_cost_usd, Some(0.44)); + let metadata = event.data.request_metadata.as_ref().expect("metadata"); + let pricing = metadata + .pointer("/settlement_snapshot/pricing_snapshot") + .expect("pricing snapshot"); + assert_eq!(pricing["billing_processing_tier"], "flex"); + assert_eq!(pricing["tiered_pricing_source"], "global_default"); + assert_eq!(pricing["price_per_request_source"], "provider_override"); + assert_eq!(pricing["pricing_source"], "mixed"); + assert_eq!( + metadata + .pointer("/billing_snapshot/resolved_variables/image_output_price_per_image") + .and_then(Value::as_f64), + Some(0.2) + ); + assert_eq!( + metadata + .pointer("/billing_snapshot/cost_breakdown/image_output_cost") + .and_then(Value::as_f64), + Some(0.4) + ); + assert_eq!( + metadata + .pointer("/billing_snapshot/cost_breakdown/request_cost") + .and_then(Value::as_f64), + Some(0.04) + ); + } + #[tokio::test] async fn image_usage_uses_image_count_for_request_cost() { let lookup = TestLookup { diff --git a/crates/aether-billing/src/lib.rs b/crates/aether-billing/src/lib.rs index ca37ace25..d7930aaeb 100644 --- a/crates/aether-billing/src/lib.rs +++ b/crates/aether-billing/src/lib.rs @@ -22,7 +22,10 @@ pub use precision::{ quantize_cost, quantize_display, quantize_value, BILLING_DISPLAY_PRECISION, BILLING_STORAGE_PRECISION, }; -pub use pricing::{BillingComputation, BillingModelPricingSnapshot, BillingUsageInput}; +pub use pricing::{ + BillingAuthorizationEstimateInput, BillingComputation, BillingModelPricingSnapshot, + BillingPricingResolution, BillingPricingSource, BillingUsageInput, +}; pub use schema::{ BillingSnapshot, BillingSnapshotStatus, CostResult, BILLING_SNAPSHOT_SCHEMA_VERSION, }; diff --git a/crates/aether-billing/src/pricing.rs b/crates/aether-billing/src/pricing.rs index 1c584f770..871ad957a 100644 --- a/crates/aether-billing/src/pricing.rs +++ b/crates/aether-billing/src/pricing.rs @@ -1,6 +1,67 @@ +use std::collections::BTreeSet; + +use aether_data_contracts::repository::{ + billing::StoredBillingModelContext, usage::normalize_provider_service_tier, +}; use serde::{Deserialize, Serialize}; use serde_json::Value; +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum BillingPricingSource { + ProviderOverride, + GlobalDefault, +} + +impl BillingPricingSource { + pub fn as_str(self) -> &'static str { + match self { + Self::ProviderOverride => "provider_override", + Self::GlobalDefault => "global_default", + } + } +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct BillingPricingResolution { + pub requested_processing_tier: Option, + pub actual_processing_tier: Option, + pub billing_processing_tier: Option, + pub tiered_pricing: Option, + pub tiered_pricing_source: Option, + pub price_per_request: Option, + pub price_per_request_source: Option, +} + +impl BillingPricingResolution { + pub fn requires_actual_processing_tier(&self) -> bool { + self.billing_processing_tier.is_none() + } + + pub fn pricing_source(&self) -> &'static str { + match (self.tiered_pricing_source, self.price_per_request_source) { + (Some(tiered), Some(request)) if tiered != request => "mixed", + (Some(source), _) | (_, Some(source)) => source.as_str(), + (None, None) => "unpriced", + } + } + + pub fn bills_standard_processing_tier(&self) -> bool { + self.billing_processing_tier + .as_deref() + .is_some_and(processing_tier_is_standard) + } + + pub fn bills_requested_processing_tier(&self) -> bool { + let requested = self + .requested_processing_tier + .as_deref() + .map(canonical_processing_tier) + .unwrap_or_else(|| "standard".to_string()); + self.billing_processing_tier.as_deref() == Some(requested.as_str()) + } +} + #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct BillingModelPricingSnapshot { pub provider_id: String, @@ -21,32 +82,159 @@ pub struct BillingModelPricingSnapshot { } impl BillingModelPricingSnapshot { - pub fn effective_tiered_pricing(&self) -> Option<&Value> { + pub fn resolve_pricing( + &self, + requested_processing_tier: Option<&str>, + actual_processing_tier: Option<&str>, + ) -> BillingPricingResolution { + let requested_processing_tier = normalize_processing_tier(requested_processing_tier); + let actual_processing_tier = normalize_processing_tier(actual_processing_tier); + let billing_processing_tier = actual_processing_tier + .as_deref() + .map(canonical_processing_tier) + .or_else(|| { + requested_processing_tier.as_deref().map_or( + Some("standard".to_string()), + |requested| { + processing_tier_is_standard(requested).then(|| "standard".to_string()) + }, + ) + }); + + let (tiered_pricing, tiered_pricing_source) = billing_processing_tier + .as_deref() + .and_then(|tier| self.resolve_tiered_pricing(tier)) + .map_or((None, None), |(pricing, source)| { + (Some(pricing.clone()), Some(source)) + }); + let (price_per_request, price_per_request_source) = self + .resolve_price_per_request() + .map_or((None, None), |(price, source)| (Some(price), Some(source))); + + BillingPricingResolution { + requested_processing_tier, + actual_processing_tier, + billing_processing_tier, + tiered_pricing, + tiered_pricing_source, + price_per_request, + price_per_request_source, + } + } + + pub fn resolve_authorization_pricing_candidates( + &self, + requested_processing_tier: Option<&str>, + ) -> Option> { + let requested_processing_tier = normalize_processing_tier(requested_processing_tier); + let requested_billing_tier = requested_processing_tier + .as_deref() + .map(canonical_processing_tier) + .unwrap_or_else(|| "standard".to_string()); + let requested_resolution = self.authorization_pricing_for_tier( + requested_processing_tier.clone(), + Some(requested_billing_tier.clone()), + ); + if !processing_tier_is_standard(&requested_billing_tier) + && requested_resolution.tiered_pricing.is_none() + { + return None; + } + + let mut billing_tiers = BTreeSet::from(["standard".to_string(), requested_billing_tier]); + for pricing in [ + self.model_tiered_pricing.as_ref(), + self.default_tiered_pricing.as_ref(), + ] + .into_iter() + .flatten() + { + let Some(processing_tiers) = pricing.get("processing_tiers").and_then(Value::as_object) + else { + continue; + }; + billing_tiers.extend(processing_tiers.keys().filter_map(|tier| { + normalize_processing_tier(Some(tier)).map(|tier| canonical_processing_tier(&tier)) + })); + } + + let candidates = billing_tiers + .into_iter() + .filter_map(|billing_tier| { + let resolution = self.authorization_pricing_for_tier( + requested_processing_tier.clone(), + Some(billing_tier), + ); + (resolution.bills_standard_processing_tier() || resolution.tiered_pricing.is_some()) + .then_some(resolution) + }) + .collect::>(); + (!candidates.is_empty()).then_some(candidates) + } + + fn authorization_pricing_for_tier( + &self, + requested_processing_tier: Option, + billing_processing_tier: Option, + ) -> BillingPricingResolution { + let (tiered_pricing, tiered_pricing_source) = billing_processing_tier + .as_deref() + .and_then(|tier| self.resolve_tiered_pricing(tier)) + .map_or((None, None), |(pricing, source)| { + (Some(pricing.clone()), Some(source)) + }); + let (price_per_request, price_per_request_source) = self + .resolve_price_per_request() + .map_or((None, None), |(price, source)| (Some(price), Some(source))); + + BillingPricingResolution { + requested_processing_tier, + actual_processing_tier: None, + billing_processing_tier, + tiered_pricing, + tiered_pricing_source, + price_per_request, + price_per_request_source, + } + } + + fn resolve_tiered_pricing( + &self, + processing_tier: &str, + ) -> Option<(&Value, BillingPricingSource)> { + if processing_tier_is_standard(processing_tier) { + return self + .model_tiered_pricing + .as_ref() + .filter(|value| has_pricing_data(value)) + .map(|value| (value, BillingPricingSource::ProviderOverride)) + .or_else(|| { + self.default_tiered_pricing + .as_ref() + .filter(|value| has_pricing_data(value)) + .map(|value| (value, BillingPricingSource::GlobalDefault)) + }); + } + self.model_tiered_pricing .as_ref() - .filter(|value| has_pricing_data(value)) - .or(self.default_tiered_pricing.as_ref()) + .and_then(|pricing| processing_tier_overlay(pricing, processing_tier)) + .map(|value| (value, BillingPricingSource::ProviderOverride)) + .or_else(|| { + self.default_tiered_pricing + .as_ref() + .and_then(|pricing| processing_tier_overlay(pricing, processing_tier)) + .map(|value| (value, BillingPricingSource::GlobalDefault)) + }) } - pub fn effective_price_per_request(&self) -> Option { + fn resolve_price_per_request(&self) -> Option<(f64, BillingPricingSource)> { self.model_price_per_request - .or(self.default_price_per_request) - } - - pub fn pricing_source(&self) -> &'static str { - if self - .model_tiered_pricing - .as_ref() - .is_some_and(has_pricing_data) - || self.model_price_per_request.is_some() - { - "provider_override" - } else if self.default_tiered_pricing.is_some() || self.default_price_per_request.is_some() - { - "global_default" - } else { - "unpriced" - } + .map(|price| (price, BillingPricingSource::ProviderOverride)) + .or_else(|| { + self.default_price_per_request + .map(|price| (price, BillingPricingSource::GlobalDefault)) + }) } pub fn is_free_tier(&self) -> bool { @@ -75,6 +263,74 @@ impl BillingModelPricingSnapshot { } } +impl From<&StoredBillingModelContext> for BillingModelPricingSnapshot { + fn from(context: &StoredBillingModelContext) -> Self { + Self { + provider_id: context.provider_id.clone(), + provider_billing_type: context.provider_billing_type.clone(), + provider_api_key_id: context.provider_api_key_id.clone(), + provider_api_key_rate_multipliers: context.provider_api_key_rate_multipliers.clone(), + provider_api_key_cache_ttl_minutes: context.provider_api_key_cache_ttl_minutes, + global_model_id: context.global_model_id.clone(), + global_model_name: context.global_model_name.clone(), + global_model_config: context.global_model_config.clone(), + default_price_per_request: context.default_price_per_request, + default_tiered_pricing: context.default_tiered_pricing.clone(), + model_id: context.model_id.clone(), + model_provider_model_name: context.model_provider_model_name.clone(), + model_config: context.model_config.clone(), + model_price_per_request: context.model_price_per_request, + model_tiered_pricing: context.model_tiered_pricing.clone(), + } + } +} + +impl From for BillingModelPricingSnapshot { + fn from(context: StoredBillingModelContext) -> Self { + Self { + provider_id: context.provider_id, + provider_billing_type: context.provider_billing_type, + provider_api_key_id: context.provider_api_key_id, + provider_api_key_rate_multipliers: context.provider_api_key_rate_multipliers, + provider_api_key_cache_ttl_minutes: context.provider_api_key_cache_ttl_minutes, + global_model_id: context.global_model_id, + global_model_name: context.global_model_name, + global_model_config: context.global_model_config, + default_price_per_request: context.default_price_per_request, + default_tiered_pricing: context.default_tiered_pricing, + model_id: context.model_id, + model_provider_model_name: context.model_provider_model_name, + model_config: context.model_config, + model_price_per_request: context.model_price_per_request, + model_tiered_pricing: context.model_tiered_pricing, + } + } +} + +fn normalize_processing_tier(value: Option<&str>) -> Option { + value.and_then(normalize_provider_service_tier) +} + +fn canonical_processing_tier(value: &str) -> String { + if processing_tier_is_standard(value) { + "standard".to_string() + } else { + value.to_string() + } +} + +fn processing_tier_is_standard(value: &str) -> bool { + matches!(value, "auto" | "default" | "standard") +} + +fn processing_tier_overlay<'a>(pricing: &'a Value, tier: &str) -> Option<&'a Value> { + pricing + .get("processing_tiers") + .and_then(Value::as_object) + .and_then(|tiers| tiers.get(tier)) + .filter(|value| has_pricing_data(value)) +} + fn has_pricing_data(value: &Value) -> bool { value .get("tiers") @@ -104,7 +360,7 @@ fn value_has_entries(value: &Value) -> bool { mod tests { use serde_json::json; - use super::BillingModelPricingSnapshot; + use super::{BillingModelPricingSnapshot, BillingPricingSource}; fn snapshot( model_tiered_pricing: Option, @@ -135,11 +391,21 @@ mod tests { json!({"tiers":[{"up_to":null,"input_price_per_1m":3.0,"output_price_per_1m":15.0}]}); let pricing = snapshot(Some(json!({})), Some(default_pricing.clone())); - assert_eq!(pricing.effective_tiered_pricing(), Some(&default_pricing)); + let resolution = pricing.resolve_pricing(None, None); + assert_eq!(resolution.tiered_pricing, Some(default_pricing.clone())); + assert_eq!( + resolution.tiered_pricing_source, + Some(BillingPricingSource::GlobalDefault) + ); let pricing = snapshot(Some(json!({"tiers": []})), Some(default_pricing.clone())); - assert_eq!(pricing.effective_tiered_pricing(), Some(&default_pricing)); + let resolution = pricing.resolve_pricing(None, None); + assert_eq!(resolution.tiered_pricing, Some(default_pricing)); + assert_eq!( + resolution.tiered_pricing_source, + Some(BillingPricingSource::GlobalDefault) + ); } #[test] @@ -150,7 +416,140 @@ mod tests { json!({"tiers":[{"up_to":null,"input_price_per_1m":3.0,"output_price_per_1m":15.0}]}); let pricing = snapshot(Some(provider_pricing.clone()), Some(default_pricing)); - assert_eq!(pricing.effective_tiered_pricing(), Some(&provider_pricing)); + let resolution = pricing.resolve_pricing(None, None); + assert_eq!(resolution.tiered_pricing, Some(provider_pricing)); + assert_eq!( + resolution.tiered_pricing_source, + Some(BillingPricingSource::ProviderOverride) + ); + } + + #[test] + fn explicit_nonstandard_request_requires_actual_tier() { + let pricing = snapshot( + None, + Some(json!({ + "tiers": [{"up_to": null, "input_price_per_1m": 3.0}], + "processing_tiers": { + "priority": {"tiers": [{"up_to": 272000, "input_price_per_1m": 6.0}]} + } + })), + ); + + let resolution = pricing.resolve_pricing(Some("Priority"), None); + + assert!(resolution.requires_actual_processing_tier()); + assert_eq!( + resolution.requested_processing_tier.as_deref(), + Some("priority") + ); + assert_eq!(resolution.billing_processing_tier, None); + assert_eq!(resolution.tiered_pricing, None); + } + + #[test] + fn actual_tier_selects_exact_catalog_and_source() { + let pricing = snapshot( + Some(json!({ + "processing_tiers": { + "priority": {"tiers": [{"up_to": 272000, "input_price_per_1m": 9.0}]} + } + })), + Some(json!({ + "tiers": [{"up_to": null, "input_price_per_1m": 3.0}], + "processing_tiers": { + "flex": {"tiers": [{"up_to": null, "input_price_per_1m": 1.5}]} + } + })), + ); + + let flex = pricing.resolve_pricing(Some("priority"), Some("flex")); + assert_eq!(flex.billing_processing_tier.as_deref(), Some("flex")); + assert_eq!( + flex.tiered_pricing_source, + Some(BillingPricingSource::GlobalDefault) + ); + assert_eq!( + flex.tiered_pricing + .as_ref() + .and_then(|value| value.pointer("/tiers/0/input_price_per_1m")) + .and_then(serde_json::Value::as_f64), + Some(1.5) + ); + + let standard = pricing.resolve_pricing(Some("priority"), Some("Default")); + assert_eq!(standard.actual_processing_tier.as_deref(), Some("default")); + assert_eq!( + standard.billing_processing_tier.as_deref(), + Some("standard") + ); + assert_eq!( + standard.tiered_pricing_source, + Some(BillingPricingSource::GlobalDefault) + ); + } + + #[test] + fn tiered_and_fixed_price_sources_are_recorded_independently() { + let mut pricing = snapshot( + None, + Some(json!({"tiers": [{"up_to": null, "input_price_per_1m": 3.0}]})), + ); + pricing.model_price_per_request = Some(0.02); + + let resolution = pricing.resolve_pricing(None, None); + + assert_eq!( + resolution.tiered_pricing_source, + Some(BillingPricingSource::GlobalDefault) + ); + assert_eq!( + resolution.price_per_request_source, + Some(BillingPricingSource::ProviderOverride) + ); + assert_eq!(resolution.pricing_source(), "mixed"); + } + + #[test] + fn authorization_candidates_include_requested_catalog_without_inventing_actual_tier() { + let pricing = snapshot( + None, + Some(json!({ + "tiers": [{"up_to": null, "input_price_per_1m": 3.0}], + "processing_tiers": { + "priority": {"tiers": [{"up_to": 272000, "input_price_per_1m": 6.0}]} + } + })), + ); + + let candidates = pricing + .resolve_authorization_pricing_candidates(Some("Priority")) + .expect("authorization catalogs should resolve"); + let resolution = candidates + .iter() + .find(|resolution| resolution.billing_processing_tier.as_deref() == Some("priority")) + .expect("priority catalog should be included"); + + assert_eq!( + resolution.requested_processing_tier.as_deref(), + Some("priority") + ); + assert_eq!(resolution.actual_processing_tier, None); + assert_eq!( + resolution.billing_processing_tier.as_deref(), + Some("priority") + ); + assert_eq!( + resolution + .tiered_pricing + .as_ref() + .and_then(|value| value.pointer("/tiers/0/input_price_per_1m")) + .and_then(serde_json::Value::as_f64), + Some(6.0) + ); + assert!(candidates.iter().any(|resolution| { + resolution.billing_processing_tier.as_deref() == Some("standard") + })); } } @@ -158,6 +557,10 @@ mod tests { pub struct BillingUsageInput { pub task_type: String, pub api_format: Option, + #[serde(default)] + pub requested_processing_tier: Option, + #[serde(default)] + pub actual_processing_tier: Option, pub request_count: i64, pub input_tokens: i64, pub output_tokens: i64, @@ -172,11 +575,37 @@ pub struct BillingUsageInput { pub cache_ttl_minutes: Option, } +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct BillingAuthorizationEstimateInput { + pub task_type: String, + pub api_format: Option, + pub requested_processing_tier: Option, + #[serde(default)] + pub cache_ttl_minutes: Option, + pub input_tokens: i64, + pub max_output_tokens: Option, +} + +impl BillingAuthorizationEstimateInput { + pub fn new(task_type: impl Into, input_tokens: i64) -> Self { + Self { + task_type: task_type.into(), + api_format: None, + requested_processing_tier: None, + cache_ttl_minutes: None, + input_tokens: input_tokens.max(0), + max_output_tokens: None, + } + } +} + impl BillingUsageInput { pub fn new(task_type: impl Into) -> Self { Self { task_type: task_type.into(), api_format: None, + requested_processing_tier: None, + actual_processing_tier: None, request_count: 1, input_tokens: 0, output_tokens: 0, @@ -199,4 +628,5 @@ pub struct BillingComputation { pub actual_total_cost: f64, pub rate_multiplier: f64, pub is_free_tier: bool, + pub pricing_resolution: BillingPricingResolution, } diff --git a/crates/aether-billing/src/service.rs b/crates/aether-billing/src/service.rs index 5672cfc11..377e0eebf 100644 --- a/crates/aether-billing/src/service.rs +++ b/crates/aether-billing/src/service.rs @@ -8,7 +8,10 @@ use crate::default_rule::{ explicit_image_output_price_ranges, normalize_task_type, DefaultBillingRuleGenerator, }; use crate::precision::quantize_cost; -use crate::pricing::{BillingComputation, BillingModelPricingSnapshot, BillingUsageInput}; +use crate::pricing::{ + BillingAuthorizationEstimateInput, BillingComputation, BillingModelPricingSnapshot, + BillingPricingResolution, BillingUsageInput, +}; use crate::schema::{ BillingSnapshot, BillingSnapshotStatus, CostResult, BILLING_SNAPSHOT_SCHEMA_VERSION, }; @@ -33,39 +36,209 @@ impl BillingService { pricing: &BillingModelPricingSnapshot, input: &BillingUsageInput, ) -> Result { - let Some(rule) = - DefaultBillingRuleGenerator::generate_for_pricing(pricing, &input.task_type) - else { - return Ok(BillingComputation { - cost_result: CostResult { - cost: 0.0, - status: BillingSnapshotStatus::NoRule, - snapshot: BillingSnapshot { - schema_version: BILLING_SNAPSHOT_SCHEMA_VERSION.to_string(), - rule_id: None, - rule_name: None, - scope: None, - expression: None, - resolved_dimensions: build_dimensions(input, pricing), - resolved_variables: BTreeMap::new(), - cost_breakdown: BTreeMap::new(), - total_cost: 0.0, - tier_index: None, - tier_info: None, - missing_required: Vec::new(), - status: BillingSnapshotStatus::NoRule, - calculated_at: now_marker(), - engine_version: "2.0".to_string(), - }, - }, - actual_total_cost: 0.0, - rate_multiplier: pricing - .rate_multiplier_for_api_format(input.api_format.as_deref()), - is_free_tier: pricing.is_free_tier(), - }); + let pricing_resolution = pricing.resolve_pricing( + input.requested_processing_tier.as_deref(), + input.actual_processing_tier.as_deref(), + ); + self.calculate_with_resolution(pricing, input, pricing_resolution) + } + + pub fn estimate_authorization_cost_upper_bound( + &self, + pricing: &BillingModelPricingSnapshot, + estimate: &BillingAuthorizationEstimateInput, + ) -> Result, ExpressionEvaluationError> { + if normalize_task_type(&estimate.task_type) == "image" { + return Ok(None); + } + let Some(pricing_resolutions) = pricing.resolve_authorization_pricing_candidates( + estimate.requested_processing_tier.as_deref(), + ) else { + return Ok(None); + }; + if pricing.is_free_tier() { + return Ok(Some(0.0)); + } + if estimate.max_output_tokens.is_none() + && pricing_resolutions.iter().any(|resolution| { + resolution + .tiered_pricing + .as_ref() + .is_some_and(pricing_has_positive_output_rate) + }) + { + return Ok(None); + } + + let input_tokens = estimate.input_tokens.max(0); + let output_tokens = estimate.max_output_tokens.unwrap_or(0).max(0); + let base_input = BillingUsageInput { + task_type: estimate.task_type.clone(), + api_format: estimate.api_format.clone(), + requested_processing_tier: estimate.requested_processing_tier.clone(), + actual_processing_tier: None, + input_tokens, + output_tokens, + cache_ttl_minutes: estimate + .cache_ttl_minutes + .or(pricing.provider_api_key_cache_ttl_minutes), + ..BillingUsageInput::new(estimate.task_type.clone()) + }; + let mut scenarios = vec![base_input.clone()]; + if input_tokens > 0 + && pricing_resolutions + .iter() + .any(|resolution| resolution.tiered_pricing.is_some()) + { + let mut cache_creation = base_input.clone(); + cache_creation.cache_creation_tokens = input_tokens; + scenarios.push(cache_creation); + + if estimate.cache_ttl_minutes.is_none() { + let mut cache_creation_5m = base_input.clone(); + cache_creation_5m.cache_creation_tokens = input_tokens; + cache_creation_5m.cache_creation_ephemeral_5m_tokens = input_tokens; + cache_creation_5m.cache_ttl_minutes = Some(5); + scenarios.push(cache_creation_5m); + + let mut cache_creation_1h = base_input.clone(); + cache_creation_1h.cache_creation_tokens = input_tokens; + cache_creation_1h.cache_creation_ephemeral_1h_tokens = input_tokens; + cache_creation_1h.cache_ttl_minutes = Some(60); + scenarios.push(cache_creation_1h); + } + + let mut cache_read = base_input; + cache_read.cache_read_tokens = input_tokens; + scenarios.push(cache_read); + } + + let mut upper_bound = 0.0_f64; + 'pricing_catalogs: for pricing_resolution in pricing_resolutions { + let is_requested_catalog = pricing_resolution.bills_requested_processing_tier(); + for scenario in &scenarios { + let total_input_context = normalize_total_input_context_for_cache_hit_rate( + scenario.api_format.as_deref(), + scenario.input_tokens, + scenario.cache_creation_tokens, + scenario.cache_read_tokens, + ); + let Some(pricing_candidates) = + authorization_pricing_candidates(&pricing_resolution, total_input_context) + else { + return Ok(None); + }; + + // Validate the selected catalog and its finite coverage using the same path as + // settlement before evaluating every reachable tier as an upper-bound candidate. + let selected = + self.calculate_with_resolution(pricing, scenario, pricing_resolution.clone())?; + if !billing_computation_is_bounded(&selected) { + if !is_requested_catalog + && billing_computation_is_outside_catalog_context(&selected) + { + continue 'pricing_catalogs; + } + return Ok(None); + } + + for candidate in pricing_candidates { + let computation = + self.calculate_with_resolution(pricing, scenario, candidate)?; + if !billing_computation_is_bounded(&computation) { + return Ok(None); + } + upper_bound = upper_bound.max(computation.actual_total_cost); + } + } + } + Ok(Some(upper_bound)) + } + + fn calculate_with_resolution( + &self, + pricing: &BillingModelPricingSnapshot, + input: &BillingUsageInput, + pricing_resolution: BillingPricingResolution, + ) -> Result { + if pricing_resolution.requires_actual_processing_tier() { + return Ok(no_rule_computation( + pricing, + input, + pricing_resolution, + "actual_processing_tier", + )); + } + if !pricing_resolution.bills_standard_processing_tier() + && pricing_resolution.tiered_pricing.is_none() + { + return Ok(no_rule_computation( + pricing, + input, + pricing_resolution, + "processing_tier_catalog", + )); + } + + let total_input_context = normalize_total_input_context_for_cache_hit_rate( + input.api_format.as_deref(), + input.input_tokens, + input.cache_creation_tokens, + input.cache_read_tokens, + ); + let has_token_usage = input.input_tokens > 0 + || input.output_tokens > 0 + || input.cache_creation_tokens > 0 + || input.cache_read_tokens > 0; + if has_token_usage { + if let Some(pricing_config) = pricing_resolution.tiered_pricing.as_ref() { + let tiers = pricing_config + .get("tiers") + .and_then(Value::as_array) + .map(Vec::as_slice) + .unwrap_or_default(); + if tiers.is_empty() && pricing_resolution.price_per_request.is_none() { + return Ok(no_rule_computation( + pricing, + input, + pricing_resolution, + "token_pricing", + )); + } + if !tiers.is_empty() + && !pricing_covers_input_context(pricing_config, total_input_context) + { + return Ok(no_rule_computation( + pricing, + input, + pricing_resolution, + "input_context_tier", + )); + } + } else if pricing_resolution.price_per_request.is_none() { + return Ok(no_rule_computation( + pricing, + input, + pricing_resolution, + "token_pricing", + )); + } + } + + let Some(rule) = DefaultBillingRuleGenerator::generate_for_pricing( + &pricing.global_model_name, + &pricing_resolution, + &input.task_type, + ) else { + return Ok(no_rule_computation( + pricing, + input, + pricing_resolution, + "pricing_rule", + )); }; - let dims = build_dimensions(input, pricing); + let dims = build_dimensions(input, &pricing_resolution); let result = self.engine.evaluate( &rule.expression, Some(&rule.variables), @@ -116,10 +289,144 @@ impl BillingService { actual_total_cost, rate_multiplier, is_free_tier, + pricing_resolution, }) } } +fn pricing_has_positive_output_rate(pricing: &Value) -> bool { + pricing + .get("tiers") + .and_then(Value::as_array) + .into_iter() + .flatten() + .filter_map(|tier| tier.get("output_price_per_1m").and_then(Value::as_f64)) + .any(|price| price.is_finite() && price > 0.0) +} + +fn billing_computation_is_bounded(computation: &BillingComputation) -> bool { + computation.cost_result.status == BillingSnapshotStatus::Complete + && computation.actual_total_cost.is_finite() + && computation.actual_total_cost >= 0.0 +} + +fn billing_computation_is_outside_catalog_context(computation: &BillingComputation) -> bool { + computation.cost_result.status == BillingSnapshotStatus::NoRule + && computation.cost_result.snapshot.missing_required == ["input_context_tier"] +} + +fn authorization_pricing_candidates( + pricing: &BillingPricingResolution, + max_input_context: i64, +) -> Option> { + let Some(config) = pricing.tiered_pricing.as_ref() else { + return Some(vec![pricing.clone()]); + }; + let Some(tiers) = config.get("tiers").and_then(Value::as_array) else { + return Some(vec![pricing.clone()]); + }; + if tiers.is_empty() { + return Some(vec![pricing.clone()]); + } + + let max_input_context = max_input_context.max(0); + let mut previous_up_to: Option = None; + let mut candidates = Vec::new(); + for tier in tiers { + let tier_object = tier.as_object()?; + let up_to = match tier_object.get("up_to") { + None | Some(Value::Null) => None, + Some(value) => Some(nonnegative_i64(value)?), + }; + if let (Some(previous), Some(current)) = (previous_up_to, up_to) { + if current < previous { + return None; + } + } + + let lower_bound = previous_up_to.map_or(0, |value| value.saturating_add(1)); + if lower_bound <= max_input_context { + let mut candidate_config = config.clone(); + let candidate_object = candidate_config.as_object_mut()?; + let mut candidate_tier = tier.clone(); + candidate_tier + .as_object_mut()? + .insert("up_to".to_string(), Value::Null); + candidate_object.insert("tiers".to_string(), Value::Array(vec![candidate_tier])); + + let mut candidate = pricing.clone(); + candidate.tiered_pricing = Some(candidate_config); + candidates.push(candidate); + } + + match up_to { + Some(up_to) if max_input_context > up_to => previous_up_to = Some(up_to), + _ => break, + } + } + + (!candidates.is_empty()).then_some(candidates) +} + +fn nonnegative_i64(value: &Value) -> Option { + value + .as_i64() + .or_else(|| value.as_u64().and_then(|value| i64::try_from(value).ok())) + .filter(|value| *value >= 0) +} + +fn no_rule_computation( + pricing: &BillingModelPricingSnapshot, + input: &BillingUsageInput, + pricing_resolution: BillingPricingResolution, + missing_required: &str, +) -> BillingComputation { + let resolved_dimensions = build_dimensions(input, &pricing_resolution); + BillingComputation { + cost_result: CostResult { + cost: 0.0, + status: BillingSnapshotStatus::NoRule, + snapshot: BillingSnapshot { + schema_version: BILLING_SNAPSHOT_SCHEMA_VERSION.to_string(), + rule_id: None, + rule_name: None, + scope: None, + expression: None, + resolved_dimensions, + resolved_variables: BTreeMap::new(), + cost_breakdown: BTreeMap::new(), + total_cost: 0.0, + tier_index: None, + tier_info: None, + missing_required: vec![missing_required.to_string()], + status: BillingSnapshotStatus::NoRule, + calculated_at: now_marker(), + engine_version: "2.0".to_string(), + }, + }, + actual_total_cost: 0.0, + rate_multiplier: pricing.rate_multiplier_for_api_format(input.api_format.as_deref()), + is_free_tier: pricing.is_free_tier(), + pricing_resolution, + } +} + +fn pricing_covers_input_context(pricing: &Value, total_input_context: i64) -> bool { + let Some(tiers) = pricing.get("tiers").and_then(Value::as_array) else { + return true; + }; + let Some(last_tier) = tiers.last() else { + return true; + }; + match last_tier.get("up_to") { + None | Some(Value::Null) => true, + Some(value) => value + .as_i64() + .or_else(|| value.as_u64().and_then(|value| i64::try_from(value).ok())) + .is_some_and(|up_to| total_input_context <= up_to), + } +} + impl Default for BillingService { fn default() -> Self { Self::new() @@ -128,7 +435,7 @@ impl Default for BillingService { fn build_dimensions( input: &BillingUsageInput, - pricing: &BillingModelPricingSnapshot, + pricing: &BillingPricingResolution, ) -> BTreeMap { let normalized_input_tokens = normalize_input_tokens_for_billing( input.api_format.as_deref(), @@ -149,8 +456,9 @@ fn build_dimensions( input.cache_creation_tokens, input.cache_read_tokens, ); - let image_output_pricing = image_output_pricing_state(pricing); - let image_output_resolution = resolve_image_output_price_resolution(pricing, input); + let pricing_config = pricing.tiered_pricing.as_ref(); + let image_output_pricing = image_output_pricing_state(pricing_config); + let image_output_resolution = resolve_image_output_price_resolution(pricing_config, input); let mut out = BTreeMap::from([ ("input_tokens".to_string(), json!(normalized_input_tokens)), @@ -216,6 +524,18 @@ fn build_dimensions( "effective_task_type".to_string(), json!(normalize_task_type(&input.task_type)), ), + ( + "requested_processing_tier".to_string(), + json!(pricing.requested_processing_tier), + ), + ( + "actual_processing_tier".to_string(), + json!(pricing.actual_processing_tier), + ), + ( + "billing_processing_tier".to_string(), + json!(pricing.billing_processing_tier), + ), ]); out.insert( @@ -304,7 +624,7 @@ struct ParsedImageOutputPriceRange { prices: BTreeMap, } -fn image_output_pricing_state(pricing: &BillingModelPricingSnapshot) -> ImageOutputPricingState { +fn image_output_pricing_state(pricing: Option<&Value>) -> ImageOutputPricingState { let matrix_enabled = pricing_has_image_output_matrix(pricing); let range_enabled = pricing_has_image_output_ranges(pricing); let default_enabled = pricing_has_image_output_default_price(pricing); @@ -316,11 +636,10 @@ fn image_output_pricing_state(pricing: &BillingModelPricingSnapshot) -> ImageOut } fn resolve_image_output_price_resolution( - pricing: &BillingModelPricingSnapshot, + pricing: Option<&Value>, input: &BillingUsageInput, ) -> ImageOutputPriceResolution { - let pricing_config = pricing.effective_tiered_pricing(); - let default_price = explicit_image_output_price_default(pricing_config); + let default_price = explicit_image_output_price_default(pricing); let image_size = input .image_size .as_deref() @@ -335,7 +654,7 @@ fn resolve_image_output_price_resolution( if let (Some(size), Some(entries)) = ( image_size.as_deref(), - explicit_image_output_price_entries(pricing_config), + explicit_image_output_price_entries(pricing), ) { for key in image_price_lookup_keys(size, image_quality.as_deref()) { if let Some(price) = entries.get(&key).and_then(Value::as_f64) { @@ -351,7 +670,7 @@ fn resolve_image_output_price_resolution( if let Some(pixels) = image_pixels { if let Some((price, bucket)) = resolve_image_output_range_price( - explicit_image_output_price_ranges(pricing_config).unwrap_or_default(), + explicit_image_output_price_ranges(pricing).unwrap_or_default(), pixels, image_quality.as_deref(), default_price, @@ -382,8 +701,8 @@ fn resolve_image_output_price_resolution( } } -fn pricing_has_image_output_matrix(pricing: &BillingModelPricingSnapshot) -> bool { - let Some(config) = pricing.effective_tiered_pricing() else { +fn pricing_has_image_output_matrix(pricing: Option<&Value>) -> bool { + let Some(config) = pricing else { return false; }; [ @@ -400,13 +719,12 @@ fn pricing_has_image_output_matrix(pricing: &BillingModelPricingSnapshot) -> boo }) } -fn pricing_has_image_output_ranges(pricing: &BillingModelPricingSnapshot) -> bool { - explicit_image_output_price_ranges(pricing.effective_tiered_pricing()) - .is_some_and(|ranges| !ranges.is_empty()) +fn pricing_has_image_output_ranges(pricing: Option<&Value>) -> bool { + explicit_image_output_price_ranges(pricing).is_some_and(|ranges| !ranges.is_empty()) } -fn pricing_has_image_output_default_price(pricing: &BillingModelPricingSnapshot) -> bool { - let Some(config) = pricing.effective_tiered_pricing() else { +fn pricing_has_image_output_default_price(pricing: Option<&Value>) -> bool { + let Some(config) = pricing else { return false; }; config @@ -586,7 +904,10 @@ mod tests { use serde_json::json; use super::BillingService; - use crate::{BillingModelPricingSnapshot, BillingSnapshotStatus, BillingUsageInput}; + use crate::{ + BillingAuthorizationEstimateInput, BillingModelPricingSnapshot, BillingPricingSource, + BillingSnapshotStatus, BillingUsageInput, + }; fn pricing() -> BillingModelPricingSnapshot { BillingModelPricingSnapshot { @@ -616,6 +937,49 @@ mod tests { } } + fn processing_pricing() -> BillingModelPricingSnapshot { + BillingModelPricingSnapshot { + provider_api_key_rate_multipliers: None, + default_price_per_request: None, + default_tiered_pricing: Some(json!({ + "tiers": [ + {"up_to": 272000, "input_price_per_1m": 5.0, "cache_creation_price_per_1m": 6.25, "cache_read_price_per_1m": 0.5, "output_price_per_1m": 30.0}, + {"up_to": null, "input_price_per_1m": 10.0, "cache_creation_price_per_1m": 12.5, "cache_read_price_per_1m": 1.0, "output_price_per_1m": 45.0} + ], + "processing_tiers": { + "flex": {"tiers": [ + {"up_to": 272000, "input_price_per_1m": 2.5, "cache_creation_price_per_1m": 3.125, "cache_read_price_per_1m": 0.25, "output_price_per_1m": 15.0}, + {"up_to": null, "input_price_per_1m": 5.0, "cache_creation_price_per_1m": 6.25, "cache_read_price_per_1m": 0.5, "output_price_per_1m": 22.5} + ]} + } + })), + model_tiered_pricing: Some(json!({ + "processing_tiers": { + "priority": {"tiers": [ + {"up_to": 272000, "input_price_per_1m": 10.0, "cache_creation_price_per_1m": 12.5, "cache_read_price_per_1m": 1.0, "output_price_per_1m": 60.0} + ]} + } + })), + ..pricing() + } + } + + fn processing_usage( + requested: Option<&str>, + actual: Option<&str>, + input_tokens: i64, + ) -> BillingUsageInput { + BillingUsageInput { + api_format: Some("openai:responses".to_string()), + requested_processing_tier: requested.map(ToOwned::to_owned), + actual_processing_tier: actual.map(ToOwned::to_owned), + input_tokens, + cache_creation_tokens: 10, + cache_ttl_minutes: Some(30), + ..BillingUsageInput::new("chat") + } + } + #[test] fn calculates_complete_snapshot_for_usage() { let result = BillingService::new() @@ -624,6 +988,8 @@ mod tests { &BillingUsageInput { task_type: "chat".to_string(), api_format: Some("openai:chat".to_string()), + requested_processing_tier: None, + actual_processing_tier: None, request_count: 1, input_tokens: 1_000, output_tokens: 500, @@ -654,6 +1020,8 @@ mod tests { &BillingUsageInput { task_type: "chat".to_string(), api_format: Some("openai:responses".to_string()), + requested_processing_tier: None, + actual_processing_tier: None, request_count: 1, input_tokens: 1_000, output_tokens: 10, @@ -696,6 +1064,8 @@ mod tests { &BillingUsageInput { task_type: "chat".to_string(), api_format: Some("openai:responses".to_string()), + requested_processing_tier: None, + actual_processing_tier: None, request_count: 1, input_tokens: 1_000, output_tokens: 10, @@ -730,6 +1100,415 @@ mod tests { assert!(costs.get("cache_read_cost").copied().unwrap_or_default() > 0.0); } + #[test] + fn nonstandard_request_without_actual_tier_fails_closed() { + let result = BillingService::new() + .calculate( + &processing_pricing(), + &processing_usage(Some("priority"), None, 100), + ) + .expect("billing should calculate"); + + assert_eq!(result.cost_result.status, BillingSnapshotStatus::NoRule); + assert_eq!( + result.cost_result.snapshot.missing_required, + vec!["actual_processing_tier"] + ); + assert_eq!( + result.cost_result.snapshot.resolved_dimensions["billing_processing_tier"], + json!(null) + ); + } + + #[test] + fn actual_tier_controls_standard_flex_and_priority_catalogs() { + let cases = [ + ( + "default", + 100, + 5.0, + 6.25, + BillingPricingSource::GlobalDefault, + ), + ("flex", 100, 2.5, 3.125, BillingPricingSource::GlobalDefault), + ( + "priority", + 100, + 10.0, + 12.5, + BillingPricingSource::ProviderOverride, + ), + ]; + + for (actual, input_tokens, input_price, cache_write_price, source) in cases { + let result = BillingService::new() + .calculate( + &processing_pricing(), + &processing_usage(Some("priority"), Some(actual), input_tokens), + ) + .expect("processing tier should resolve"); + + assert_eq!(result.cost_result.status, BillingSnapshotStatus::Complete); + assert_eq!( + result.cost_result.snapshot.resolved_variables["input_price_per_1m"], + json!(input_price) + ); + assert_eq!( + result.cost_result.snapshot.resolved_variables["cache_creation_price_per_1m"], + json!(cache_write_price) + ); + assert_eq!( + result.pricing_resolution.tiered_pricing_source, + Some(source) + ); + } + } + + #[test] + fn finite_processing_catalog_and_unknown_actual_tier_fail_closed() { + let priority = BillingService::new() + .calculate( + &processing_pricing(), + &processing_usage(Some("priority"), Some("priority"), 300_000), + ) + .expect("billing should calculate"); + assert_eq!(priority.cost_result.status, BillingSnapshotStatus::NoRule); + assert_eq!( + priority.cost_result.snapshot.missing_required, + vec!["input_context_tier"] + ); + + let unknown = BillingService::new() + .calculate( + &processing_pricing(), + &processing_usage(Some("priority"), Some("expedited"), 100), + ) + .expect("billing should calculate"); + assert_eq!(unknown.cost_result.status, BillingSnapshotStatus::NoRule); + assert_eq!( + unknown.cost_result.snapshot.missing_required, + vec!["processing_tier_catalog"] + ); + } + + #[test] + fn processing_catalog_boundaries_match_context_and_priority_contracts() { + let cases = [ + ( + "default", + 272_000, + BillingSnapshotStatus::Complete, + Some(5.0), + ), + ( + "default", + 272_001, + BillingSnapshotStatus::Complete, + Some(10.0), + ), + ("flex", 272_000, BillingSnapshotStatus::Complete, Some(2.5)), + ("flex", 272_001, BillingSnapshotStatus::Complete, Some(5.0)), + ( + "priority", + 272_000, + BillingSnapshotStatus::Complete, + Some(10.0), + ), + ("priority", 272_001, BillingSnapshotStatus::NoRule, None), + ]; + + for (actual, input_tokens, status, input_price) in cases { + let result = BillingService::new() + .calculate( + &processing_pricing(), + &processing_usage(Some(actual), Some(actual), input_tokens), + ) + .expect("processing boundary should resolve"); + assert_eq!( + result.cost_result.status, status, + "{actual} at {input_tokens}" + ); + if let Some(input_price) = input_price { + assert_eq!( + result.cost_result.snapshot.resolved_variables["input_price_per_1m"], + json!(input_price), + "{actual} at {input_tokens}" + ); + } else { + assert_eq!( + result.cost_result.snapshot.missing_required, + vec!["input_context_tier"] + ); + } + } + } + + #[test] + fn authorization_estimate_uses_known_request_cache_ttl() { + let pricing = BillingModelPricingSnapshot { + provider_api_key_rate_multipliers: None, + default_price_per_request: None, + default_tiered_pricing: Some(json!({ + "tiers": [{ + "up_to": null, + "input_price_per_1m": 1.0, + "output_price_per_1m": 0.0, + "cache_creation_price_per_1m": 1.25, + "cache_read_price_per_1m": 0.1, + "cache_ttl_pricing": [{ + "ttl_minutes": 60, + "cache_creation_price_per_1m": 100.0, + "cache_read_price_per_1m": 100.0 + }] + }] + })), + ..pricing() + }; + let service = BillingService::new(); + let mut estimate = BillingAuthorizationEstimateInput::new("chat", 1_000_000); + estimate.api_format = Some("openai:responses".to_string()); + estimate.max_output_tokens = Some(0); + estimate.cache_ttl_minutes = Some(30); + + assert_eq!( + service + .estimate_authorization_cost_upper_bound(&pricing, &estimate) + .expect("known TTL estimate should calculate"), + Some(1.25) + ); + + estimate.cache_ttl_minutes = None; + assert_eq!( + service + .estimate_authorization_cost_upper_bound(&pricing, &estimate) + .expect("unknown TTL estimate should calculate"), + Some(100.0) + ); + } + + #[test] + fn authorization_estimate_uses_only_processing_catalogs_eligible_for_context() { + let service = BillingService::new(); + let mut estimate = BillingAuthorizationEstimateInput::new("chat", 300_000); + estimate.api_format = Some("openai:responses".to_string()); + estimate.max_output_tokens = Some(0); + estimate.cache_ttl_minutes = Some(30); + + for requested_processing_tier in [None, Some("standard"), Some("flex")] { + estimate.requested_processing_tier = requested_processing_tier.map(ToOwned::to_owned); + assert_eq!( + service + .estimate_authorization_cost_upper_bound(&processing_pricing(), &estimate) + .expect("eligible processing catalogs should calculate"), + Some(3.75), + "requested tier: {requested_processing_tier:?}" + ); + } + + estimate.requested_processing_tier = Some("priority".to_string()); + assert_eq!( + service + .estimate_authorization_cost_upper_bound(&processing_pricing(), &estimate) + .expect("ineligible requested catalog should resolve"), + None + ); + } + + #[test] + fn unknown_actual_tier_cannot_fall_back_to_fixed_request_price() { + let pricing = BillingModelPricingSnapshot { + default_price_per_request: Some(0.02), + ..processing_pricing() + }; + let result = BillingService::new() + .calculate( + &pricing, + &processing_usage(Some("priority"), Some("expedited"), 100), + ) + .expect("billing should calculate"); + + assert_eq!(result.cost_result.status, BillingSnapshotStatus::NoRule); + assert_eq!( + result.cost_result.snapshot.missing_required, + vec!["processing_tier_catalog"] + ); + assert_eq!(result.pricing_resolution.price_per_request, Some(0.02)); + } + + #[test] + fn authorization_estimate_bounds_requested_and_provider_actual_catalogs() { + let service = BillingService::new(); + let mut estimate = BillingAuthorizationEstimateInput::new("chat", 100_000); + estimate.api_format = Some("openai:responses".to_string()); + estimate.max_output_tokens = Some(1_000_000); + + estimate.requested_processing_tier = Some("priority".to_string()); + let priority = service + .estimate_authorization_cost_upper_bound(&processing_pricing(), &estimate) + .expect("priority estimate should calculate") + .expect("priority estimate should be bounded"); + + estimate.requested_processing_tier = Some("flex".to_string()); + let flex = service + .estimate_authorization_cost_upper_bound(&processing_pricing(), &estimate) + .expect("flex estimate should calculate") + .expect("flex estimate should be bounded"); + + assert_eq!(priority, 61.25); + assert_eq!(flex, 61.25); + assert_eq!(priority, flex); + } + + #[test] + fn authorization_estimate_returns_none_when_the_bound_cannot_be_proven() { + let service = BillingService::new(); + let mut estimate = BillingAuthorizationEstimateInput::new("chat", 100); + estimate.api_format = Some("openai:responses".to_string()); + estimate.requested_processing_tier = Some("priority".to_string()); + + assert_eq!( + service + .estimate_authorization_cost_upper_bound(&processing_pricing(), &estimate) + .expect("unbounded output estimate should resolve"), + None + ); + + estimate.max_output_tokens = Some(10); + estimate.requested_processing_tier = Some("expedited".to_string()); + assert_eq!( + service + .estimate_authorization_cost_upper_bound(&processing_pricing(), &estimate) + .expect("unknown tier estimate should resolve"), + None + ); + + estimate.requested_processing_tier = Some("priority".to_string()); + estimate.input_tokens = 300_000; + assert_eq!( + service + .estimate_authorization_cost_upper_bound(&processing_pricing(), &estimate) + .expect("finite catalog estimate should resolve"), + None + ); + } + + #[test] + fn authorization_estimate_supports_standard_fixed_price_and_free_tier() { + let service = BillingService::new(); + let estimate = BillingAuthorizationEstimateInput::new("chat", 1_000); + let fixed_pricing = BillingModelPricingSnapshot { + default_tiered_pricing: None, + default_price_per_request: Some(0.02), + provider_api_key_rate_multipliers: None, + ..pricing() + }; + assert_eq!( + service + .estimate_authorization_cost_upper_bound(&fixed_pricing, &estimate) + .expect("fixed estimate should calculate"), + Some(0.02) + ); + + let free_pricing = BillingModelPricingSnapshot { + provider_billing_type: Some("free_tier".to_string()), + ..processing_pricing() + }; + assert_eq!( + service + .estimate_authorization_cost_upper_bound(&free_pricing, &estimate) + .expect("free estimate should calculate"), + Some(0.0) + ); + } + + #[test] + fn authorization_estimate_checks_every_reachable_non_monotonic_price_tier() { + let pricing = BillingModelPricingSnapshot { + default_price_per_request: None, + default_tiered_pricing: Some(json!({ + "tiers": [ + { + "up_to": 10_000, + "input_price_per_1m": 100.0, + "output_price_per_1m": 100.0, + "cache_creation_price_per_1m": 100.0, + "cache_read_price_per_1m": 100.0 + }, + { + "up_to": null, + "input_price_per_1m": 1.0, + "output_price_per_1m": 1.0, + "cache_creation_price_per_1m": 1.0, + "cache_read_price_per_1m": 1.0 + } + ] + })), + model_tiered_pricing: None, + ..pricing() + }; + let mut estimate = BillingAuthorizationEstimateInput::new("chat", 100_000); + estimate.api_format = Some("openai:responses".to_string()); + estimate.max_output_tokens = Some(0); + + assert_eq!( + BillingService::new() + .estimate_authorization_cost_upper_bound(&pricing, &estimate) + .expect("non-monotonic catalog should calculate"), + Some(10.0) + ); + } + + #[test] + fn authorization_estimate_uses_api_key_cache_read_ttl_price() { + let pricing = BillingModelPricingSnapshot { + provider_api_key_cache_ttl_minutes: Some(60), + default_price_per_request: None, + default_tiered_pricing: Some(json!({ + "tiers": [{ + "up_to": null, + "input_price_per_1m": 1.0, + "output_price_per_1m": 1.0, + "cache_creation_price_per_1m": 1.25, + "cache_read_price_per_1m": 0.1, + "cache_ttl_pricing": [{ + "ttl_minutes": 60, + "cache_creation_price_per_1m": 1.25, + "cache_read_price_per_1m": 200.0 + }] + }] + })), + model_tiered_pricing: None, + ..pricing() + }; + let mut estimate = BillingAuthorizationEstimateInput::new("chat", 100_000); + estimate.api_format = Some("openai:responses".to_string()); + estimate.max_output_tokens = Some(0); + + assert_eq!( + BillingService::new() + .estimate_authorization_cost_upper_bound(&pricing, &estimate) + .expect("cache read TTL catalog should calculate"), + Some(20.0) + ); + } + + #[test] + fn fixed_request_pricing_remains_independent_for_standard_usage() { + let pricing = BillingModelPricingSnapshot { + default_tiered_pricing: None, + default_price_per_request: Some(0.02), + ..pricing() + }; + let result = BillingService::new() + .calculate(&pricing, &processing_usage(None, None, 1_000)) + .expect("fixed request pricing should calculate"); + + assert_eq!(result.cost_result.status, BillingSnapshotStatus::Complete); + assert_eq!(result.cost_result.cost, 0.02); + assert_eq!(result.pricing_resolution.tiered_pricing, None); + assert_eq!(result.pricing_resolution.price_per_request, Some(0.02)); + } + #[test] fn image_token_usage_without_image_output_price_bills_tokens_only() { let pricing = BillingModelPricingSnapshot { @@ -750,6 +1529,8 @@ mod tests { &BillingUsageInput { task_type: "image".to_string(), api_format: Some("openai:image".to_string()), + requested_processing_tier: None, + actual_processing_tier: None, request_count: 1, input_tokens: 1_000, output_tokens: 20_000, @@ -814,6 +1595,8 @@ mod tests { &BillingUsageInput { task_type: "image".to_string(), api_format: Some("openai:image".to_string()), + requested_processing_tier: None, + actual_processing_tier: None, request_count: 1, input_tokens: 1_000, output_tokens: 20_000, @@ -865,6 +1648,8 @@ mod tests { &BillingUsageInput { task_type: "image".to_string(), api_format: Some("openai:image".to_string()), + requested_processing_tier: None, + actual_processing_tier: None, request_count: 1, input_tokens: 0, output_tokens: 0, @@ -912,6 +1697,8 @@ mod tests { &BillingUsageInput { task_type: "image".to_string(), api_format: Some("openai:image".to_string()), + requested_processing_tier: None, + actual_processing_tier: None, request_count: 1, input_tokens: 0, output_tokens: 0, @@ -972,6 +1759,8 @@ mod tests { &BillingUsageInput { task_type: "image".to_string(), api_format: Some("openai:image".to_string()), + requested_processing_tier: None, + actual_processing_tier: None, request_count: 1, input_tokens: 1_000, output_tokens: 20_000, @@ -1037,6 +1826,8 @@ mod tests { &BillingUsageInput { task_type: "image".to_string(), api_format: Some("openai:image".to_string()), + requested_processing_tier: None, + actual_processing_tier: None, request_count: 1, input_tokens: 1_000, output_tokens: 20_000, @@ -1135,6 +1926,8 @@ mod tests { &BillingUsageInput { task_type: "chat".to_string(), api_format: None, + requested_processing_tier: None, + actual_processing_tier: None, request_count: 1, input_tokens: 1_000, output_tokens: 10, @@ -1208,6 +2001,8 @@ mod tests { &BillingUsageInput { task_type: "chat".to_string(), api_format: None, + requested_processing_tier: None, + actual_processing_tier: None, request_count: 1, input_tokens: 1_000, output_tokens: 10, diff --git a/crates/aether-contracts/src/lib.rs b/crates/aether-contracts/src/lib.rs index a84b35e04..c450eac37 100644 --- a/crates/aether-contracts/src/lib.rs +++ b/crates/aether-contracts/src/lib.rs @@ -11,7 +11,9 @@ 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, TRANSPORT_BACKEND_BROWSER_WREQ, + EXECUTION_REQUEST_HTTP1_ONLY_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, TRANSPORT_HTTP_MODE_H2C_PRIOR_KNOWLEDGE, TRANSPORT_HTTP_MODE_HTTP1_ONLY, TRANSPORT_POOL_SCOPE_KEY, diff --git a/crates/aether-contracts/src/plan.rs b/crates/aether-contracts/src/plan.rs index 2e58736a6..3609fb561 100644 --- a/crates/aether-contracts/src/plan.rs +++ b/crates/aether-contracts/src/plan.rs @@ -7,6 +7,11 @@ 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 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, Default, Serialize, Deserialize, PartialEq, Eq)] #[serde(default)] @@ -115,6 +120,7 @@ pub struct ExecutionPlan { #[serde(default, skip_serializing_if = "Option::is_none")] pub content_encoding: Option, pub body: RequestBody, + /// Whether the upstream API uses a streaming response protocol. #[serde(default)] pub stream: bool, pub client_api_format: String, diff --git a/crates/aether-contracts/src/tunnel.rs b/crates/aether-contracts/src/tunnel.rs index 0e8033432..97528adbc 100644 --- a/crates/aether-contracts/src/tunnel.rs +++ b/crates/aether-contracts/src/tunnel.rs @@ -12,6 +12,7 @@ pub const TUNNEL_PROTOCOL_VERSION_HEADER: &str = "x-aether-tunnel-protocol-versi pub const TUNNEL_NODE_NAME_B64_HEADER: &str = "x-aether-tunnel-node-name-b64"; pub const CURRENT_TUNNEL_PROTOCOL_VERSION: u8 = 3; pub const CURRENT_TUNNEL_PROTOCOL_VERSION_STR: &str = "3"; +pub const MAX_TUNNEL_RELAY_META_LEN: usize = 256 * 1024; pub mod flags { pub const END_STREAM: u8 = 0x01; @@ -220,6 +221,63 @@ pub struct RequestMeta { pub transport_profile: Option, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ResolvedTunnelRequestTimeouts { + pub first_byte_ms: u64, + pub response_body_ms: Option, +} + +pub fn resolve_tunnel_request_timeouts(meta: &RequestMeta) -> ResolvedTunnelRequestTimeouts { + let legacy_timeout_ms = meta.timeout.saturating_mul(1_000); + let first_byte_ms = if meta.stream { + meta.stream_first_byte_timeout_ms + .unwrap_or(legacy_timeout_ms) + } else { + meta.request_timeout_ms + .or(meta.stream_first_byte_timeout_ms) + .unwrap_or(legacy_timeout_ms) + }; + let response_body_ms = (!meta.stream).then_some(first_byte_ms); + + ResolvedTunnelRequestTimeouts { + first_byte_ms: if meta.stream { + clamp_stream_first_byte_timeout_ms(first_byte_ms) + } else { + clamp_upstream_request_timeout_ms(first_byte_ms) + }, + response_body_ms: response_body_ms.map(clamp_upstream_request_timeout_ms), + } +} + +pub fn try_decode_tunnel_relay_request_meta( + buffer: &[u8], +) -> Result, String> { + if buffer.len() < 4 { + return Ok(None); + } + let meta_len = u32::from_be_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]) as usize; + if meta_len > MAX_TUNNEL_RELAY_META_LEN { + return Err("relay metadata too large".to_string()); + } + let meta_end = 4usize + .checked_add(meta_len) + .ok_or_else(|| "relay envelope length overflow".to_string())?; + if buffer.len() < meta_end { + return Ok(None); + } + let meta = serde_json::from_slice::(&buffer[4..meta_end]) + .map_err(|error| format!("invalid relay metadata: {error}"))?; + Ok(Some((meta, meta_end))) +} + +fn clamp_upstream_request_timeout_ms(timeout_ms: u64) -> u64 { + timeout_ms.clamp(1, crate::MAX_EXECUTION_REQUEST_TIMEOUT_MS) +} + +fn clamp_stream_first_byte_timeout_ms(timeout_ms: u64) -> u64 { + timeout_ms.clamp(1, crate::MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_MS) +} + fn default_timeout() -> u64 { 60 } @@ -461,13 +519,115 @@ fn compress_gzip(data: &[u8]) -> Result { mod tests { use super::{ compress_payload, decode_payload, encode_frame, encode_goaway_v3, encode_ping, - encode_reset_stream, encode_window_update, raw_payload, Frame, FrameHeader, GoAwayPayload, - MsgType, RequestMeta, ResetStreamPayload, WindowUpdatePayload, - CURRENT_TUNNEL_PROTOCOL_VERSION, CURRENT_TUNNEL_PROTOCOL_VERSION_STR, FLAG_GZIP_COMPRESSED, + encode_reset_stream, encode_window_update, raw_payload, resolve_tunnel_request_timeouts, + try_decode_tunnel_relay_request_meta, Frame, FrameHeader, GoAwayPayload, MsgType, + RequestMeta, ResetStreamPayload, WindowUpdatePayload, CURRENT_TUNNEL_PROTOCOL_VERSION, + CURRENT_TUNNEL_PROTOCOL_VERSION_STR, FLAG_GZIP_COMPRESSED, MAX_TUNNEL_RELAY_META_LEN, REQUEST_HEADERS, TUNNEL_PROTOCOL_VERSION_HEADER, }; use bytes::Bytes; + fn request_meta(stream: bool) -> RequestMeta { + RequestMeta { + provider_id: None, + endpoint_id: None, + key_id: None, + method: "POST".to_string(), + url: "https://example.com/responses".to_string(), + headers: std::collections::HashMap::new(), + stream, + request_timeout_ms: None, + stream_first_byte_timeout_ms: None, + timeout: 60, + follow_redirects: None, + http1_only: false, + transport_profile: None, + } + } + + #[test] + fn tunnel_request_timeouts_preserve_non_stream_total_timeout() { + let mut meta = request_meta(false); + meta.request_timeout_ms = Some(900_000); + meta.stream_first_byte_timeout_ms = Some(12_000); + + let resolved = resolve_tunnel_request_timeouts(&meta); + + assert_eq!(resolved.first_byte_ms, 900_000); + assert_eq!(resolved.response_body_ms, Some(900_000)); + } + + #[test] + fn tunnel_request_timeouts_keep_stream_body_unbounded() { + let mut meta = request_meta(true); + meta.request_timeout_ms = Some(900_000); + meta.stream_first_byte_timeout_ms = Some(12_000); + + let resolved = resolve_tunnel_request_timeouts(&meta); + + assert_eq!(resolved.first_byte_ms, 12_000); + assert_eq!(resolved.response_body_ms, None); + } + + #[test] + fn tunnel_request_timeouts_keep_stream_first_byte_protocol_limit() { + let mut meta = request_meta(true); + meta.stream_first_byte_timeout_ms = + Some(crate::MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_MS + 1); + + let resolved = resolve_tunnel_request_timeouts(&meta); + + assert_eq!( + resolved.first_byte_ms, + crate::MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_MS + ); + assert_eq!(resolved.response_body_ms, None); + } + + #[test] + fn tunnel_request_timeouts_clamp_only_out_of_range_protocol_values() { + let mut meta = request_meta(false); + meta.request_timeout_ms = Some(crate::MAX_EXECUTION_REQUEST_TIMEOUT_MS + 1); + + let resolved = resolve_tunnel_request_timeouts(&meta); + + assert_eq!( + resolved.first_byte_ms, + crate::MAX_EXECUTION_REQUEST_TIMEOUT_MS + ); + assert_eq!( + resolved.response_body_ms, + Some(crate::MAX_EXECUTION_REQUEST_TIMEOUT_MS) + ); + } + + #[test] + fn tunnel_relay_request_meta_decodes_from_a_partial_prefix() { + let meta = request_meta(false); + let encoded_meta = serde_json::to_vec(&meta).expect("meta should encode"); + let mut envelope = Vec::new(); + envelope.extend_from_slice(&(encoded_meta.len() as u32).to_be_bytes()); + envelope.extend_from_slice(&encoded_meta); + envelope.extend_from_slice(b"request-body"); + + assert!(try_decode_tunnel_relay_request_meta(&envelope[..3]) + .expect("partial prefix should be valid") + .is_none()); + let (decoded, body_offset) = try_decode_tunnel_relay_request_meta(&envelope) + .expect("envelope should be valid") + .expect("metadata should be complete"); + + assert_eq!(decoded.request_timeout_ms, meta.request_timeout_ms); + assert_eq!(&envelope[body_offset..], b"request-body"); + } + + #[test] + fn tunnel_relay_request_meta_rejects_oversized_prefix() { + let oversized = (MAX_TUNNEL_RELAY_META_LEN as u32 + 1).to_be_bytes(); + + assert!(try_decode_tunnel_relay_request_meta(&oversized).is_err()); + } + #[test] fn request_meta_accepts_integer_timeout() { let raw = br#"{"method":"GET","url":"https://example.com","headers":{},"timeout":15}"#; diff --git a/crates/aether-contracts/src/usage.rs b/crates/aether-contracts/src/usage.rs index c1bf82f54..8ca5e730a 100644 --- a/crates/aether-contracts/src/usage.rs +++ b/crates/aether-contracts/src/usage.rs @@ -132,6 +132,8 @@ pub struct ExecutionStreamTerminalSummary { pub response_id: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub model: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub provider_actual_service_tier: Option, #[serde(default)] pub observed_finish: bool, #[serde(default, skip_serializing_if = "is_zero_u64")] diff --git a/crates/aether-data-contracts/src/repository/candidate_selection/types.rs b/crates/aether-data-contracts/src/repository/candidate_selection/types.rs index cc153ffda..309773285 100644 --- a/crates/aether-data-contracts/src/repository/candidate_selection/types.rs +++ b/crates/aether-data-contracts/src/repository/candidate_selection/types.rs @@ -97,13 +97,13 @@ impl StoredMinimalCandidateSelectionRow { None => true, Some(formats) => formats .iter() - .any(|value| api_format_matches(value, api_format)), + .any(|value| api_format_permission_covers(value, api_format)), } } } -fn api_format_matches(left: &str, right: &str) -> bool { - aether_ai_formats::api_format_alias_matches(left, right) +fn api_format_permission_covers(allowed: &str, requested: &str) -> bool { + aether_ai_formats::api_format_permission_covers(allowed, requested) } #[async_trait] 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 13e237216..8143b9880 100644 --- a/crates/aether-data-contracts/src/repository/provider_catalog/mod.rs +++ b/crates/aether-data-contracts/src/repository/provider_catalog/mod.rs @@ -2,7 +2,8 @@ mod types; pub use types::{ ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, ProviderCatalogReadRepository, - ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey, + 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 d89ad58d0..701aaae5d 100644 --- a/crates/aether-data-contracts/src/repository/provider_catalog/types.rs +++ b/crates/aether-data-contracts/src/repository/provider_catalog/types.rs @@ -1,5 +1,11 @@ use async_trait::async_trait; +#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] +pub struct ProviderCatalogUpstreamMetadataNamespaceUpdate { + pub namespace: String, + pub value: serde_json::Value, +} + #[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] pub struct StoredProviderCatalogProvider { pub id: String, @@ -685,6 +691,32 @@ pub trait ProviderCatalogWriteRepository: Send + Sync { updated_at_unix_secs: Option, ) -> Result; + async fn upsert_key_upstream_metadata_namespace( + &self, + key_id: &str, + namespace: &str, + value: &serde_json::Value, + updated_at_unix_secs: Option, + ) -> Result; + + async fn update_key_model_fetch_state( + &self, + key_id: &str, + allowed_models: Option<&serde_json::Value>, + last_models_fetch_at_unix_secs: Option, + last_models_fetch_error: Option<&str>, + updated_at_unix_secs: Option, + ) -> Result; + + async fn update_key_model_fetch_success( + &self, + key_id: &str, + allowed_models: Option<&serde_json::Value>, + last_models_fetch_at_unix_secs: u64, + upstream_metadata_updates: &[ProviderCatalogUpstreamMetadataNamespaceUpdate], + updated_at_unix_secs: Option, + ) -> Result; + async fn delete_key(&self, key_id: &str) -> Result; async fn clear_key_oauth_invalid_marker( diff --git a/crates/aether-data-contracts/src/repository/usage/mod.rs b/crates/aether-data-contracts/src/repository/usage/mod.rs index ff009a63b..f20a0b13e 100644 --- a/crates/aether-data-contracts/src/repository/usage/mod.rs +++ b/crates/aether-data-contracts/src/repository/usage/mod.rs @@ -1,8 +1,10 @@ mod types; pub use types::{ - extract_provider_reasoning_effort_from_body, extract_provider_service_tier_from_body, - parse_usage_body_ref, usage_body_ref, usage_request_metadata_client_family, + extract_provider_actual_service_tier_from_response, + extract_provider_cache_ttl_minutes_from_metadata, extract_provider_reasoning_effort_from_body, + extract_provider_service_tier_from_body, normalize_provider_service_tier, parse_usage_body_ref, + resolve_provider_cache_ttl_minutes, usage_body_ref, usage_request_metadata_client_family, ApiKeyLastUsedDelta, ManagementTokenCounterDelta, PendingUsageCleanupSummary, ProviderApiKeyWindowUsageRequest, ProxyNodeCounterDelta, StoredProviderApiKeyUsageSummary, StoredProviderApiKeyWindowUsageSummary, StoredProviderUsageSummary, StoredProviderUsageWindow, @@ -28,6 +30,7 @@ pub use types::{ UsageLeaderboardQuery, UsageMonitoringErrorCountQuery, UsageMonitoringErrorListQuery, UsagePerformancePercentilesQuery, UsageProviderPerformanceQuery, UsageReadRepository, UsageRepository, UsageSettledCostSummaryQuery, UsageTimeSeriesGranularity, - UsageTimeSeriesQuery, UsageWriteRepository, PROVIDER_REASONING_EFFORT_METADATA_KEY, + UsageTimeSeriesQuery, UsageWriteRepository, PROVIDER_ACTUAL_SERVICE_TIER_METADATA_KEY, + PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY, PROVIDER_REASONING_EFFORT_METADATA_KEY, PROVIDER_SERVICE_TIER_METADATA_KEY, }; diff --git a/crates/aether-data-contracts/src/repository/usage/types.rs b/crates/aether-data-contracts/src/repository/usage/types.rs index 644deb470..69d732283 100644 --- a/crates/aether-data-contracts/src/repository/usage/types.rs +++ b/crates/aether-data-contracts/src/repository/usage/types.rs @@ -4,6 +4,8 @@ use serde_json::Value; pub const PROVIDER_REASONING_EFFORT_METADATA_KEY: &str = "provider_reasoning_effort"; pub const PROVIDER_SERVICE_TIER_METADATA_KEY: &str = "provider_service_tier"; +pub const PROVIDER_ACTUAL_SERVICE_TIER_METADATA_KEY: &str = "provider_actual_service_tier"; +pub const PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY: &str = "provider_cache_ttl_minutes"; pub fn extract_provider_reasoning_effort_from_body(value: Option<&Value>) -> Option { let object = value.and_then(Value::as_object)?; @@ -43,7 +45,26 @@ pub fn extract_provider_service_tier_from_body(value: Option<&Value>) -> Option< .and_then(normalize_provider_service_tier) } -fn normalize_provider_service_tier(value: &str) -> Option { +pub fn extract_provider_actual_service_tier_from_response(value: Option<&Value>) -> Option { + let value = value?; + value + .get("chunks") + .and_then(Value::as_array) + .and_then(|chunks| { + chunks + .iter() + .rev() + .find_map(|chunk| extract_provider_actual_service_tier_from_response(Some(chunk))) + }) + .or_else(|| { + value + .get("response") + .and_then(|response| extract_provider_service_tier_from_body(Some(response))) + }) + .or_else(|| extract_provider_service_tier_from_body(Some(value))) +} + +pub fn normalize_provider_service_tier(value: &str) -> Option { let normalized = value.trim().to_ascii_lowercase(); if normalized.is_empty() || normalized.len() > 64 { return None; @@ -51,6 +72,48 @@ fn normalize_provider_service_tier(value: &str) -> Option { Some(normalized) } +pub fn resolve_provider_cache_ttl_minutes( + provider_api_format: Option<&str>, + provider_model: Option<&str>, + source_model: Option<&str>, + provider_request_body: Option<&Value>, +) -> Option { + let provider_api_format = provider_api_format?.trim(); + let provider_request_body = provider_request_body?; + let provider_model = provider_request_body + .get("model") + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .or_else(|| { + provider_model + .map(str::trim) + .filter(|value| !value.is_empty()) + })?; + let source_model = source_model + .map(str::trim) + .filter(|value| !value.is_empty()) + .unwrap_or(provider_model); + aether_ai_formats::resolve_openai_prompt_cache_ttl_minutes( + provider_api_format, + provider_model, + source_model, + provider_request_body, + ) +} + +pub fn extract_provider_cache_ttl_minutes_from_metadata(value: Option<&Value>) -> Option { + value + .and_then(Value::as_object) + .and_then(|metadata| metadata.get(PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY)) + .and_then(|value| { + value + .as_i64() + .or_else(|| value.as_u64().and_then(|value| i64::try_from(value).ok())) + }) + .filter(|value| *value > 0) +} + /// Joined usage read model assembled from the accounting row plus the newer audit/snapshot /// satellite tables. /// @@ -454,6 +517,28 @@ impl StoredRequestUsageAudit { .and_then(normalize_provider_service_tier) } + pub fn provider_actual_service_tier(&self) -> Option { + self.request_metadata_string(PROVIDER_ACTUAL_SERVICE_TIER_METADATA_KEY) + .and_then(normalize_provider_service_tier) + .or_else(|| { + extract_provider_actual_service_tier_from_response(self.response_body.as_ref()) + }) + } + + pub fn provider_cache_ttl_minutes(&self) -> Option { + resolve_provider_cache_ttl_minutes( + self.endpoint_api_format + .as_deref() + .or(self.api_format.as_deref()), + self.target_model.as_deref().or(Some(self.model.as_str())), + Some(self.model.as_str()), + self.provider_request_body.as_ref(), + ) + .or_else(|| { + extract_provider_cache_ttl_minutes_from_metadata(self.request_metadata.as_ref()) + }) + } + pub fn body_ref(&self, field: UsageBodyField) -> Option<&str> { match field { UsageBodyField::RequestBody => self.request_body_ref.as_deref(), @@ -2069,6 +2154,7 @@ fn parse_timestamp(value: i64, field_name: &str) -> Result, + ) = sqlx::query_as( + r#" +SELECT + allowed_models, + upstream_metadata, + EXTRACT(EPOCH FROM last_models_fetch_at)::bigint, + last_models_fetch_error +FROM public.provider_api_keys +WHERE id = 'metadata-migration-key' +"#, + ) + .fetch_one(&pool) + .await + .expect("provider model fetch state should load"); + assert_eq!(allowed_models, serde_json::json!(["gpt-5.6-sol"])); + assert_eq!(fetched_at, 1_740_000_002); + assert_eq!(fetch_error, None); + assert_eq!( + metadata, + serde_json::json!({ + "admin": {"source": "manual"}, + "catalog": { + "model": "gpt-5.6", + "nested": [1, true, null] + }, + "codex_models": { + "cards": { + "gpt-5.6-sol": { + "slug": "gpt-5.6-sol", + "use_responses_lite": true + } + } + } + }) + ); + + query("DELETE FROM public.providers WHERE id = 'metadata-migration-provider'") + .execute(&pool) + .await + .expect("provider migration fixture should clean up"); +} + #[tokio::test] async fn mysql_migrations_create_core_config_tables_when_url_is_set() { let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL") diff --git a/crates/aether-data/src/repository/auth/types.rs b/crates/aether-data/src/repository/auth/types.rs index c12e3a1e2..f3b7039b4 100644 --- a/crates/aether-data/src/repository/auth/types.rs +++ b/crates/aether-data/src/repository/auth/types.rs @@ -247,7 +247,7 @@ impl ResolvedAuthApiKeySnapshot { &mut self.user_allowed_providers, &mut self.api_key_allowed_providers, ); - constrain_api_key_list_policy_to_user_policy( + constrain_api_key_api_format_policy_to_user_policy( &mut self.user_allowed_api_formats, &mut self.api_key_allowed_api_formats, ); @@ -277,6 +277,22 @@ fn constrain_api_key_list_policy_to_user_policy( *api_key_policy = Some(effective); } +fn constrain_api_key_api_format_policy_to_user_policy( + user_policy: &mut Option>, + api_key_policy: &mut Option>, +) { + let Some(api_key_values) = api_key_policy.as_ref() else { + return; + }; + let Some(user_values) = user_policy.as_ref() else { + return; + }; + let effective = + aether_ai_formats::intersect_api_format_allowed_lists(api_key_values, user_values); + *user_policy = Some(effective.clone()); + *api_key_policy = Some(effective); +} + fn intersect_allowed_lists(left: &[String], right: &[String]) -> Vec { let right_values = right.iter().collect::>(); left.iter() @@ -810,6 +826,29 @@ mod tests { StoredAuthApiKeySnapshot, }; + #[test] + fn api_format_policy_intersection_preserves_search_companion_scope() { + assert_eq!( + aether_ai_formats::intersect_api_format_allowed_lists( + &["openai:search".to_string()], + &["openai:responses".to_string()], + ), + vec!["openai:search".to_string()] + ); + assert_eq!( + aether_ai_formats::intersect_api_format_allowed_lists( + &["openai:responses".to_string()], + &["openai:search".to_string()], + ), + vec!["openai:search".to_string()] + ); + assert!(aether_ai_formats::intersect_api_format_allowed_lists( + &["openai:search".to_string()], + &["openai:chat".to_string()], + ) + .is_empty()); + } + #[test] fn rejects_non_array_allowed_providers() { assert!(StoredAuthApiKeySnapshot::new( diff --git a/crates/aether-data/src/repository/candidate_selection/memory.rs b/crates/aether-data/src/repository/candidate_selection/memory.rs index 044847f55..4b136bba4 100644 --- a/crates/aether-data/src/repository/candidate_selection/memory.rs +++ b/crates/aether-data/src/repository/candidate_selection/memory.rs @@ -230,7 +230,7 @@ fn row_matches_requested_model( mapping.api_formats.as_ref().is_none_or(|formats| { formats .iter() - .any(|value| api_format_matches(value, api_format)) + .any(|value| api_format_scope_covers(value, api_format)) }) && mapping.endpoint_ids.as_ref().is_none_or(|endpoint_ids| { endpoint_ids .iter() @@ -286,7 +286,7 @@ fn mapping_scope_matches( mapping.api_formats.as_ref().is_none_or(|formats| { formats .iter() - .any(|value| api_format_matches(value, api_format)) + .any(|value| api_format_scope_covers(value, api_format)) }) && mapping.endpoint_ids.as_ref().is_none_or(|endpoint_ids| { endpoint_ids .iter() @@ -294,6 +294,10 @@ fn mapping_scope_matches( }) } +fn api_format_scope_covers(allowed: &str, requested: &str) -> bool { + aether_ai_formats::api_format_permission_covers(allowed, requested) +} + fn key_auth_channel_matches(row: &StoredMinimalCandidateSelectionRow, api_format: &str) -> bool { let provider_type = row.provider_type.trim().to_ascii_lowercase(); let auth_type = row.key_auth_type.trim().to_ascii_lowercase(); @@ -303,7 +307,10 @@ fn key_auth_channel_matches(row: &StoredMinimalCandidateSelectionRow, api_format auth_type == "oauth" && matches!( api_format.as_str(), - "openai:responses" | "openai:responses:compact" | "openai:image" + "openai:responses" + | "openai:responses:compact" + | "openai:search" + | "openai:image" ) } "chatgpt_web" => { @@ -434,6 +441,47 @@ mod tests { assert_eq!(rows[0].provider_id, "provider-1"); } + #[tokio::test] + async fn search_uses_responses_key_and_model_permissions_with_exact_endpoint_identity() { + let mut search = sample_row( + "provider-search", + "openai:search", + "global-search-model", + 10, + ); + search.provider_type = "codex".to_string(); + search.key_auth_type = "oauth".to_string(); + search.key_api_formats = Some(vec!["openai:responses".to_string()]); + search.model_provider_model_name = "upstream-search-model".to_string(); + search.model_provider_model_mappings = Some(vec![StoredProviderModelMapping { + name: "gpt-5.6-sol".to_string(), + priority: 0, + api_formats: Some(vec!["openai:responses".to_string()]), + endpoint_ids: None, + }]); + + let mut responses = search.clone(); + responses.provider_id = "provider-responses".to_string(); + responses.endpoint_id = "endpoint-responses".to_string(); + responses.endpoint_api_format = "openai:responses".to_string(); + responses.key_id = "key-responses".to_string(); + responses.model_id = "model-responses".to_string(); + + let repository = + InMemoryMinimalCandidateSelectionReadRepository::seed(vec![responses, search]); + let rows = repository + .list_for_exact_api_format_and_requested_model("openai:search", "gpt-5.6-sol") + .await + .expect("Search candidate should load"); + + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].endpoint_api_format, "openai:search"); + assert_eq!( + rows[0].key_api_formats, + Some(vec!["openai:responses".to_string()]) + ); + } + #[tokio::test] async fn includes_grok_oauth_rows_for_chat_models() { let mut row = sample_row( diff --git a/crates/aether-data/src/repository/candidate_selection/mysql.rs b/crates/aether-data/src/repository/candidate_selection/mysql.rs index 691bc1ce4..9ae01cb74 100644 --- a/crates/aether-data/src/repository/candidate_selection/mysql.rs +++ b/crates/aether-data/src/repository/candidate_selection/mysql.rs @@ -409,7 +409,7 @@ fn mapping_scope_matches( mapping.api_formats.as_ref().is_none_or(|formats| { formats .iter() - .any(|value| api_format_matches(value, api_format)) + .any(|value| api_format_scope_covers(value, api_format)) }) && mapping.endpoint_ids.as_ref().is_none_or(|endpoint_ids| { endpoint_ids .iter() @@ -426,7 +426,10 @@ fn key_auth_channel_matches(row: &CandidateSelectionRow, api_format: &str) -> bo auth_type == "oauth" && matches!( api_format.as_str(), - "openai:responses" | "openai:responses:compact" | "openai:image" + "openai:responses" + | "openai:responses:compact" + | "openai:search" + | "openai:image" ) } "chatgpt_web" => { @@ -763,6 +766,10 @@ fn api_format_matches(left: &str, right: &str) -> bool { aether_ai_formats::api_format_alias_matches(left, right) } +fn api_format_scope_covers(allowed: &str, requested: &str) -> bool { + aether_ai_formats::api_format_permission_covers(allowed, requested) +} + fn sql_match_aliases(api_formats: &[String]) -> Vec { api_formats .iter() diff --git a/crates/aether-data/src/repository/candidate_selection/postgres.rs b/crates/aether-data/src/repository/candidate_selection/postgres.rs index 227214e16..ccf7d5835 100644 --- a/crates/aether-data/src/repository/candidate_selection/postgres.rs +++ b/crates/aether-data/src/repository/candidate_selection/postgres.rs @@ -71,7 +71,7 @@ INNER JOIN LATERAL ( ( LOWER(BTRIM(p.provider_type)) = 'codex' AND LOWER(BTRIM(pak.auth_type)) = 'oauth' - AND LOWER($3) IN ('openai:responses', 'openai:responses:compact', 'openai:image') + AND LOWER($3) IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image') ) OR ( LOWER(BTRIM(p.provider_type)) = 'chatgpt_web' @@ -164,7 +164,7 @@ WHERE p.is_active = TRUE ( LOWER(BTRIM(p.provider_type)) = 'codex' AND LOWER(BTRIM(pak.auth_type)) = 'oauth' - AND LOWER($3) IN ('openai:responses', 'openai:responses:compact', 'openai:image') + AND LOWER($3) IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image') ) OR ( LOWER(BTRIM(p.provider_type)) = 'chatgpt_web' @@ -350,7 +350,7 @@ INNER JOIN LATERAL ( ( LOWER(BTRIM(p.provider_type)) = 'codex' AND LOWER(BTRIM(pak.auth_type)) = 'oauth' - AND LOWER($4) IN ('openai:responses', 'openai:responses:compact', 'openai:image') + AND LOWER($4) IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image') ) OR ( LOWER(BTRIM(p.provider_type)) = 'chatgpt_web' @@ -444,7 +444,7 @@ WHERE p.is_active = TRUE ( LOWER(BTRIM(p.provider_type)) = 'codex' AND LOWER(BTRIM(pak.auth_type)) = 'oauth' - AND LOWER($4) IN ('openai:responses', 'openai:responses:compact', 'openai:image') + AND LOWER($4) IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image') ) OR ( LOWER(BTRIM(p.provider_type)) = 'chatgpt_web' @@ -638,7 +638,7 @@ WHERE p.is_active = TRUE ( LOWER(BTRIM(p.provider_type)) = 'codex' AND LOWER(BTRIM(pak.auth_type)) = 'oauth' - AND LOWER($6) IN ('openai:responses', 'openai:responses:compact', 'openai:image') + AND LOWER($6) IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image') ) OR ( LOWER(BTRIM(p.provider_type)) = 'chatgpt_web' @@ -784,7 +784,8 @@ impl SqlxMinimalCandidateSelectionReadRepository { let mut rows = Vec::new(); let canonical_api_format = normalize_api_format(api_format); let storage_aliases = api_format_aliases(&canonical_api_format); - let sql_match_aliases = sql_match_aliases(&storage_aliases); + let sql_match_aliases = + sql_match_aliases(&api_format_permission_aliases(&canonical_api_format)); for api_format in storage_aliases { rows.extend( Self::collect_query_rows( @@ -809,7 +810,8 @@ impl SqlxMinimalCandidateSelectionReadRepository { let mut rows = Vec::new(); let canonical_api_format = normalize_api_format(api_format); let storage_aliases = api_format_aliases(&canonical_api_format); - let sql_match_aliases = sql_match_aliases(&storage_aliases); + let sql_match_aliases = + sql_match_aliases(&api_format_permission_aliases(&canonical_api_format)); for api_format in storage_aliases { rows.extend( Self::collect_query_rows( @@ -835,7 +837,8 @@ impl SqlxMinimalCandidateSelectionReadRepository { let mut rows = Vec::new(); let canonical_api_format = normalize_api_format(api_format); let storage_aliases = api_format_aliases(&canonical_api_format); - let sql_match_aliases = sql_match_aliases(&storage_aliases); + let sql_match_aliases = + sql_match_aliases(&api_format_permission_aliases(&canonical_api_format)); let sql = requested_model_selection_sql(); for api_format in storage_aliases { rows.extend( @@ -861,7 +864,8 @@ impl SqlxMinimalCandidateSelectionReadRepository { let mut rows = Vec::new(); let canonical_api_format = normalize_api_format(&query.api_format); let storage_aliases = api_format_aliases(&canonical_api_format); - let sql_match_aliases = sql_match_aliases(&storage_aliases); + let sql_match_aliases = + sql_match_aliases(&api_format_permission_aliases(&canonical_api_format)); let limit = i64::from(query.limit.max(1)); let offset = i64::from(query.offset); let sql = requested_model_selection_page_sql(); @@ -891,7 +895,8 @@ impl SqlxMinimalCandidateSelectionReadRepository { let mut rows = Vec::new(); let canonical_api_format = normalize_api_format(&query.api_format); let storage_aliases = api_format_aliases(&canonical_api_format); - let sql_match_aliases = sql_match_aliases(&storage_aliases); + let sql_match_aliases = + sql_match_aliases(&api_format_permission_aliases(&canonical_api_format)); let limit = i64::from(query.limit.max(1)); let offset = i64::from(query.offset); let sql = pool_key_candidate_selection_sql(&query.order); @@ -929,7 +934,8 @@ impl SqlxMinimalCandidateSelectionReadRepository { let mut rows = Vec::new(); let canonical_api_format = normalize_api_format(&query.api_format); let storage_aliases = api_format_aliases(&canonical_api_format); - let sql_match_aliases = sql_match_aliases(&storage_aliases); + let sql_match_aliases = + sql_match_aliases(&api_format_permission_aliases(&canonical_api_format)); let sql = pool_key_candidate_selection_by_key_ids_sql(); for api_format in storage_aliases { rows.extend( @@ -1111,6 +1117,10 @@ fn api_format_aliases(api_format: &str) -> Vec { aether_ai_formats::api_format_storage_aliases(api_format) } +fn api_format_permission_aliases(api_format: &str) -> Vec { + aether_ai_formats::api_format_permission_storage_aliases(api_format) +} + fn normalize_api_format(api_format: &str) -> String { aether_ai_formats::normalize_api_format_alias(api_format) } diff --git a/crates/aether-data/src/repository/candidate_selection/sqlite.rs b/crates/aether-data/src/repository/candidate_selection/sqlite.rs index fd9f4f474..4190d56af 100644 --- a/crates/aether-data/src/repository/candidate_selection/sqlite.rs +++ b/crates/aether-data/src/repository/candidate_selection/sqlite.rs @@ -125,7 +125,8 @@ impl SqliteMinimalCandidateSelectionReadRepository { ) -> Result, DataLayerError> { let canonical_api_format = normalize_api_format(api_format); let storage_aliases = api_format_aliases(&canonical_api_format); - let match_aliases = sql_match_aliases(&storage_aliases); + let match_aliases = + sql_match_aliases(&api_format_permission_aliases(&canonical_api_format)); let mut rows = Vec::new(); for storage_api_format in storage_aliases { @@ -270,7 +271,8 @@ impl MinimalCandidateSelectionReadRepository for SqliteMinimalCandidateSelection ) -> Result, DataLayerError> { let canonical_api_format = normalize_api_format(&query.api_format); let storage_aliases = api_format_aliases(&canonical_api_format); - let match_aliases = sql_match_aliases(&storage_aliases); + let match_aliases = + sql_match_aliases(&api_format_permission_aliases(&canonical_api_format)); let mut rows = Vec::::new(); let page_in_sql = !matches!(query.order, StoredPoolKeyCandidateOrder::LoadBalance { .. }); @@ -337,7 +339,8 @@ impl MinimalCandidateSelectionReadRepository for SqliteMinimalCandidateSelection .collect::>(); let canonical_api_format = normalize_api_format(&query.api_format); let storage_aliases = api_format_aliases(&canonical_api_format); - let match_aliases = sql_match_aliases(&storage_aliases); + let match_aliases = + sql_match_aliases(&api_format_permission_aliases(&canonical_api_format)); let mut rows = Vec::new(); for storage_api_format in storage_aliases { @@ -493,7 +496,7 @@ fn push_key_auth_channel_sql_filter( ); builder.push_bind(api_format.clone()); builder.push( - r#" IN ('openai:responses', 'openai:responses:compact', 'openai:image') + r#" IN ('openai:responses', 'openai:responses:compact', 'openai:search', 'openai:image') ) OR ( LOWER(TRIM(p.provider_type)) = 'chatgpt_web' @@ -808,7 +811,7 @@ fn mapping_scope_matches( mapping.api_formats.as_ref().is_none_or(|formats| { formats .iter() - .any(|value| api_format_matches(value, api_format)) + .any(|value| api_format_scope_covers(value, api_format)) }) && mapping.endpoint_ids.as_ref().is_none_or(|endpoint_ids| { endpoint_ids .iter() @@ -825,7 +828,10 @@ fn key_auth_channel_matches(row: &CandidateSelectionRow, api_format: &str) -> bo auth_type == "oauth" && matches!( api_format.as_str(), - "openai:responses" | "openai:responses:compact" | "openai:image" + "openai:responses" + | "openai:responses:compact" + | "openai:search" + | "openai:image" ) } "chatgpt_web" => { @@ -1145,6 +1151,10 @@ fn api_format_aliases(api_format: &str) -> Vec { aether_ai_formats::api_format_storage_aliases(api_format) } +fn api_format_permission_aliases(api_format: &str) -> Vec { + aether_ai_formats::api_format_permission_storage_aliases(api_format) +} + fn normalize_api_format(api_format: &str) -> String { aether_ai_formats::normalize_api_format_alias(api_format) } @@ -1153,6 +1163,10 @@ fn api_format_matches(left: &str, right: &str) -> bool { aether_ai_formats::api_format_alias_matches(left, right) } +fn api_format_scope_covers(allowed: &str, requested: &str) -> bool { + aether_ai_formats::api_format_permission_covers(allowed, requested) +} + fn sql_match_aliases(api_formats: &[String]) -> Vec { api_formats .iter() @@ -1264,6 +1278,21 @@ mod tests { .collect::>(), vec!["key-chatgpt-web-oauth", "key-chatgpt-web-bearer"] ); + + let search_rows = repository + .list_for_exact_api_format_and_requested_model_page( + &StoredRequestedModelCandidateRowsQuery { + api_format: "openai:search".to_string(), + requested_model_name: "gpt-5.6-sol".to_string(), + offset: 0, + limit: 10, + }, + ) + .await + .expect("Codex Search rows should load through Responses permissions"); + assert_eq!(search_rows.len(), 1); + assert_eq!(search_rows[0].key_id, "key-codex-search"); + assert_eq!(search_rows[0].endpoint_api_format, "openai:search"); } async fn seed_candidate_selection(pool: &sqlx::SqlitePool) { @@ -1296,6 +1325,11 @@ INSERT INTO providers ( ) VALUES ('provider-windsurf', 'Windsurf', 'windsurf', 15, 1, 1, 1); +INSERT INTO providers ( + id, name, provider_type, provider_priority, is_active, created_at, updated_at +) +VALUES ('provider-codex-search', 'Codex Search', 'codex', 12, 1, 1, 1); + INSERT INTO provider_endpoints ( id, provider_id, name, base_url, api_format, is_active, created_at, updated_at ) @@ -1312,6 +1346,14 @@ VALUES ( 'https://server.codeium.com', 'openai:chat', 1, 1, 1 ); +INSERT INTO provider_endpoints ( + id, provider_id, name, base_url, api_format, is_active, created_at, updated_at +) +VALUES ( + 'endpoint-codex-search', 'provider-codex-search', 'Codex Search', + 'https://chatgpt.com/backend-api/codex', 'openai:search', 1, 1, 1 +); + INSERT INTO provider_api_keys ( id, provider_id, name, auth_type, api_formats, internal_priority, is_active, created_at, updated_at ) @@ -1327,13 +1369,22 @@ VALUES ( 'key-windsurf-oauth', 'provider-windsurf', 'OAuth', 'oauth', '["openai:chat"]', 10, 1, 1, 1 ); +INSERT INTO provider_api_keys ( + id, provider_id, name, auth_type, api_formats, internal_priority, is_active, created_at, updated_at +) +VALUES ( + 'key-codex-search', 'provider-codex-search', 'OAuth', 'oauth', + '["openai:responses"]', 10, 1, 1, 1 +); + INSERT INTO global_models ( id, name, config, is_active, created_at, updated_at ) VALUES ('global-1', 'gpt-5', '{"model_mappings":["alias-global"],"streaming":true}', 1, 1, 1), ('global-image-1', 'gpt-image-2', NULL, 1, 1, 1), - ('global-windsurf-1', 'claude-opus-4-7', '{"streaming":true}', 1, 1, 1); + ('global-windsurf-1', 'claude-opus-4-7', '{"streaming":true}', 1, 1, 1), + ('global-codex-search-1', 'search-global', '{"streaming":false}', 1, 1, 1); INSERT INTO models ( id, provider_id, global_model_id, provider_model_name, provider_model_mappings, @@ -1351,6 +1402,11 @@ VALUES ( ( 'model-windsurf-opus', 'provider-windsurf', 'global-windsurf-1', 'claude-opus-4-7', NULL, NULL, 1, 1, 1, 1 +), +( + 'model-codex-search', 'provider-codex-search', 'global-codex-search-1', 'search-upstream', + '[{"name":"gpt-5.6-sol","api_formats":["openai:responses"],"priority":1}]', + 0, 1, 1, 1, 1 ); "#, ) diff --git a/crates/aether-data/src/repository/provider_catalog/memory.rs b/crates/aether-data/src/repository/provider_catalog/memory.rs index e4b28bb46..cae53c76d 100644 --- a/crates/aether-data/src/repository/provider_catalog/memory.rs +++ b/crates/aether-data/src/repository/provider_catalog/memory.rs @@ -7,7 +7,8 @@ use serde_json::{json, Map, Value}; use super::{ ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, ProviderCatalogReadRepository, - ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey, + ProviderCatalogUpstreamMetadataNamespaceUpdate, ProviderCatalogWriteRepository, + StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider, }; @@ -748,6 +749,111 @@ impl ProviderCatalogWriteRepository for InMemoryProviderCatalogReadRepository { Ok(true) } + async fn upsert_key_upstream_metadata_namespace( + &self, + key_id: &str, + namespace: &str, + value: &serde_json::Value, + updated_at_unix_secs: Option, + ) -> Result { + if namespace.trim().is_empty() { + return Err(DataLayerError::InvalidInput( + "provider catalog upstream metadata namespace is empty".to_string(), + )); + } + let mut index = self + .index + .write() + .expect("provider catalog repository lock"); + let Some(key) = index.keys.get_mut(key_id) else { + return Ok(false); + }; + let metadata = key + .upstream_metadata + .get_or_insert_with(|| serde_json::json!({})); + let Some(metadata) = metadata.as_object_mut() else { + return Err(DataLayerError::UnexpectedValue( + "provider catalog upstream metadata must be an object".to_string(), + )); + }; + metadata.insert(namespace.to_string(), value.clone()); + key.updated_at_unix_secs = updated_at_unix_secs; + Ok(true) + } + + async fn update_key_model_fetch_state( + &self, + key_id: &str, + allowed_models: Option<&serde_json::Value>, + last_models_fetch_at_unix_secs: Option, + last_models_fetch_error: Option<&str>, + updated_at_unix_secs: Option, + ) -> Result { + let mut index = self + .index + .write() + .expect("provider catalog repository lock"); + let Some(key) = index.keys.get_mut(key_id) else { + return Ok(false); + }; + key.allowed_models = allowed_models.cloned(); + key.last_models_fetch_at_unix_secs = last_models_fetch_at_unix_secs; + key.last_models_fetch_error = last_models_fetch_error.map(str::to_string); + key.updated_at_unix_secs = updated_at_unix_secs; + Ok(true) + } + + async fn update_key_model_fetch_success( + &self, + key_id: &str, + allowed_models: Option<&serde_json::Value>, + last_models_fetch_at_unix_secs: u64, + upstream_metadata_updates: &[ProviderCatalogUpstreamMetadataNamespaceUpdate], + updated_at_unix_secs: Option, + ) -> Result { + if upstream_metadata_updates + .iter() + .any(|update| update.namespace.trim().is_empty()) + { + return Err(DataLayerError::InvalidInput( + "provider catalog upstream metadata namespace is empty".to_string(), + )); + } + let mut index = self + .index + .write() + .expect("provider catalog repository lock"); + let Some(key) = index.keys.get_mut(key_id) else { + return Ok(false); + }; + if !upstream_metadata_updates.is_empty() + && key + .upstream_metadata + .as_ref() + .is_some_and(|metadata| !metadata.is_object()) + { + return Err(DataLayerError::UnexpectedValue( + "provider catalog upstream metadata must be an object".to_string(), + )); + } + + key.allowed_models = allowed_models.cloned(); + key.last_models_fetch_at_unix_secs = Some(last_models_fetch_at_unix_secs); + key.last_models_fetch_error = None; + key.updated_at_unix_secs = updated_at_unix_secs; + if !upstream_metadata_updates.is_empty() { + let metadata = key + .upstream_metadata + .get_or_insert_with(|| serde_json::json!({})) + .as_object_mut() + .expect("upstream metadata object was validated"); + for update in upstream_metadata_updates { + metadata.insert(update.namespace.clone(), update.value.clone()); + } + } + Ok(true) + } + async fn delete_key(&self, key_id: &str) -> Result { let mut index = self .index diff --git a/crates/aether-data/src/repository/provider_catalog/mod.rs b/crates/aether-data/src/repository/provider_catalog/mod.rs index e792049f3..639d9ae9c 100644 --- a/crates/aether-data/src/repository/provider_catalog/mod.rs +++ b/crates/aether-data/src/repository/provider_catalog/mod.rs @@ -6,7 +6,8 @@ mod sqlite; #[allow(unused_imports)] pub(crate) use aether_data_contracts::repository::provider_catalog::{ ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, ProviderCatalogReadRepository, - ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey, + ProviderCatalogUpstreamMetadataNamespaceUpdate, ProviderCatalogWriteRepository, + StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider, }; diff --git a/crates/aether-data/src/repository/provider_catalog/mysql.rs b/crates/aether-data/src/repository/provider_catalog/mysql.rs index 0b4a08b9f..f954cef28 100644 --- a/crates/aether-data/src/repository/provider_catalog/mysql.rs +++ b/crates/aether-data/src/repository/provider_catalog/mysql.rs @@ -3,9 +3,10 @@ use sqlx::{mysql::MySqlRow, MySql, QueryBuilder, Row}; use super::{ InMemoryProviderCatalogReadRepository, ProviderCatalogKeyListQuery, - ProviderCatalogReadRepository, ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, - StoredProviderCatalogKey, StoredProviderCatalogKeyMaintenanceSummary, - StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider, + ProviderCatalogReadRepository, ProviderCatalogUpstreamMetadataNamespaceUpdate, + ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey, + StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage, + StoredProviderCatalogKeyStats, StoredProviderCatalogProvider, }; use crate::driver::mysql::MysqlPool; use crate::error::SqlResultExt; @@ -800,6 +801,164 @@ WHERE id = ? Ok(rows_affected > 0) } + pub async fn upsert_key_upstream_metadata_namespace( + &self, + key_id: &str, + namespace: &str, + value: &serde_json::Value, + updated_at_unix_secs: Option, + ) -> Result { + validate_non_empty(key_id, "provider catalog key_id")?; + validate_non_empty(namespace, "provider catalog upstream metadata namespace")?; + let value_json = serde_json::to_string(value).map_err(|err| { + DataLayerError::UnexpectedValue(format!( + "provider_api_keys.upstream_metadata namespace is not serializable: {err}" + )) + })?; + let namespace_path = format!( + "$.{}", + serde_json::to_string(namespace).map_err(|err| { + DataLayerError::UnexpectedValue(format!( + "provider_api_keys.upstream_metadata namespace is not serializable: {err}" + )) + })? + ); + let rows_affected = sqlx::query( + r#" +UPDATE provider_api_keys +SET upstream_metadata = JSON_SET( + COALESCE(NULLIF(upstream_metadata, ''), '{}'), + ?, CAST(? AS JSON) + ), + updated_at = ? +WHERE id = ? +"#, + ) + .bind(namespace_path) + .bind(value_json) + .bind(updated_at_unix_secs.unwrap_or_else(current_unix_secs) as i64) + .bind(key_id) + .execute(&self.pool) + .await + .map_sql_err()? + .rows_affected(); + Ok(rows_affected > 0) + } + + pub async fn update_key_model_fetch_state( + &self, + key_id: &str, + allowed_models: Option<&serde_json::Value>, + last_models_fetch_at_unix_secs: Option, + last_models_fetch_error: Option<&str>, + updated_at_unix_secs: Option, + ) -> Result { + validate_non_empty(key_id, "provider catalog key_id")?; + let rows_affected = sqlx::query( + r#" +UPDATE provider_api_keys +SET allowed_models = ?, last_models_fetch_at = ?, last_models_fetch_error = ?, updated_at = ? +WHERE id = ? +"#, + ) + .bind(optional_json_ref_to_string( + allowed_models, + "provider_api_keys.allowed_models", + )?) + .bind(optional_i64_from_u64( + last_models_fetch_at_unix_secs, + "provider_api_keys.last_models_fetch_at", + )?) + .bind(last_models_fetch_error) + .bind(updated_at_unix_secs.unwrap_or_else(current_unix_secs) as i64) + .bind(key_id) + .execute(&self.pool) + .await + .map_sql_err()? + .rows_affected(); + Ok(rows_affected > 0) + } + + pub async fn update_key_model_fetch_success( + &self, + key_id: &str, + allowed_models: Option<&serde_json::Value>, + last_models_fetch_at_unix_secs: u64, + upstream_metadata_updates: &[ProviderCatalogUpstreamMetadataNamespaceUpdate], + updated_at_unix_secs: Option, + ) -> Result { + validate_non_empty(key_id, "provider catalog key_id")?; + let allowed_models = + optional_json_ref_to_string(allowed_models, "provider_api_keys.allowed_models")?; + let namespace_updates = upstream_metadata_updates + .iter() + .map(|update| { + validate_non_empty( + &update.namespace, + "provider catalog upstream metadata namespace", + )?; + let path = format!( + "$.{}", + serde_json::to_string(&update.namespace).map_err(|err| { + DataLayerError::UnexpectedValue(format!( + "provider_api_keys.upstream_metadata namespace is not serializable: {err}" + )) + })? + ); + let value = serde_json::to_string(&update.value).map_err(|err| { + DataLayerError::UnexpectedValue(format!( + "provider_api_keys.upstream_metadata namespace is not serializable: {err}" + )) + })?; + Ok((path, value)) + }) + .collect::, DataLayerError>>()?; + let updated_at = updated_at_unix_secs.unwrap_or_else(current_unix_secs) as i64; + let mut tx = self.pool.begin().await.map_sql_err()?; + let rows_affected = sqlx::query( + r#" +UPDATE provider_api_keys +SET allowed_models = ?, last_models_fetch_at = ?, last_models_fetch_error = NULL, updated_at = ? +WHERE id = ? +"#, + ) + .bind(allowed_models) + .bind(optional_i64_from_u64( + Some(last_models_fetch_at_unix_secs), + "provider_api_keys.last_models_fetch_at", + )?) + .bind(updated_at) + .bind(key_id) + .execute(&mut *tx) + .await + .map_sql_err()? + .rows_affected(); + if rows_affected == 0 { + tx.rollback().await.map_sql_err()?; + return Ok(false); + } + for (path, value) in namespace_updates { + sqlx::query( + r#" +UPDATE provider_api_keys +SET upstream_metadata = JSON_SET( + COALESCE(NULLIF(upstream_metadata, ''), '{}'), + ?, CAST(? AS JSON) + ) +WHERE id = ? +"#, + ) + .bind(path) + .bind(value) + .bind(key_id) + .execute(&mut *tx) + .await + .map_sql_err()?; + } + tx.commit().await.map_sql_err()?; + Ok(true) + } + pub async fn clear_key_oauth_invalid_marker( &self, key_id: &str, @@ -1102,6 +1261,61 @@ impl ProviderCatalogWriteRepository for MysqlProviderCatalogReadRepository { .await } + async fn upsert_key_upstream_metadata_namespace( + &self, + key_id: &str, + namespace: &str, + value: &serde_json::Value, + updated_at_unix_secs: Option, + ) -> Result { + Self::upsert_key_upstream_metadata_namespace( + self, + key_id, + namespace, + value, + updated_at_unix_secs, + ) + .await + } + + async fn update_key_model_fetch_state( + &self, + key_id: &str, + allowed_models: Option<&serde_json::Value>, + last_models_fetch_at_unix_secs: Option, + last_models_fetch_error: Option<&str>, + updated_at_unix_secs: Option, + ) -> Result { + Self::update_key_model_fetch_state( + self, + key_id, + allowed_models, + last_models_fetch_at_unix_secs, + last_models_fetch_error, + updated_at_unix_secs, + ) + .await + } + + async fn update_key_model_fetch_success( + &self, + key_id: &str, + allowed_models: Option<&serde_json::Value>, + last_models_fetch_at_unix_secs: u64, + upstream_metadata_updates: &[ProviderCatalogUpstreamMetadataNamespaceUpdate], + updated_at_unix_secs: Option, + ) -> Result { + Self::update_key_model_fetch_success( + self, + key_id, + allowed_models, + last_models_fetch_at_unix_secs, + upstream_metadata_updates, + updated_at_unix_secs, + ) + .await + } + async fn delete_key(&self, key_id: &str) -> Result { Self::delete_key(self, key_id).await } diff --git a/crates/aether-data/src/repository/provider_catalog/postgres.rs b/crates/aether-data/src/repository/provider_catalog/postgres.rs index ae3598c3c..be46bae0b 100644 --- a/crates/aether-data/src/repository/provider_catalog/postgres.rs +++ b/crates/aether-data/src/repository/provider_catalog/postgres.rs @@ -4,7 +4,8 @@ use sqlx::{postgres::PgRow, PgPool, Postgres, QueryBuilder, Row}; use super::{ ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, ProviderCatalogReadRepository, - ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey, + ProviderCatalogUpstreamMetadataNamespaceUpdate, ProviderCatalogWriteRepository, + StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider, }; @@ -1857,6 +1858,147 @@ WHERE id = $1 Ok(rows_affected > 0) } + pub async fn upsert_key_upstream_metadata_namespace( + &self, + key_id: &str, + namespace: &str, + value: &serde_json::Value, + updated_at_unix_secs: Option, + ) -> Result { + if key_id.trim().is_empty() || namespace.trim().is_empty() { + return Err(DataLayerError::InvalidInput( + "provider catalog key_id and upstream metadata namespace are required".to_string(), + )); + } + let rows_affected = sqlx::query( + r#" +UPDATE provider_api_keys +SET + upstream_metadata = COALESCE(upstream_metadata, '{}'::jsonb) || jsonb_build_object($2, $3::jsonb), + updated_at = CASE + WHEN $4::double precision IS NULL THEN NOW() + ELSE TO_TIMESTAMP($4::double precision) + END +WHERE id = $1 +"#, + ) + .bind(key_id) + .bind(namespace) + .bind(value) + .bind(updated_at_unix_secs.map(|value| value as f64)) + .execute(&self.pool) + .await + .map_postgres_err()? + .rows_affected(); + Ok(rows_affected > 0) + } + + pub async fn update_key_model_fetch_state( + &self, + key_id: &str, + allowed_models: Option<&serde_json::Value>, + last_models_fetch_at_unix_secs: Option, + last_models_fetch_error: Option<&str>, + updated_at_unix_secs: Option, + ) -> Result { + if key_id.trim().is_empty() { + return Err(DataLayerError::InvalidInput( + "provider catalog key_id is empty".to_string(), + )); + } + let rows_affected = sqlx::query( + r#" +UPDATE provider_api_keys +SET + allowed_models = $2, + last_models_fetch_at = CASE + WHEN $3::double precision IS NULL THEN NULL + ELSE TO_TIMESTAMP($3::double precision) + END, + last_models_fetch_error = $4, + updated_at = CASE + WHEN $5::double precision IS NULL THEN NOW() + ELSE TO_TIMESTAMP($5::double precision) + END +WHERE id = $1 +"#, + ) + .bind(key_id) + .bind(allowed_models) + .bind(last_models_fetch_at_unix_secs.map(|value| value as f64)) + .bind(last_models_fetch_error) + .bind(updated_at_unix_secs.map(|value| value as f64)) + .execute(&self.pool) + .await + .map_postgres_err()? + .rows_affected(); + Ok(rows_affected > 0) + } + + pub async fn update_key_model_fetch_success( + &self, + key_id: &str, + allowed_models: Option<&serde_json::Value>, + last_models_fetch_at_unix_secs: u64, + upstream_metadata_updates: &[ProviderCatalogUpstreamMetadataNamespaceUpdate], + updated_at_unix_secs: Option, + ) -> Result { + if key_id.trim().is_empty() + || upstream_metadata_updates + .iter() + .any(|update| update.namespace.trim().is_empty()) + { + return Err(DataLayerError::InvalidInput( + "provider catalog key_id and upstream metadata namespaces are required".to_string(), + )); + } + let mut tx = self.pool.begin().await.map_postgres_err()?; + let rows_affected = sqlx::query( + r#" +UPDATE provider_api_keys +SET + allowed_models = $2, + last_models_fetch_at = TO_TIMESTAMP($3::double precision), + last_models_fetch_error = NULL, + updated_at = CASE + WHEN $4::double precision IS NULL THEN NOW() + ELSE TO_TIMESTAMP($4::double precision) + END +WHERE id = $1 +"#, + ) + .bind(key_id) + .bind(allowed_models) + .bind(last_models_fetch_at_unix_secs as f64) + .bind(updated_at_unix_secs.map(|value| value as f64)) + .execute(&mut *tx) + .await + .map_postgres_err()? + .rows_affected(); + if rows_affected == 0 { + tx.rollback().await.map_postgres_err()?; + return Ok(false); + } + for update in upstream_metadata_updates { + sqlx::query( + r#" +UPDATE provider_api_keys +SET upstream_metadata = COALESCE(upstream_metadata, '{}'::jsonb) + || jsonb_build_object($2, $3::jsonb) +WHERE id = $1 +"#, + ) + .bind(key_id) + .bind(&update.namespace) + .bind(&update.value) + .execute(&mut *tx) + .await + .map_postgres_err()?; + } + tx.commit().await.map_postgres_err()?; + Ok(true) + } + pub async fn update_key_health_state( &self, key_id: &str, @@ -2047,6 +2189,61 @@ impl ProviderCatalogWriteRepository for SqlxProviderCatalogReadRepository { .await } + async fn upsert_key_upstream_metadata_namespace( + &self, + key_id: &str, + namespace: &str, + value: &serde_json::Value, + updated_at_unix_secs: Option, + ) -> Result { + Self::upsert_key_upstream_metadata_namespace( + self, + key_id, + namespace, + value, + updated_at_unix_secs, + ) + .await + } + + async fn update_key_model_fetch_state( + &self, + key_id: &str, + allowed_models: Option<&serde_json::Value>, + last_models_fetch_at_unix_secs: Option, + last_models_fetch_error: Option<&str>, + updated_at_unix_secs: Option, + ) -> Result { + Self::update_key_model_fetch_state( + self, + key_id, + allowed_models, + last_models_fetch_at_unix_secs, + last_models_fetch_error, + updated_at_unix_secs, + ) + .await + } + + async fn update_key_model_fetch_success( + &self, + key_id: &str, + allowed_models: Option<&serde_json::Value>, + last_models_fetch_at_unix_secs: u64, + upstream_metadata_updates: &[ProviderCatalogUpstreamMetadataNamespaceUpdate], + updated_at_unix_secs: Option, + ) -> Result { + Self::update_key_model_fetch_success( + self, + key_id, + allowed_models, + last_models_fetch_at_unix_secs, + upstream_metadata_updates, + updated_at_unix_secs, + ) + .await + } + async fn delete_key(&self, key_id: &str) -> Result { Self::delete_key(self, key_id).await } diff --git a/crates/aether-data/src/repository/provider_catalog/sqlite.rs b/crates/aether-data/src/repository/provider_catalog/sqlite.rs index 24784d80c..f90bcb00c 100644 --- a/crates/aether-data/src/repository/provider_catalog/sqlite.rs +++ b/crates/aether-data/src/repository/provider_catalog/sqlite.rs @@ -3,7 +3,8 @@ use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite}; use super::{ ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, ProviderCatalogReadRepository, - ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey, + ProviderCatalogUpstreamMetadataNamespaceUpdate, ProviderCatalogWriteRepository, + StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider, }; @@ -1224,6 +1225,164 @@ WHERE id = ? Ok(rows_affected > 0) } + pub async fn upsert_key_upstream_metadata_namespace( + &self, + key_id: &str, + namespace: &str, + value: &serde_json::Value, + updated_at_unix_secs: Option, + ) -> Result { + validate_non_empty(key_id, "provider catalog key_id")?; + validate_non_empty(namespace, "provider catalog upstream metadata namespace")?; + let value_json = serde_json::to_string(value).map_err(|err| { + DataLayerError::UnexpectedValue(format!( + "provider_api_keys.upstream_metadata namespace is not serializable: {err}" + )) + })?; + let namespace_path = format!( + "$.{}", + serde_json::to_string(namespace).map_err(|err| { + DataLayerError::UnexpectedValue(format!( + "provider_api_keys.upstream_metadata namespace is not serializable: {err}" + )) + })? + ); + let rows_affected = sqlx::query( + r#" +UPDATE provider_api_keys +SET upstream_metadata = json_set( + COALESCE(NULLIF(upstream_metadata, ''), '{}'), + ?, json(?) + ), + updated_at = ? +WHERE id = ? +"#, + ) + .bind(namespace_path) + .bind(value_json) + .bind(updated_at_unix_secs.unwrap_or_else(current_unix_secs) as i64) + .bind(key_id) + .execute(&self.pool) + .await + .map_sql_err()? + .rows_affected(); + Ok(rows_affected > 0) + } + + pub async fn update_key_model_fetch_state( + &self, + key_id: &str, + allowed_models: Option<&serde_json::Value>, + last_models_fetch_at_unix_secs: Option, + last_models_fetch_error: Option<&str>, + updated_at_unix_secs: Option, + ) -> Result { + validate_non_empty(key_id, "provider catalog key_id")?; + let rows_affected = sqlx::query( + r#" +UPDATE provider_api_keys +SET allowed_models = ?, last_models_fetch_at = ?, last_models_fetch_error = ?, updated_at = ? +WHERE id = ? +"#, + ) + .bind(optional_json_ref_to_string( + allowed_models, + "provider_api_keys.allowed_models", + )?) + .bind(optional_i64_from_u64( + last_models_fetch_at_unix_secs, + "provider_api_keys.last_models_fetch_at", + )?) + .bind(last_models_fetch_error) + .bind(updated_at_unix_secs.unwrap_or_else(current_unix_secs) as i64) + .bind(key_id) + .execute(&self.pool) + .await + .map_sql_err()? + .rows_affected(); + Ok(rows_affected > 0) + } + + pub async fn update_key_model_fetch_success( + &self, + key_id: &str, + allowed_models: Option<&serde_json::Value>, + last_models_fetch_at_unix_secs: u64, + upstream_metadata_updates: &[ProviderCatalogUpstreamMetadataNamespaceUpdate], + updated_at_unix_secs: Option, + ) -> Result { + validate_non_empty(key_id, "provider catalog key_id")?; + let allowed_models = + optional_json_ref_to_string(allowed_models, "provider_api_keys.allowed_models")?; + let namespace_updates = upstream_metadata_updates + .iter() + .map(|update| { + validate_non_empty( + &update.namespace, + "provider catalog upstream metadata namespace", + )?; + let path = format!( + "$.{}", + serde_json::to_string(&update.namespace).map_err(|err| { + DataLayerError::UnexpectedValue(format!( + "provider_api_keys.upstream_metadata namespace is not serializable: {err}" + )) + })? + ); + let value = serde_json::to_string(&update.value).map_err(|err| { + DataLayerError::UnexpectedValue(format!( + "provider_api_keys.upstream_metadata namespace is not serializable: {err}" + )) + })?; + Ok((path, value)) + }) + .collect::, DataLayerError>>()?; + let updated_at = updated_at_unix_secs.unwrap_or_else(current_unix_secs) as i64; + let mut tx = self.pool.begin().await.map_sql_err()?; + let rows_affected = sqlx::query( + r#" +UPDATE provider_api_keys +SET allowed_models = ?, last_models_fetch_at = ?, last_models_fetch_error = NULL, updated_at = ? +WHERE id = ? +"#, + ) + .bind(allowed_models) + .bind(optional_i64_from_u64( + Some(last_models_fetch_at_unix_secs), + "provider_api_keys.last_models_fetch_at", + )?) + .bind(updated_at) + .bind(key_id) + .execute(&mut *tx) + .await + .map_sql_err()? + .rows_affected(); + if rows_affected == 0 { + tx.rollback().await.map_sql_err()?; + return Ok(false); + } + for (path, value) in namespace_updates { + sqlx::query( + r#" +UPDATE provider_api_keys +SET upstream_metadata = json_set( + COALESCE(NULLIF(upstream_metadata, ''), '{}'), + ?, json(?) + ) +WHERE id = ? +"#, + ) + .bind(path) + .bind(value) + .bind(key_id) + .execute(&mut *tx) + .await + .map_sql_err()?; + } + tx.commit().await.map_sql_err()?; + Ok(true) + } + pub async fn clear_key_oauth_invalid_marker( &self, key_id: &str, @@ -1511,6 +1670,61 @@ impl ProviderCatalogWriteRepository for SqliteProviderCatalogReadRepository { .await } + async fn upsert_key_upstream_metadata_namespace( + &self, + key_id: &str, + namespace: &str, + value: &serde_json::Value, + updated_at_unix_secs: Option, + ) -> Result { + Self::upsert_key_upstream_metadata_namespace( + self, + key_id, + namespace, + value, + updated_at_unix_secs, + ) + .await + } + + async fn update_key_model_fetch_state( + &self, + key_id: &str, + allowed_models: Option<&serde_json::Value>, + last_models_fetch_at_unix_secs: Option, + last_models_fetch_error: Option<&str>, + updated_at_unix_secs: Option, + ) -> Result { + Self::update_key_model_fetch_state( + self, + key_id, + allowed_models, + last_models_fetch_at_unix_secs, + last_models_fetch_error, + updated_at_unix_secs, + ) + .await + } + + async fn update_key_model_fetch_success( + &self, + key_id: &str, + allowed_models: Option<&serde_json::Value>, + last_models_fetch_at_unix_secs: u64, + upstream_metadata_updates: &[ProviderCatalogUpstreamMetadataNamespaceUpdate], + updated_at_unix_secs: Option, + ) -> Result { + Self::update_key_model_fetch_success( + self, + key_id, + allowed_models, + last_models_fetch_at_unix_secs, + upstream_metadata_updates, + updated_at_unix_secs, + ) + .await + } + async fn delete_key(&self, key_id: &str) -> Result { Self::delete_key(self, key_id).await } @@ -2117,7 +2331,8 @@ mod tests { use super::SqliteProviderCatalogReadRepository; use crate::lifecycle::migrate::run_sqlite_migrations; use crate::repository::provider_catalog::{ - ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, StoredProviderCatalogEndpoint, + ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, + ProviderCatalogUpstreamMetadataNamespaceUpdate, StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider, }; use serde_json::json; @@ -2354,11 +2569,38 @@ mod tests { assert!(repository .update_key_upstream_metadata( "key-write-1", - Some(&json!({"models":["gpt-4.1-mini"]})), + Some(&json!({ + "codex": { + "quota_by_model": { + "gpt-5.6-sol": {"remaining_fraction": 0.75} + } + }, + "codex_models": {"cards": {"old": {"slug": "old"}}} + })), Some(1_740_000_000), ) .await .expect("upstream metadata should update")); + assert!(repository + .update_key_model_fetch_success( + "key-write-1", + Some(&json!(["gpt-5.6-sol"])), + 1_740_000_002, + &[ProviderCatalogUpstreamMetadataNamespaceUpdate { + namespace: "codex_models".to_string(), + value: json!({ + "cards": { + "gpt-5.6-sol": { + "slug": "gpt-5.6-sol", + "use_responses_lite": true + } + } + }), + }], + Some(1_740_000_002), + ) + .await + .expect("model fetch success should update atomically")); assert!(repository .update_key_oauth_credentials( "key-write-1", @@ -2394,7 +2636,26 @@ mod tests { ); assert_eq!( reloaded_key.upstream_metadata, - Some(json!({"models":["gpt-4.1-mini"]})) + Some(json!({ + "codex": { + "quota_by_model": { + "gpt-5.6-sol": {"remaining_fraction": 0.75} + } + }, + "codex_models": { + "cards": { + "gpt-5.6-sol": { + "slug": "gpt-5.6-sol", + "use_responses_lite": true + } + } + } + })) + ); + assert_eq!(reloaded_key.allowed_models, Some(json!(["gpt-5.6-sol"]))); + assert_eq!( + reloaded_key.last_models_fetch_at_unix_secs, + Some(1_740_000_002) ); assert!(reloaded_key.is_active); diff --git a/crates/aether-data/src/repository/usage/memory.rs b/crates/aether-data/src/repository/usage/memory.rs index a263906f5..e775b7ec4 100644 --- a/crates/aether-data/src/repository/usage/memory.rs +++ b/crates/aether-data/src/repository/usage/memory.rs @@ -2761,14 +2761,18 @@ impl UsageWriteRepository for InMemoryUsageReadRepository { }) .unwrap_or_default(); if existing.as_ref().is_some_and(|existing| { - usage_status_is_finalized(existing.status.as_str()) - && usage_status_is_lifecycle(usage.status.as_str()) - && !usage_can_recover_terminal_failure( - existing.status.as_str(), - existing.billing_status.as_str(), - usage.status.as_str(), - usage.billing_status.as_str(), - ) + let can_recover = usage_can_recover_terminal_failure( + existing.status.as_str(), + existing.billing_status.as_str(), + usage.status.as_str(), + usage.billing_status.as_str(), + ); + let finalized_lifecycle_regression = usage_status_is_finalized(&existing.status) + && usage_status_is_lifecycle(&usage.status); + let completed_terminal_failure_recovery = existing.billing_status == "void" + && matches!(existing.status.as_str(), "failed" | "cancelled") + && usage.status == "completed"; + (finalized_lifecycle_regression || completed_terminal_failure_recovery) && !can_recover }) { return Ok(existing.expect("existing usage should be present").clone()); } @@ -3589,7 +3593,7 @@ mod tests { } #[tokio::test] - async fn upsert_allows_streaming_recovery_after_void_failure() { + async fn upsert_allows_completed_recovery_after_void_failure() { let repository = InMemoryUsageReadRepository::default(); repository .upsert(UpsertUsageRecord { @@ -3698,12 +3702,12 @@ mod tests { output_price_per_1m: None, total_cost_usd: None, actual_total_cost_usd: None, - status_code: None, + status_code: Some(200), error_message: None, error_category: None, response_time_ms: Some(45), first_byte_time_ms: Some(12), - status: "streaming".to_string(), + status: "completed".to_string(), billing_status: "pending".to_string(), request_headers: None, request_body: None, @@ -3732,7 +3736,7 @@ mod tests { request_metadata: Some(json!({ "trace_id": "trace-recovered" })), - finalized_at_unix_secs: None, + finalized_at_unix_secs: Some(102), created_at_unix_ms: Some(100), updated_at_unix_secs: 102, }) @@ -3744,11 +3748,11 @@ mod tests { .await .expect("usage lookup should succeed") .expect("usage should exist"); - assert_eq!(stored.status, "streaming"); + assert_eq!(stored.status, "completed"); assert_eq!(stored.billing_status, "pending"); - assert_eq!(stored.status_code, None); + assert_eq!(stored.status_code, Some(200)); assert_eq!(stored.error_message, None); - assert_eq!(stored.finalized_at_unix_secs, None); + assert_eq!(stored.finalized_at_unix_secs, Some(102)); assert_eq!( stored.request_metadata, Some(json!({ "trace_id": "trace-recovered" })) @@ -3756,6 +3760,55 @@ mod tests { assert_eq!(stored.total_tokens, 10); } + #[tokio::test] + async fn upsert_rejects_non_authoritative_void_failure_recovery() { + let repository = InMemoryUsageReadRepository::default(); + for (request_id, status, billing_status, status_code) in [ + ("req-late-active-1", "streaming", "pending", None), + ( + "req-late-response-start-1", + "streaming", + "pending", + Some(200), + ), + ( + "req-settled-completion-1", + "completed", + "settled", + Some(200), + ), + ] { + repository + .upsert(UpsertUsageRecord { + status: "failed".to_string(), + billing_status: "void".to_string(), + status_code: Some(503), + finalized_at_unix_secs: Some(101), + updated_at_unix_secs: 101, + ..sample_upsert_usage_record(request_id) + }) + .await + .expect("failed usage should upsert"); + + let stored = repository + .upsert(UpsertUsageRecord { + status: status.to_string(), + billing_status: billing_status.to_string(), + status_code, + finalized_at_unix_secs: None, + updated_at_unix_secs: 102, + ..sample_upsert_usage_record(request_id) + }) + .await + .expect("non-authoritative recovery should be ignored"); + + assert_eq!(stored.status, "failed"); + assert_eq!(stored.billing_status, "void"); + assert_eq!(stored.status_code, Some(503)); + assert_eq!(stored.finalized_at_unix_secs, Some(101)); + } + } + #[tokio::test] async fn stale_pending_update_does_not_reopen_void_failure() { let repository = InMemoryUsageReadRepository::default(); @@ -4924,9 +4977,9 @@ mod tests { }) .await .expect("dashboard should summarize"); - assert_eq!(dashboard.effective_input_tokens, 20); + assert_eq!(dashboard.effective_input_tokens, 0); assert_eq!(dashboard.cache_creation_tokens, 20); - assert_eq!(dashboard.total_tokens, 140); + assert_eq!(dashboard.total_tokens, 120); let leaderboard = repository .summarize_usage_leaderboard(&UsageLeaderboardQuery { @@ -4940,7 +4993,7 @@ mod tests { .await .expect("leaderboard should summarize"); assert_eq!(leaderboard.len(), 1); - assert_eq!(leaderboard[0].total_tokens, 140); + assert_eq!(leaderboard[0].total_tokens, 120); } #[tokio::test] diff --git a/crates/aether-data/src/repository/usage/mod.rs b/crates/aether-data/src/repository/usage/mod.rs index 162992117..94b2a7db7 100644 --- a/crates/aether-data/src/repository/usage/mod.rs +++ b/crates/aether-data/src/repository/usage/mod.rs @@ -583,10 +583,7 @@ pub(crate) fn incoming_usage_can_recover_terminal_failure( incoming_status: &str, incoming_billing_status: &str, ) -> bool { - incoming_billing_status == "pending" - // Late pending placeholders are not authoritative enough to reopen a void terminal row; - // they can otherwise regress a real failure back to pending when background writes race. - && matches!(incoming_status, "streaming" | "completed") + incoming_billing_status == "pending" && incoming_status == "completed" } pub(crate) fn usage_can_recover_terminal_failure( @@ -827,52 +824,58 @@ mod tests { } #[test] - fn incoming_usage_recovery_requires_streaming_or_completed_state() { + fn incoming_usage_recovery_requires_completed_state() { assert!(incoming_usage_can_recover_terminal_failure( "completed", - "pending" + "pending", )); - assert!(incoming_usage_can_recover_terminal_failure( + assert!(!incoming_usage_can_recover_terminal_failure( "streaming", - "pending" + "pending", )); assert!(!incoming_usage_can_recover_terminal_failure( - "pending", "pending" + "pending", "pending", )); assert!(!incoming_usage_can_recover_terminal_failure( - "failed", "void" + "failed", "void", )); assert!(!incoming_usage_can_recover_terminal_failure( "completed", - "settled" + "settled", )); } #[test] - fn usage_recovery_requires_void_failure_to_be_followed_by_streaming_or_completed_state() { + fn usage_recovery_requires_void_failure_and_completed_state() { assert!(usage_can_recover_terminal_failure( "failed", "void", "completed", - "pending" + "pending", )); assert!(usage_can_recover_terminal_failure( "cancelled", "void", - "streaming", - "pending" + "completed", + "pending", )); assert!(!usage_can_recover_terminal_failure( - "failed", "void", "pending", "pending" + "failed", + "void", + "streaming", + "pending", + )); + assert!(!usage_can_recover_terminal_failure( + "failed", "void", "pending", "pending", )); assert!(!usage_can_recover_terminal_failure( "completed", "pending", "completed", - "pending" + "pending", )); assert!(!usage_can_recover_terminal_failure( - "failed", "void", "failed", "void" + "failed", "void", "failed", "void", )); } diff --git a/crates/aether-data/src/repository/usage/postgres/queries/list_recent_usage_audits_prefix.sql b/crates/aether-data/src/repository/usage/postgres/queries/list_recent_usage_audits_prefix.sql index a2d786624..fbfdf7e0c 100644 --- a/crates/aether-data/src/repository/usage/postgres/queries/list_recent_usage_audits_prefix.sql +++ b/crates/aether-data/src/repository/usage/postgres/queries/list_recent_usage_audits_prefix.sql @@ -116,6 +116,7 @@ SELECT OR NULLIF(BTRIM("usage".request_metadata->>'request_path_and_query'), '') IS NOT NULL OR NULLIF(BTRIM("usage".request_metadata->>'provider_reasoning_effort'), '') IS NOT NULL OR NULLIF(BTRIM("usage".request_metadata->>'provider_service_tier'), '') IS NOT NULL + OR NULLIF(BTRIM("usage".request_metadata->>'provider_actual_service_tier'), '') IS NOT NULL OR ("usage".request_metadata->>'client_requested_stream') IN ('true', 'false') OR ("usage".request_metadata->>'upstream_is_stream') IN ('true', 'false') THEN jsonb_strip_nulls(jsonb_build_object( @@ -131,6 +132,8 @@ SELECT NULLIF(BTRIM("usage".request_metadata->>'provider_reasoning_effort'), ''), 'provider_service_tier', NULLIF(BTRIM("usage".request_metadata->>'provider_service_tier'), ''), + 'provider_actual_service_tier', + NULLIF(BTRIM("usage".request_metadata->>'provider_actual_service_tier'), ''), 'client_requested_stream', CASE WHEN ("usage".request_metadata->>'client_requested_stream') IN ('true', 'false') diff --git a/crates/aether-data/src/repository/usage/postgres/queries/list_usage_audits_prefix.sql b/crates/aether-data/src/repository/usage/postgres/queries/list_usage_audits_prefix.sql index a2d786624..fbfdf7e0c 100644 --- a/crates/aether-data/src/repository/usage/postgres/queries/list_usage_audits_prefix.sql +++ b/crates/aether-data/src/repository/usage/postgres/queries/list_usage_audits_prefix.sql @@ -116,6 +116,7 @@ SELECT OR NULLIF(BTRIM("usage".request_metadata->>'request_path_and_query'), '') IS NOT NULL OR NULLIF(BTRIM("usage".request_metadata->>'provider_reasoning_effort'), '') IS NOT NULL OR NULLIF(BTRIM("usage".request_metadata->>'provider_service_tier'), '') IS NOT NULL + OR NULLIF(BTRIM("usage".request_metadata->>'provider_actual_service_tier'), '') IS NOT NULL OR ("usage".request_metadata->>'client_requested_stream') IN ('true', 'false') OR ("usage".request_metadata->>'upstream_is_stream') IN ('true', 'false') THEN jsonb_strip_nulls(jsonb_build_object( @@ -131,6 +132,8 @@ SELECT NULLIF(BTRIM("usage".request_metadata->>'provider_reasoning_effort'), ''), 'provider_service_tier', NULLIF(BTRIM("usage".request_metadata->>'provider_service_tier'), ''), + 'provider_actual_service_tier', + NULLIF(BTRIM("usage".request_metadata->>'provider_actual_service_tier'), ''), 'client_requested_stream', CASE WHEN ("usage".request_metadata->>'client_requested_stream') IN ('true', 'false') diff --git a/crates/aether-data/src/repository/usage/postgres/tests.rs b/crates/aether-data/src/repository/usage/postgres/tests.rs index e92966b93..403c24318 100644 --- a/crates/aether-data/src/repository/usage/postgres/tests.rs +++ b/crates/aether-data/src/repository/usage/postgres/tests.rs @@ -733,6 +733,8 @@ fn usage_sql_uses_json_null_placeholders_for_usage_payload_columns() { assert!(sql.contains("request_metadata->>'client_ip'")); assert!(sql.contains("'user_agent'")); assert!(sql.contains("request_metadata->>'user_agent'")); + assert!(sql.contains("request_metadata->>'provider_service_tier'")); + assert!(sql.contains("request_metadata->>'provider_actual_service_tier'")); assert!(sql.contains("AS client_family")); assert!(sql.contains("request_metadata->'client_session_affinity'->>'client_family'")); assert!(sql.contains("request_metadata->>'client_family'")); diff --git a/crates/aether-data/src/repository/usage/sqlite.rs b/crates/aether-data/src/repository/usage/sqlite.rs index 29ab28d74..d61540936 100644 --- a/crates/aether-data/src/repository/usage/sqlite.rs +++ b/crates/aether-data/src/repository/usage/sqlite.rs @@ -4598,6 +4598,45 @@ mod tests { assert_eq!(existing.updated_at_unix_secs, 1_000); } + #[tokio::test] + async fn sqlite_usage_write_repository_does_not_reopen_void_failure_from_late_streaming() { + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .expect("sqlite pool should connect"); + run_sqlite_migrations(&pool) + .await + .expect("sqlite migrations should run"); + seed_stats_targets(&pool).await; + + let repository = SqliteUsageWriteRepository::new(pool); + for (request_id, status_code) in [ + ("request-late-active", None), + ("request-late-response-start", Some(200)), + ] { + let mut failed = sample_usage(request_id, "failed", "void", 1_000); + failed.status_code = Some(503); + repository + .upsert(failed) + .await + .expect("failed usage should upsert"); + + let mut late_streaming = sample_usage(request_id, "streaming", "pending", 1_001); + late_streaming.status_code = status_code; + late_streaming.finalized_at_unix_secs = None; + let current = repository + .upsert(late_streaming) + .await + .expect("late streaming usage should be ignored"); + + assert_eq!(current.status, "failed"); + assert_eq!(current.billing_status, "void"); + assert_eq!(current.status_code, Some(503)); + assert_eq!(current.finalized_at_unix_secs, Some(1_000)); + } + } + #[tokio::test] async fn sqlite_usage_write_repository_does_not_regress_terminal_usage_from_late_streaming() { let pool = sqlx::sqlite::SqlitePoolOptions::new() diff --git a/crates/aether-model-fetch/src/lib.rs b/crates/aether-model-fetch/src/lib.rs index 4f54ac9e7..6ce0e69c6 100644 --- a/crates/aether-model-fetch/src/lib.rs +++ b/crates/aether-model-fetch/src/lib.rs @@ -13,10 +13,12 @@ pub use config::{ pub use logic::{ aggregate_models_for_cache, apply_model_filters, build_models_fetch_url, deepseek_anthropic_models_fetch_uses_openai_auth, endpoint_supports_rust_models_fetch, - extract_error_message, json_string_list, merge_upstream_metadata, parse_models_response, - parse_models_response_page, parse_windsurf_model_configs_response, preset_models_for_provider, + extract_error_message, json_string_list, merge_upstream_metadata, + model_catalog_upstream_metadata, parse_models_response, parse_models_response_page, + parse_windsurf_model_configs_response, preset_models_for_provider, provider_type_uses_preset_models, select_models_fetch_endpoint, - selected_models_fetch_endpoints, ModelFetchRunSummary, ModelsFetchPage, ModelsFetchSuccess, + selected_models_fetch_endpoints, upstream_metadata_namespace_updates, ModelFetchRunSummary, + ModelsFetchPage, ModelsFetchSuccess, }; pub use strategy::{ fetch_models_from_transports, ModelFetchStrategy, ModelFetchStrategyKind, ModelsFetchOutcome, diff --git a/crates/aether-model-fetch/src/logic.rs b/crates/aether-model-fetch/src/logic.rs index 866af9a62..a7d425129 100644 --- a/crates/aether-model-fetch/src/logic.rs +++ b/crates/aether-model-fetch/src/logic.rs @@ -373,13 +373,7 @@ pub fn preset_models_for_provider(provider_type: &str) -> Option> { preset_model("claude-sonnet-4-5-20250929", "anthropic", "Claude Sonnet 4.5", "claude:messages"), preset_model("claude-haiku-4-5-20251001", "anthropic", "Claude Haiku 4.5", "claude:messages"), ], - "codex" => vec![ - preset_model("gpt-5.5", "openai", "GPT-5.5", "openai:responses"), - preset_model("gpt-5.4", "openai", "GPT-5.4", "openai:responses"), - preset_model("gpt-5.4-mini", "openai", "GPT-5.4 Mini", "openai:responses"), - preset_model("gpt-5.3-codex", "openai", "GPT-5.3 Codex", "openai:responses"), - preset_model("gpt-5.3-codex-spark", "openai", "GPT-5.3 Codex Spark", "openai:responses"), - ], + "codex" => aether_ai_formats::bundled_codex_model_cards().to_vec(), "grok" => vec![ preset_model("grok-4.20-0309-non-reasoning", "xai", "Grok 4.20 0309 Non-Reasoning", "openai:chat"), preset_model("grok-4.20-0309", "xai", "Grok 4.20 0309", "openai:chat"), @@ -451,6 +445,35 @@ pub fn merge_upstream_metadata(current: Option<&Value>, incoming: &Value) -> Val Value::Object(merged) } +pub fn model_catalog_upstream_metadata( + provider_type: &str, + cached_models: &[Value], +) -> Option { + provider_type.trim().eq_ignore_ascii_case("codex").then(|| { + let cards = aether_ai_formats::effective_codex_model_cards(cached_models); + aether_ai_formats::build_codex_model_catalog_metadata(&cards) + }) +} + +pub fn upstream_metadata_namespace_updates( + current: Option<&Value>, + incoming: &Value, +) -> Vec<(String, Value)> { + let Some(incoming) = incoming.as_object() else { + return Vec::new(); + }; + let merged = merge_upstream_metadata(current, &Value::Object(incoming.clone())); + incoming + .keys() + .filter_map(|namespace| { + merged + .get(namespace) + .cloned() + .map(|value| (namespace.clone(), value)) + }) + .collect() +} + pub fn apply_model_filters( fetched_model_ids: &[String], locked_models: Vec, @@ -691,7 +714,8 @@ fn build_codex_models_url(base_url: &str) -> Option { if !has_client_version { let separator = if url.contains('?') { '&' } else { '?' }; url.push(separator); - url.push_str("client_version=0.128.0-alpha.1"); + url.push_str("client_version="); + url.push_str(aether_ai_formats::CODEX_CLIENT_VERSION); } Some(url) } @@ -1006,8 +1030,7 @@ mod tests { "https://chatgpt.com/backend-api/codex" ), Some(( - "https://chatgpt.com/backend-api/codex/models?client_version=0.128.0-alpha.1" - .to_string(), + "https://chatgpt.com/backend-api/codex/models?client_version=0.144.1".to_string(), "openai:responses".to_string() )) ); @@ -1135,6 +1158,33 @@ mod tests { ); } + #[test] + fn parse_models_response_preserves_gpt_5_6_model_card_capabilities() { + let card = json!({ + "slug": "gpt-5.6-sol", + "default_reasoning_level": "low", + "supported_reasoning_levels": [ + {"effort": "low"}, + {"effort": "max"}, + {"effort": "ultra"} + ], + "multi_agent_version": "v2", + "supports_image_detail_original": true, + "future_capability": {"mode": "preserve-me"} + }); + let parsed = parse_models_response("openai:responses", &json!({"models": [card]})) + .expect("Codex model card should parse"); + + let cached = &parsed.cached_models[0]; + assert_eq!(cached["id"], "gpt-5.6-sol"); + assert_eq!(cached["default_reasoning_level"], "low"); + assert_eq!(cached["supported_reasoning_levels"][2]["effort"], "ultra"); + assert_eq!(cached["multi_agent_version"], "v2"); + assert_eq!(cached["supports_image_detail_original"], true); + assert_eq!(cached["future_capability"]["mode"], "preserve-me"); + assert_eq!(cached["api_formats"], json!(["openai:responses"])); + } + #[test] fn parse_models_response_page_reads_claude_pagination_state() { let parsed = parse_models_response_page( @@ -1243,13 +1293,71 @@ mod tests { assert_eq!( model_ids, vec![ + "gpt-5.6-sol", + "gpt-5.6-terra", + "gpt-5.6-luna", "gpt-5.5", "gpt-5.4", "gpt-5.4-mini", - "gpt-5.3-codex", - "gpt-5.3-codex-spark", + "gpt-5.2", + "codex-auto-review", ] ); + let sol = models + .iter() + .find(|model| model["id"] == "gpt-5.6-sol") + .expect("Sol preset"); + assert_eq!(sol["default_reasoning_level"], "low"); + assert_eq!( + sol["supported_reasoning_levels"] + .as_array() + .expect("reasoning levels") + .iter() + .filter_map(|level| level["effort"].as_str()) + .collect::>(), + vec!["low", "medium", "high", "xhigh", "max", "ultra"] + ); + assert_eq!(sol["multi_agent_version"], "v2"); + assert_eq!(sol["supports_image_detail_original"], true); + assert_eq!(sol["context_window"], 372_000); + + for model_id in ["gpt-5.6-sol", "gpt-5.6-terra", "gpt-5.6-luna"] { + let model = models + .iter() + .find(|model| model["id"] == model_id) + .expect("GPT-5.6 Codex preset"); + assert_eq!(model["shell_type"], "shell_command"); + assert_eq!(model["comp_hash"], "3000"); + assert_eq!(model["experimental_supported_tools"], json!([])); + assert_eq!(model["tool_mode"], "code_mode_only"); + assert_eq!(model["prefer_websockets"], true); + assert_eq!(model["reasoning_summary_format"], "experimental"); + assert_eq!(model["truncation_policy"]["limit"], 10_000); + assert_eq!(model["minimal_client_version"], "0.144.0"); + assert!(model.get("effective_context_window_percent").is_none()); + } + + let luna = models + .iter() + .find(|model| model["id"] == "gpt-5.6-luna") + .expect("Luna preset"); + assert_eq!(luna["default_reasoning_level"], "medium"); + assert_eq!(luna["multi_agent_version"], "v1"); + assert!(!luna["supported_reasoning_levels"] + .as_array() + .expect("reasoning levels") + .iter() + .any(|level| level["effort"] == "ultra")); + + let auto_review = models + .iter() + .find(|model| model["id"] == "codex-auto-review") + .expect("Codex auto review preset"); + assert_eq!(auto_review["visibility"], "hide"); + assert_eq!(auto_review["supported_in_api"], true); + assert_eq!(auto_review["default_reasoning_level"], "medium"); + assert_eq!(auto_review["default_reasoning_summary"], "none"); + assert_eq!(auto_review["use_responses_lite"], false); } #[test] diff --git a/crates/aether-model-fetch/src/strategy.rs b/crates/aether-model-fetch/src/strategy.rs index c98bfdc01..71e25a644 100644 --- a/crates/aether-model-fetch/src/strategy.rs +++ b/crates/aether-model-fetch/src/strategy.rs @@ -170,7 +170,7 @@ async fn execute_model_fetch_strategy( true, )), ModelFetchStrategyKind::StandardTransport => { - fetch_standard_models(runtime, transports).await + fetch_standard_models(runtime, transports, strategy.provider_id()).await } ModelFetchStrategyKind::Vertex => fetch_vertex_models(runtime, transports).await, ModelFetchStrategyKind::Antigravity => { @@ -192,6 +192,7 @@ async fn execute_model_fetch_strategy( async fn fetch_standard_models( runtime: &(impl ModelFetchTransportRuntime + ?Sized), transports: &[GatewayProviderTransportSnapshot], + provider_type: &str, ) -> Result { let mut all_models = Vec::new(); let mut errors = Vec::new(); @@ -208,7 +209,9 @@ async fn fetch_standard_models( } let merged_models = aggregate_models_for_cache(&all_models); - Ok(build_success_outcome(merged_models, None, has_success).with_errors(errors)) + let upstream_metadata = + crate::logic::model_catalog_upstream_metadata(provider_type, &merged_models); + Ok(build_success_outcome(merged_models, upstream_metadata, has_success).with_errors(errors)) } async fn fetch_standard_models_for_transport( @@ -1996,7 +1999,11 @@ mod tests { executed_urls: Arc::clone(&executed_urls), response_body: json!({ "models": [{ - "id": "gpt-5.4-upstream" + "id": "gpt-5.6-future", + "slug": "gpt-5.6-future", + "default_reasoning_level": "high", + "supported_reasoning_levels": [{"effort": "high"}], + "future_capability": {"mode": "preserve-me"} }] }), status_code: 200, @@ -2008,10 +2015,16 @@ mod tests { let urls = executed_urls.lock().expect("executed_urls lock"); assert_eq!( urls.as_slice(), - &["https://chatgpt.com/backend-api/codex/models?client_version=0.128.0-alpha.1"] + &["https://chatgpt.com/backend-api/codex/models?client_version=0.144.1"] ); - assert_eq!(outcome.fetched_model_ids, vec!["gpt-5.4-upstream"]); + assert_eq!(outcome.fetched_model_ids, vec!["gpt-5.6-future"]); assert_eq!(outcome.cached_models.len(), 1); + let card = &outcome + .upstream_metadata + .as_ref() + .expect("Codex model catalog metadata")["codex_models"]["cards"]["gpt-5.6-future"]; + assert_eq!(card["default_reasoning_level"], "high"); + assert_eq!(card["future_capability"]["mode"], "preserve-me"); } #[tokio::test] diff --git a/crates/aether-model-fetch/src/transport.rs b/crates/aether-model-fetch/src/transport.rs index 1e18e2705..94a7e6857 100644 --- a/crates/aether-model-fetch/src/transport.rs +++ b/crates/aether-model-fetch/src/transport.rs @@ -21,11 +21,10 @@ use aether_provider_transport::{ GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth, }; use async_trait::async_trait; -use serde_json::{json, Value}; +use serde_json::json; use crate::{build_models_fetch_url, deepseek_anthropic_models_fetch_uses_openai_auth}; -const OPENAI_RESPONSES_USER_AGENT: &str = "openai-codex/1.0"; const CLAUDE_CLI_USER_AGENT: &str = "claude-code/1.0.1"; const GEMINI_CLI_USER_AGENT: &str = "GeminiCLI/0.1.5 (Windows; AMD64)"; const CLAUDE_VERSION_HEADER: &str = "2023-06-01"; @@ -130,7 +129,10 @@ pub async fn build_standard_models_fetch_execution_plan( &auth_header_value, ); if is_codex_openai_models_fetch { - if let Some(account_id) = extract_codex_account_id(transport) { + let auth_identity = aether_ai_formats::parse_codex_auth_identity( + transport.key.decrypted_auth_config.as_deref(), + ); + if let Some(account_id) = auth_identity.account_id { insert_non_empty_auth_header( &mut headers, &mut protected_headers, @@ -138,6 +140,14 @@ pub async fn build_standard_models_fetch_execution_plan( &account_id, ); } + if auth_identity.is_fedramp { + insert_non_empty_auth_header( + &mut headers, + &mut protected_headers, + "x-openai-fedramp", + "true", + ); + } } headers = apply_fetch_header_rules(transport, headers, &protected_headers)?; ensure_upstream_auth_header(&mut headers, &auth_header_name, &auth_header_value); @@ -582,10 +592,19 @@ fn standard_models_fetch_headers( let api_format = aether_ai_formats::normalize_api_format_alias(api_format); let provider_type = provider_type.trim().to_ascii_lowercase(); match api_format.as_str() { - "openai:responses" | "openai:responses:compact" => BTreeMap::from([( - "user-agent".to_string(), - OPENAI_RESPONSES_USER_AGENT.to_string(), - )]), + "openai:responses" | "openai:responses:compact" => { + let mut headers = BTreeMap::from([( + "user-agent".to_string(), + aether_ai_formats::CODEX_CLIENT_USER_AGENT.to_string(), + )]); + if provider_type == "codex" { + headers.insert( + "originator".to_string(), + aether_ai_formats::CODEX_CLIENT_ORIGINATOR.to_string(), + ); + } + headers + } "claude:messages" => { let mut headers = BTreeMap::from([( "anthropic-version".to_string(), @@ -669,23 +688,6 @@ fn append_query_param(mut url: String, key: &str, value: &str) -> String { url } -fn extract_codex_account_id(transport: &GatewayProviderTransportSnapshot) -> Option { - let raw = transport.key.decrypted_auth_config.as_deref()?.trim(); - if raw.is_empty() { - return None; - } - - serde_json::from_str::(raw).ok().and_then(|value| { - value - .get("account_id") - .or_else(|| value.get("chatgpt_account_id")) - .and_then(Value::as_str) - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(ToOwned::to_owned) - }) -} - fn insert_non_empty_auth_header( headers: &mut BTreeMap, protected_headers: &mut Vec, @@ -839,7 +841,7 @@ mod tests { assert_eq!(plan.url, "https://example.com/models"); assert_eq!( plan.headers.get("user-agent").map(String::as_str), - Some("openai-codex/1.0") + Some(aether_ai_formats::CODEX_CLIENT_USER_AGENT) ); assert_eq!( plan.headers.get("authorization").map(String::as_str), @@ -910,7 +912,7 @@ mod tests { } #[tokio::test] - async fn builds_codex_models_fetch_plan_with_account_header() { + async fn builds_codex_models_fetch_plan_with_auth_identity_headers() { let runtime = TestRuntime { oauth_auth: Some( aether_provider_transport::LocalResolvedOAuthRequestAuth::Header { @@ -922,7 +924,12 @@ mod tests { }; let mut transport = sample_transport("codex", "openai:responses", "oauth"); transport.endpoint.base_url = "https://chatgpt.com/backend-api/codex".to_string(); - transport.key.decrypted_auth_config = Some(r#"{"account_id":"account-1"}"#.to_string()); + transport.endpoint.header_rules = Some(json!([ + {"op": "set", "name": "chatgpt-account-id", "value": "spoofed-account"}, + {"op": "set", "name": "x-openai-fedramp", "value": "false"} + ])); + transport.key.decrypted_auth_config = + Some(r#"{"account_id":"account-1","chatgpt_account_is_fedramp":true}"#.to_string()); let plan = build_models_fetch_execution_plan(&runtime, &transport) .await @@ -930,7 +937,7 @@ mod tests { assert_eq!( plan.url, - "https://chatgpt.com/backend-api/codex/models?client_version=0.128.0-alpha.1" + "https://chatgpt.com/backend-api/codex/models?client_version=0.144.1" ); assert_eq!( plan.headers.get("authorization").map(String::as_str), @@ -940,10 +947,19 @@ mod tests { plan.headers.get("chatgpt-account-id").map(String::as_str), Some("account-1") ); + assert_eq!( + plan.headers.get("x-openai-fedramp").map(String::as_str), + Some("true") + ); assert_eq!( plan.headers.get("accept").map(String::as_str), Some("application/json") ); + assert_eq!( + plan.headers.get("originator").map(String::as_str), + Some("codex_cli_rs") + ); + assert!(!plan.headers.contains_key("version")); } #[tokio::test] diff --git a/crates/aether-oauth/src/provider/providers/generic.rs b/crates/aether-oauth/src/provider/providers/generic.rs index dbb72d150..b68769cf4 100644 --- a/crates/aether-oauth/src/provider/providers/generic.rs +++ b/crates/aether-oauth/src/provider/providers/generic.rs @@ -461,6 +461,7 @@ fn enrich_generic_identity( "plan_type", "user_id", "account_name", + "is_fedramp", ] { if !auth_config.contains_key(field) { if let Some(value) = object.get(field).cloned() { @@ -506,6 +507,9 @@ fn enrich_generic_identity( .entry("organizations".to_string()) .or_insert(value); } + if let Some(value) = auth.get("chatgpt_account_is_fedramp").cloned() { + auth_config.entry("is_fedramp".to_string()).or_insert(value); + } } if let Some(profile) = claims .get("https://api.openai.com/profile") @@ -599,11 +603,12 @@ fn decode_jwt_claims(token: &str) -> Option> { #[cfg(test)] mod tests { - use super::{template_for_provider_type, GenericProviderOAuthAdapter}; + use super::{enrich_generic_identity, template_for_provider_type, GenericProviderOAuthAdapter}; use crate::network::{OAuthHttpExecutor, OAuthHttpRequest, OAuthHttpResponse}; use crate::provider::ProviderOAuthAdapter; use crate::provider::{ProviderOAuthAccount, ProviderOAuthTransportContext}; use async_trait::async_trait; + use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _}; use serde_json::json; use std::collections::BTreeMap; use std::sync::{Arc, Mutex}; @@ -623,6 +628,26 @@ mod tests { assert!(adapter.capabilities().supports_refresh_token_import); } + #[test] + fn codex_identity_extracts_fedramp_workspace_claim() { + let claims = json!({ + "https://api.openai.com/auth": { + "chatgpt_account_id": "acct-fedramp", + "chatgpt_account_is_fedramp": true + } + }); + let token = format!( + "header.{}.signature", + URL_SAFE_NO_PAD.encode(serde_json::to_vec(&claims).expect("claims should encode")) + ); + let mut auth_config = serde_json::Map::new(); + + enrich_generic_identity("codex", &mut auth_config, &json!({"access_token": token})); + + assert_eq!(auth_config.get("account_id"), Some(&json!("acct-fedramp"))); + assert_eq!(auth_config.get("is_fedramp"), Some(&json!(true))); + } + #[derive(Debug, Clone)] struct StaticExecutor { seen_request: Arc>>, diff --git a/crates/aether-provider-transport/src/auth.rs b/crates/aether-provider-transport/src/auth.rs index 57b41311f..707686ea8 100644 --- a/crates/aether-provider-transport/src/auth.rs +++ b/crates/aether-provider-transport/src/auth.rs @@ -275,7 +275,7 @@ pub fn resolve_local_gemini_auth( } } -pub(crate) fn resolve_local_auth_type_for_transport_format( +pub fn resolve_local_auth_type_for_transport_format( transport: &GatewayProviderTransportSnapshot, ) -> String { let default_auth_type = transport.key.auth_type.trim().to_ascii_lowercase(); diff --git a/crates/aether-provider-transport/src/auth_config.rs b/crates/aether-provider-transport/src/auth_config.rs index 5092a14c4..d5749b5a0 100644 --- a/crates/aether-provider-transport/src/auth_config.rs +++ b/crates/aether-provider-transport/src/auth_config.rs @@ -45,6 +45,7 @@ const IGNORABLE_AUTH_CONFIG_METADATA_KEYS: &[&str] = &[ "access_token_import_temporary", "email", "expires_at", + "is_fedramp", "model_regions", "organizations", "plan_type", diff --git a/crates/aether-provider-transport/src/conversion.rs b/crates/aether-provider-transport/src/conversion.rs index c4fd5ccbb..1e88d4da9 100644 --- a/crates/aether-provider-transport/src/conversion.rs +++ b/crates/aether-provider-transport/src/conversion.rs @@ -354,9 +354,9 @@ fn transport_key_supports_api_format( match transport.key.api_formats.as_deref() { None => true, - Some(formats) => formats - .iter() - .any(|value| aether_ai_formats::api_format_alias_matches(value, endpoint_api_format)), + Some(formats) => formats.iter().any(|value| { + aether_ai_formats::api_format_permission_covers(value, endpoint_api_format) + }), } } @@ -818,6 +818,29 @@ mod tests { ); } + #[test] + fn responses_key_permission_covers_search_without_changing_endpoint_identity() { + let mut transport = transport_snapshot("custom", "openai:search", "bearer", true, None); + transport.key.api_formats = Some(vec!["openai:responses".to_string()]); + + assert_eq!( + candidate_common_transport_skip_reason( + &transport, + candidate_facts("openai:search"), + None, + ), + None + ); + assert_eq!( + candidate_common_transport_skip_reason( + &transport, + candidate_facts("openai:responses"), + None, + ), + Some("endpoint_api_format_changed") + ); + } + #[test] fn candidate_transport_pair_policy_reports_disabled_conversion_and_unsupported_pairs() { let transport = transport_snapshot("custom", "openai:responses", "bearer", false, None); diff --git a/crates/aether-provider-transport/src/network.rs b/crates/aether-provider-transport/src/network.rs index c2678df85..725105706 100644 --- a/crates/aether-provider-transport/src/network.rs +++ b/crates/aether-provider-transport/src/network.rs @@ -454,6 +454,21 @@ mod tests { assert_eq!(timeouts.first_byte_ms, Some(30_000)); } + #[test] + fn transport_execution_timeouts_preserve_the_configurable_maximum() { + let mut transport = sample_transport(); + transport.provider.request_timeout_secs = + Some(aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_SECS as f64); + + let timeouts = resolve_transport_execution_timeouts(&transport) + .expect("provider timeouts should resolve"); + + assert_eq!( + timeouts.total_ms, + Some(aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_MS) + ); + } + #[test] fn transport_execution_timeouts_preserve_configured_first_byte_value() { let mut transport = sample_transport(); diff --git a/crates/aether-provider-transport/src/openai_image/mod.rs b/crates/aether-provider-transport/src/openai_image/mod.rs index 24acfb3dc..6f7030760 100644 --- a/crates/aether-provider-transport/src/openai_image/mod.rs +++ b/crates/aether-provider-transport/src/openai_image/mod.rs @@ -7,7 +7,7 @@ use crate::grok::{is_grok_provider_transport, resolve_grok_session_auth}; use crate::policy::local_standard_transport_unsupported_reason_with_network; use crate::rules::apply_local_header_rules_with_request_headers; use crate::snapshot::GatewayProviderTransportSnapshot; -use crate::url::{build_openai_image_url, build_openai_responses_url}; +use crate::url::build_openai_image_url; #[derive(Debug, Clone, Copy)] pub struct ProviderOpenAiImageHeadersInput<'a> { @@ -15,7 +15,7 @@ pub struct ProviderOpenAiImageHeadersInput<'a> { pub headers: &'a http::HeaderMap, pub auth_header: &'a str, pub auth_value: &'a str, - pub accept: &'a str, + pub accept: Option<&'a str>, pub header_rules: Option<&'a Value>, pub provider_request_body: &'a Value, pub original_request_body: &'a Value, @@ -44,6 +44,11 @@ fn is_dedicated_openai_image_provider(transport: &GatewayProviderTransportSnapsh .provider_type .trim() .eq_ignore_ascii_case("chatgpt_web") + || transport + .provider + .provider_type + .trim() + .eq_ignore_ascii_case("codex") || is_grok_provider_transport(transport) } @@ -61,14 +66,6 @@ pub fn build_openai_image_upstream_url( request_path: Option<&str>, request_query: Option<&str>, ) -> String { - if transport - .provider - .provider_type - .trim() - .eq_ignore_ascii_case("codex") - { - return build_openai_responses_url(&transport.endpoint.base_url, request_query, false); - } build_openai_image_url(&transport.endpoint.base_url, request_path, request_query) } @@ -82,7 +79,11 @@ pub fn build_openai_image_headers( &BTreeMap::new(), ); provider_request_headers.insert("content-type".to_string(), "application/json".to_string()); - provider_request_headers.insert("accept".to_string(), input.accept.to_string()); + if let Some(accept) = input.accept { + provider_request_headers.insert("accept".to_string(), accept.to_string()); + } else { + provider_request_headers.remove("accept"); + } crate::apply_local_auth_config_header_overrides( &mut provider_request_headers, input.transport.key.decrypted_auth_config.as_deref(), @@ -171,14 +172,25 @@ mod tests { } #[test] - fn codex_openai_image_url_stays_on_responses_surface() { + fn codex_openai_image_url_uses_images_surface() { let url = build_openai_image_upstream_url( &sample_transport(), Some("/v1/images/generations"), Some("trace=1"), ); - assert_eq!(url, "https://api.openai.com/v1/responses?trace=1"); + assert_eq!(url, "https://api.openai.com/v1/images/generations?trace=1"); + } + + #[test] + fn codex_openai_image_edit_url_uses_images_edit_surface() { + let url = build_openai_image_upstream_url( + &sample_transport(), + Some("/v1/images/edits"), + Some("trace=1"), + ); + + assert_eq!(url, "https://api.openai.com/v1/images/edits?trace=1"); } #[test] @@ -228,6 +240,19 @@ mod tests { ); } + #[test] + fn codex_is_supported_by_dedicated_openai_image_transport_policy() { + let mut transport = sample_transport(); + transport.provider.provider_type = "codex".to_string(); + transport.key.auth_type = "oauth".to_string(); + transport.key.decrypted_auth_config = Some(json!({"access_token":"token"}).to_string()); + + assert_eq!( + openai_image_transport_unsupported_reason(&transport, "openai:image"), + None + ); + } + #[test] fn grok_oauth_session_is_supported_by_dedicated_openai_image_transport_policy() { let mut transport = sample_transport(); @@ -264,7 +289,7 @@ mod tests { headers: &HeaderMap::new(), auth_header: "authorization", auth_value: "Bearer secret", - accept: "text/event-stream", + accept: Some("text/event-stream"), header_rules: Some(&json!([ {"action":"set","key":"x-image-route","value":"codex"} ])), @@ -307,7 +332,7 @@ mod tests { headers: &HeaderMap::new(), auth_header: "authorization", auth_value: "Bearer refreshed-access-token", - accept: "text/event-stream", + accept: Some("text/event-stream"), header_rules: None, provider_request_body: &json!({"model":"gpt-5.4-mini"}), original_request_body: &json!({"prompt":"draw"}), @@ -332,7 +357,7 @@ mod tests { headers: &HeaderMap::new(), auth_header: "authorization", auth_value: "Bearer secret", - accept: "application/json", + accept: Some("application/json"), header_rules: None, provider_request_body: &json!({ "model": "upstream-image-model", @@ -349,6 +374,31 @@ mod tests { ); } + #[test] + fn codex_images_omits_explicit_accept_header() { + let transport = sample_transport(); + let mut request_headers = HeaderMap::new(); + request_headers.insert("accept", "text/event-stream".parse().expect("valid header")); + let headers = build_openai_image_headers(ProviderOpenAiImageHeadersInput { + transport: &transport, + headers: &request_headers, + auth_header: "authorization", + auth_value: "Bearer secret", + accept: None, + header_rules: Some(&json!([ + {"action":"set","key":"accept","value":"application/json"} + ])), + provider_request_body: &json!({ + "model": "gpt-image-2", + "prompt": "draw a city", + }), + original_request_body: &json!({"prompt":"draw a city"}), + }) + .expect("headers should build"); + + assert!(!headers.contains_key("accept")); + } + #[test] fn standard_openai_compatible_image_url_supports_aether_api_root() { let mut transport = sample_transport(); diff --git a/crates/aether-provider-transport/src/provider_types.rs b/crates/aether-provider-transport/src/provider_types.rs index 8a3bc663b..f8efa4eaf 100644 --- a/crates/aether-provider-transport/src/provider_types.rs +++ b/crates/aether-provider-transport/src/provider_types.rs @@ -306,11 +306,17 @@ const CODEX_FIXED_PROVIDER_TEMPLATE: FixedProviderTemplate = FixedProviderTempla custom_path: None, config_defaults: EMPTY_ENDPOINT_CONFIG_DEFAULTS, }, + FixedProviderEndpointTemplate { + item_key: "openai:search", + api_format: "openai:search", + custom_path: None, + config_defaults: EMPTY_ENDPOINT_CONFIG_DEFAULTS, + }, FixedProviderEndpointTemplate { item_key: "openai:image", api_format: "openai:image", custom_path: None, - config_defaults: FORCE_STREAM_ENDPOINT_CONFIG_DEFAULTS, + config_defaults: EMPTY_ENDPOINT_CONFIG_DEFAULTS, }, ], runtime_policy: CODEX_RUNTIME_POLICY, @@ -639,7 +645,7 @@ mod tests { }; #[test] - fn codex_fixed_provider_template_includes_openai_image() { + fn codex_fixed_provider_template_includes_codex_companion_endpoints() { let template = fixed_provider_template("codex").expect("codex template should exist"); assert_eq!(template.base_url, "https://chatgpt.com/backend-api/codex"); assert_eq!(template.version, 1); @@ -652,6 +658,7 @@ mod tests { vec![ "openai:responses", "openai:responses:compact", + "openai:search", "openai:image" ] ); @@ -659,17 +666,12 @@ mod tests { let image_template = fixed_provider_endpoint_template_by_api_format("codex", "openai:image") .expect("codex image endpoint should exist"); - assert_eq!( - image_template - .config_defaults - .iter() - .map(|item| (item.key, item.value)) - .collect::>(), - vec![( - "upstream_stream_policy", - FixedProviderEndpointConfigValue::String("force_stream") - )] - ); + assert!(image_template.config_defaults.is_empty()); + + let search_template = + fixed_provider_endpoint_template_by_api_format("codex", "openai:search") + .expect("codex search endpoint should exist"); + assert!(search_template.config_defaults.is_empty()); } #[test] diff --git a/crates/aether-provider-transport/src/request_url/mod.rs b/crates/aether-provider-transport/src/request_url/mod.rs index 208aea3cd..ab77d7e4b 100644 --- a/crates/aether-provider-transport/src/request_url/mod.rs +++ b/crates/aether-provider-transport/src/request_url/mod.rs @@ -16,7 +16,8 @@ use crate::gemini_cli::{ use crate::snapshot::GatewayProviderTransportSnapshot; use crate::url::{ build_claude_messages_url, build_gemini_content_url, build_openai_chat_url, - build_openai_responses_url, build_passthrough_path_url, normalize_gemini_content_action_path, + build_openai_responses_url, build_openai_search_url, build_passthrough_path_url, + normalize_gemini_content_action_path, }; use crate::vertex::{ build_vertex_api_key_gemini_content_url, build_vertex_api_key_gemini_embedding_url, @@ -124,6 +125,10 @@ fn build_transport_request_url_inner( params.request_query, true, )), + "openai:search" => Some(build_openai_search_url( + &transport.endpoint.base_url, + params.request_query, + )), "openai:embedding" | "jina:embedding" => { build_provider_embedding_v1_url(&transport.endpoint.base_url, params.request_query) } @@ -868,6 +873,33 @@ mod tests { assert_eq!(url, "https://api.openai.example/v1/responses?tenant=demo"); } + #[test] + fn builds_openai_search_url_for_codex_provider_root() { + let transport = sample_transport( + "codex", + "openai:search", + "https://chatgpt.com/backend-api/codex", + None, + ); + + let url = build_transport_request_url( + &transport, + TransportRequestUrlParams { + provider_api_format: "openai:search", + mapped_model: Some("gpt-5.6-luna"), + upstream_is_stream: false, + request_query: Some("tenant=demo"), + kiro_api_region: None, + }, + ) + .expect("openai search url"); + + assert_eq!( + url, + "https://chatgpt.com/backend-api/codex/alpha/search?tenant=demo" + ); + } + #[test] fn expands_custom_path_templates_when_hook_does_not_apply() { let transport = sample_transport( 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 d65a24519..8dc7c2190 100644 --- a/crates/aether-provider-transport/src/same_format_provider/mod.rs +++ b/crates/aether-provider-transport/src/same_format_provider/mod.rs @@ -142,20 +142,18 @@ 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_from_endpoint_config( + 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 - || aether_ai_formats::api::force_upstream_streaming_for_provider( - transport.provider.provider_type.as_str(), - params.provider_api_format, - ), - ); - let force_body_stream_field = aether_ai_formats::endpoint_config_forces_upstream_stream_policy( - transport.endpoint.config.as_ref(), + is_kiro || is_antigravity || gemini_cli_requires_upstream_streaming, ); + 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 report_kind = if is_kiro && !params.require_streaming { "claude_cli_sync_finalize" } else if (is_gemini_cli || is_antigravity) && !params.require_streaming { @@ -238,10 +236,11 @@ fn build_same_format_provider_request_body_inner( .map(|(key, value)| (key.clone(), value.clone())), ) } else { - aether_ai_formats::convert_request_pure( + aether_ai_formats::convert_request_pure_with_context( input.client_api_format, input.provider_api_format, input.body_json, + &aether_ai_formats::FormatContext::default().with_mapped_model(input.mapped_model), ) .ok()? .value @@ -397,6 +396,30 @@ fn build_same_format_provider_request_body_inner( ), ); } + let provider_model = provider_request_body + .get("model") + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .unwrap_or(input.mapped_model) + .to_string(); + if aether_ai_formats::finalize_openai_provider_request( + &mut provider_request_body, + aether_ai_formats::OpenAiProviderRequestFinalization { + source_api_format: input.provider_api_format, + provider_api_format: input.provider_api_format, + provider_type: "", + provider_model: &provider_model, + source_model: input.source_model.unwrap_or(input.mapped_model), + body_rules: input.body_rules, + upstream_is_stream: input.upstream_is_stream, + require_body_stream_field, + }, + ) + .is_err() + { + return None; + } Some(provider_request_body) } @@ -601,6 +624,7 @@ pub fn same_format_provider_transport_unsupported_reason_for_trace( "openai:chat" => "openai:chat", "openai:responses" => "openai:responses", "openai:responses:compact" => "openai:responses:compact", + "openai:search" => "openai:search", "claude:messages" => "claude:messages", "gemini:generate_content" => "gemini:generate_content", "gemini:interactions" => "gemini:interactions", @@ -681,7 +705,9 @@ fn resolve_same_format_standard_direct_auth( transport: &GatewayProviderTransportSnapshot, provider_api_format: &str, ) -> Option<(String, String)> { - if aether_ai_formats::api_format_alias_matches(provider_api_format, "openai:embedding") { + if aether_ai_formats::api_format_alias_matches(provider_api_format, "openai:embedding") + || aether_ai_formats::api_format_alias_matches(provider_api_format, "openai:search") + { resolve_local_openai_bearer_auth(transport) } else { resolve_local_standard_auth(transport) @@ -851,6 +877,35 @@ mod tests { }, ); assert!(!sync_behavior.upstream_is_stream); + + let mut compact = sample_transport("codex"); + compact.endpoint.config = Some(json!({ + "upstream_stream_policy": "force_stream" + })); + let compact_behavior = classify_same_format_provider_request_behavior( + &compact, + SameFormatProviderRequestBehaviorParams { + require_streaming: true, + provider_api_format: "openai:responses:compact", + report_kind: "openai_responses_compact_sync_success", + }, + ); + assert!(!compact_behavior.upstream_is_stream); + + let mut search = sample_transport("codex"); + search.endpoint.config = Some(json!({ + "upstream_stream_policy": "force_stream" + })); + let search_behavior = classify_same_format_provider_request_behavior( + &search, + SameFormatProviderRequestBehaviorParams { + require_streaming: true, + provider_api_format: "openai:search", + report_kind: "openai_search_sync_success", + }, + ); + assert!(!search_behavior.upstream_is_stream); + assert!(!search_behavior.force_body_stream_field); } #[test] @@ -1000,6 +1055,31 @@ mod tests { ); } + #[test] + fn resolves_openai_search_direct_auth_with_bearer_header() { + let mut transport = sample_transport("custom"); + transport.endpoint.api_format = "openai:search".to_string(); + transport.key.auth_type = "api_key".to_string(); + let behavior = classify_same_format_provider_request_behavior( + &transport, + SameFormatProviderRequestBehaviorParams { + require_streaming: false, + provider_api_format: "openai:search", + report_kind: "openai_search_sync_success", + }, + ); + + assert_eq!( + resolve_same_format_provider_direct_auth( + &behavior, + &transport, + SameFormatProviderFamily::Standard, + "openai:search", + ), + Some(("authorization".to_string(), "Bearer secret".to_string())) + ); + } + #[test] fn keeps_claude_same_format_api_key_on_x_api_key_header() { let mut transport = sample_transport("custom"); @@ -1081,6 +1161,227 @@ mod tests { assert_eq!(body["unknown_vendor_field"]["keep"], true); } + #[test] + fn same_format_responses_body_preserves_gpt_5_6_protocol_contract_verbatim() { + let input = json!([ + { + "type": "message", + "role": "user", + "content": [ + { + "type": "input_text", + "text": "before", + "prompt_cache_breakpoint": {"mode": "explicit"} + }, + {"type": "input_text", "text": "after"} + ] + }, + { + "type": "program", + "call_id": "program-call-1", + "fingerprint": "fp-program-1" + }, + { + "type": "program_output", + "call_id": "program-call-1", + "result": "hello", + "status": "completed" + }, + { + "type": "multi_agent_call", + "call_id": "agent-call-1", + "action": "delegate", + "arguments": {"query": "release notes"}, + "agent": "researcher" + }, + { + "type": "agent_message", + "author": "researcher", + "recipient": "assistant", + "encrypted_content": "encrypted-agent-message" + }, + { + "type": "function_call_output", + "call_id": "call-1", + "output": "ok", + "caller": {"type": "multi_agent", "id": "agent-call-1"} + } + ]); + let prompt_cache_options = json!({ + "mode": "explicit", + "ttl": "30m" + }); + let body_json = json!({ + "model": "gpt-5.6-sol", + "input": input, + "reasoning": { + "effort": "max", + "mode": "pro", + "context": "all_turns" + }, + "prompt_cache_options": prompt_cache_options, + "multi_agent": { + "enabled": true, + "max_concurrent_subagents": 2 + } + }); + + let body = build_same_format_provider_request_body(SameFormatProviderRequestBodyInput { + body_json: &body_json, + mapped_model: "gpt-5.6-sol", + client_api_format: "openai:responses", + provider_api_format: "openai:responses", + source_model: Some("gpt-5.6-sol"), + 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: false, + enable_model_directives: false, + }) + .expect("valid GPT-5.6 Responses request should use raw same-format transport"); + + assert_eq!(body["input"], input); + assert_eq!(body["reasoning"], body_json["reasoning"]); + assert_eq!(body["prompt_cache_options"], prompt_cache_options); + assert_eq!(body["multi_agent"], body_json["multi_agent"]); + } + + #[test] + fn same_format_prompt_cache_capability_uses_mapped_provider_model() { + let body_json = json!({ + "model": "gpt-5.5", + "messages": [{"role": "user", "content": "hello"}], + "prompt_cache_options": {"mode": "explicit", "ttl": "30m"} + }); + let body = build_same_format_provider_request_body(SameFormatProviderRequestBodyInput { + body_json: &body_json, + mapped_model: "gpt-5.6-sol", + client_api_format: "openai:chat", + provider_api_format: "openai:chat", + source_model: Some("gpt-5.5"), + 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: false, + enable_model_directives: false, + }) + .expect("final GPT-5.6 provider model should accept prompt_cache_options"); + + assert_eq!(body["model"], "gpt-5.6-sol"); + assert_eq!( + body["prompt_cache_options"], + body_json["prompt_cache_options"] + ); + + let rejected = + build_same_format_provider_request_body(SameFormatProviderRequestBodyInput { + body_json: &body_json, + mapped_model: "gpt-5.5", + client_api_format: "openai:chat", + provider_api_format: "openai:chat", + source_model: Some("gpt-5.5"), + 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: false, + enable_model_directives: false, + }); + assert!(rejected.is_none()); + } + + #[test] + fn same_format_reasoning_capability_uses_mapped_provider_model() { + let body_json = json!({ + "model": "deployment-alias", + "input": "hello", + "reasoning": {"effort": "max"} + }); + let build = |mapped_model: &str| { + build_same_format_provider_request_body(SameFormatProviderRequestBodyInput { + body_json: &body_json, + mapped_model, + client_api_format: "openai:responses", + provider_api_format: "openai:responses", + source_model: Some("deployment-alias"), + 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: false, + enable_model_directives: false, + }) + }; + + assert!(build("gpt-5.6-sol").is_some()); + assert!(build("gpt-5.4").is_none()); + } + + #[test] + fn same_format_gpt_5_6_requires_prompt_cache_options() { + let body_json = json!({ + "model": "client-alias", + "input": "hello", + "prompt_cache_retention": "24h" + }); + let body = build_same_format_provider_request_body(SameFormatProviderRequestBodyInput { + body_json: &body_json, + mapped_model: "gpt-5.6-luna", + client_api_format: "openai:responses", + provider_api_format: "openai:responses", + source_model: Some("client-alias"), + 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: false, + enable_model_directives: false, + }); + + assert!(body.is_none()); + } + + #[test] + fn same_format_responses_body_preserves_opaque_extension_fields() { + let opaque_option = json!({"mode": "custom"}); + let body = build_same_format_provider_request_body(SameFormatProviderRequestBodyInput { + body_json: &json!({ + "model": "gpt-5.6-sol", + "input": "hello", + "multi_agent": {"enabled": true}, + "opaque_responses_option": opaque_option + }), + mapped_model: "gpt-5.6-sol", + client_api_format: "openai:responses", + provider_api_format: "openai:responses", + source_model: Some("gpt-5.6-sol"), + 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: false, + enable_model_directives: false, + }); + + let body = body.expect("same-format Responses request should remain transparent"); + assert_eq!(body["multi_agent"], json!({"enabled": true})); + assert_eq!(body["opaque_responses_option"], opaque_option); + } + #[test] fn cross_format_standard_body_fails_closed_for_lossy_chat_fields() { let body = build_same_format_provider_request_body(SameFormatProviderRequestBodyInput { @@ -1514,6 +1815,44 @@ mod tests { assert_eq!(body["metadata"]["body_rule_seen"], true); } + #[test] + fn search_body_projects_the_protocol_contract_and_keeps_fast_routing_only() { + let body = build_same_format_provider_request_body(SameFormatProviderRequestBodyInput { + body_json: &json!({ + "id": "session-1", + "model": "gpt-5.6-luna-high-fast", + "commands": {"search_query": [{"q": "Aether"}]}, + "settings": {"allowed_callers": ["direct"]}, + "store": false, + "future_extension": {"enabled": true}, + "stream": true + }), + mapped_model: "gpt-5.6-luna", + client_api_format: "openai:search", + provider_api_format: "openai:search", + source_model: Some("gpt-5.6-luna-high-fast"), + 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: false, + enable_model_directives: true, + }) + .expect("search body should build"); + + assert_eq!(body["id"], "session-1"); + assert_eq!(body["model"], "gpt-5.6-luna"); + assert_eq!(body["commands"]["search_query"][0]["q"], "Aether"); + assert_eq!(body["settings"]["allowed_callers"][0], "direct"); + assert_eq!(body["reasoning"]["effort"], "high"); + assert!(body.get("store").is_none()); + assert!(body.get("future_extension").is_none()); + assert!(body.get("service_tier").is_none()); + assert!(body.get("stream").is_none()); + } + #[test] fn builds_same_format_headers_with_auth_and_stream_accept() { let provider_request_body = json!({"model": "upstream-model"}); diff --git a/crates/aether-provider-transport/src/url.rs b/crates/aether-provider-transport/src/url.rs index 14661413b..07af0239b 100644 --- a/crates/aether-provider-transport/src/url.rs +++ b/crates/aether-provider-transport/src/url.rs @@ -28,6 +28,14 @@ pub fn build_openai_responses_url( url } +pub fn build_openai_search_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('/'); + let mut url = format!("{trimmed}/alpha/search"); + append_merged_query(&mut url, base_query, None, query, &[]); + url +} + pub fn build_openai_image_url( upstream_base_url: &str, request_path: Option<&str>, @@ -424,7 +432,7 @@ mod tests { 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_passthrough_path_url, + build_openai_responses_url, build_openai_search_url, build_passthrough_path_url, normalize_gemini_content_action_path, }; @@ -573,6 +581,21 @@ mod tests { ); } + #[test] + fn openai_search_url_preserves_api_and_codex_roots() { + assert_eq!( + build_openai_search_url( + "https://api.openai.com/v1?tenant=base", + Some("trace=1&tenant=request") + ), + "https://api.openai.com/v1/alpha/search?tenant=request&trace=1" + ); + assert_eq!( + build_openai_search_url("https://chatgpt.com/backend-api/codex/", None), + "https://chatgpt.com/backend-api/codex/alpha/search" + ); + } + #[test] fn openai_image_url_uses_images_surface() { assert_eq!( diff --git a/crates/aether-scheduler-core/src/affinity.rs b/crates/aether-scheduler-core/src/affinity.rs index a46ac6385..6b49b5914 100644 --- a/crates/aether-scheduler-core/src/affinity.rs +++ b/crates/aether-scheduler-core/src/affinity.rs @@ -165,6 +165,7 @@ mod tests { global_model_id: format!("global-model-{id}"), global_model_name: "gpt-5".to_string(), selected_provider_model_name: "gpt-5".to_string(), + supports_streaming: true, mapping_matched_model: None, } } diff --git a/crates/aether-scheduler-core/src/auth.rs b/crates/aether-scheduler-core/src/auth.rs index 659d70f5b..086b3357d 100644 --- a/crates/aether-scheduler-core/src/auth.rs +++ b/crates/aether-scheduler-core/src/auth.rs @@ -56,7 +56,7 @@ pub fn api_format_matches_allowed_value(allowed_value: &str, api_format: &str) - if allowed_value.is_empty() || api_format.is_empty() { return false; } - crate::normalize_api_format(allowed_value) == crate::normalize_api_format(api_format) + aether_ai_formats::api_format_permission_covers(allowed_value, api_format) } pub fn auth_constraints_allow_model( @@ -248,6 +248,18 @@ mod tests { "openai:responses", "openai:responses" )); + assert!(api_format_matches_allowed_value( + "openai:responses", + "openai:search" + )); + assert!(api_format_matches_allowed_value( + "openai:search", + "openai:search" + )); + assert!(!api_format_matches_allowed_value( + "openai:search", + "openai:responses" + )); assert!(!api_format_matches_allowed_value( "openai:responses", "claude:messages" diff --git a/crates/aether-scheduler-core/src/candidate/enumeration.rs b/crates/aether-scheduler-core/src/candidate/enumeration.rs index b291a71e9..07472e545 100644 --- a/crates/aether-scheduler-core/src/candidate/enumeration.rs +++ b/crates/aether-scheduler-core/src/candidate/enumeration.rs @@ -73,6 +73,7 @@ fn enumerate_minimal_candidate_selection_inner( continue; }; + let supports_streaming = row.supports_streaming(); candidates.push(SchedulerMinimalCandidateSelectionCandidate { provider_id: row.provider_id, provider_name: row.provider_name, @@ -93,6 +94,7 @@ fn enumerate_minimal_candidate_selection_inner( global_model_id: row.global_model_id, global_model_name: row.global_model_name, selected_provider_model_name, + supports_streaming, mapping_matched_model, }); } diff --git a/crates/aether-scheduler-core/src/candidate/mod.rs b/crates/aether-scheduler-core/src/candidate/mod.rs index b3cc3f688..579102284 100644 --- a/crates/aether-scheduler-core/src/candidate/mod.rs +++ b/crates/aether-scheduler-core/src/candidate/mod.rs @@ -98,6 +98,7 @@ mod tests { global_model_id: format!("global-model-{id}"), global_model_name: "gpt-5".to_string(), selected_provider_model_name: "gpt-5".to_string(), + supports_streaming: true, mapping_matched_model: None, } } @@ -211,6 +212,27 @@ mod tests { assert_eq!(candidates[0].selected_provider_model_name, "gpt-5-canary-1"); } + #[test] + fn enumeration_preserves_effective_streaming_capability() { + let mut row = sample_row("1"); + row.model_supports_streaming = Some(false); + + let candidates = + super::enumerate_minimal_candidate_selection(EnumerateMinimalCandidateSelectionInput { + rows: vec![row], + normalized_api_format: "openai:chat", + requested_model_name: "gpt-5", + resolved_global_model_name: "gpt-5", + require_streaming: false, + required_capabilities: None, + auth_constraints: None, + }) + .expect("candidate selection should build"); + + assert_eq!(candidates.len(), 1); + assert!(!candidates[0].supports_streaming); + } + #[test] fn enumeration_preserves_theoretical_candidate_order_without_final_sorting() { let mut later_priority = sample_row("1"); diff --git a/crates/aether-scheduler-core/src/candidate/types.rs b/crates/aether-scheduler-core/src/candidate/types.rs index 86f3db2aa..dc6191e08 100644 --- a/crates/aether-scheduler-core/src/candidate/types.rs +++ b/crates/aether-scheduler-core/src/candidate/types.rs @@ -25,6 +25,7 @@ pub struct SchedulerMinimalCandidateSelectionCandidate { pub global_model_id: String, pub global_model_name: String, pub selected_provider_model_name: String, + pub supports_streaming: bool, pub mapping_matched_model: Option, } diff --git a/crates/aether-scheduler-core/src/model.rs b/crates/aether-scheduler-core/src/model.rs index 044b15a90..f143645f2 100644 --- a/crates/aether-scheduler-core/src/model.rs +++ b/crates/aether-scheduler-core/src/model.rs @@ -279,7 +279,7 @@ fn mapping_scope_matches( let api_format_matches_scope = mapping.api_formats.as_ref().is_none_or(|api_formats| { api_formats .iter() - .any(|value| api_format_matches(value, api_format)) + .any(|value| api_format_scope_covers(value, api_format)) }); if !api_format_matches_scope { return false; @@ -419,6 +419,10 @@ fn api_format_matches(left: &str, right: &str) -> bool { normalize_api_format(left) == normalize_api_format(right) } +fn api_format_scope_covers(allowed: &str, requested: &str) -> bool { + aether_ai_formats::api_format_permission_covers(allowed, requested) +} + fn requested_model_name_candidates( requested_model_name: &str, enable_model_directives: bool, @@ -568,6 +572,41 @@ mod tests { ); } + #[test] + fn responses_model_mapping_scope_covers_search_in_one_direction() { + let mut row = sample_row("search-global", "search-default"); + row.endpoint_api_format = "openai:search".to_string(); + row.model_provider_model_mappings = Some(vec![StoredProviderModelMapping { + name: "gpt-5.6-sol".to_string(), + priority: 1, + api_formats: Some(vec!["openai:responses".to_string()]), + endpoint_ids: None, + }]); + + assert!(row_supports_requested_model( + &row, + "gpt-5.6-sol", + "openai:search" + )); + assert_eq!( + resolve_provider_model_name(&row, "gpt-5.6-sol", "openai:search") + .map(|resolved| resolved.0), + Some("gpt-5.6-sol".to_string()) + ); + + row.model_provider_model_mappings = Some(vec![StoredProviderModelMapping { + name: "search-only".to_string(), + priority: 1, + api_formats: Some(vec!["openai:search".to_string()]), + endpoint_ids: None, + }]); + assert!(!row_supports_requested_model( + &row, + "search-only", + "openai:responses" + )); + } + #[test] fn model_directive_suffix_prefers_exact_model_before_base_fallback() { let exact = sample_row("gpt-5.4-high", "gpt-5.4-high-upstream"); diff --git a/crates/aether-usage-runtime/src/record.rs b/crates/aether-usage-runtime/src/record.rs index 828350680..bcd9bca16 100644 --- a/crates/aether-usage-runtime/src/record.rs +++ b/crates/aether-usage-runtime/src/record.rs @@ -40,6 +40,11 @@ pub fn build_upsert_usage_record_from_event( let mut data = event.data.clone(); data.request_metadata = attach_provider_request_body_metadata( data.request_metadata, + data.endpoint_api_format + .as_deref() + .or(data.api_format.as_deref()), + data.target_model.as_deref().or(Some(data.model.as_str())), + Some(data.model.as_str()), data.provider_request_body.as_ref(), ); let now_unix_secs = event.timestamp_ms / 1_000; @@ -185,6 +190,9 @@ mod tests { "reasoning": { "effort": "max" }, "service_tier": "priority" })), + request_metadata: Some(serde_json::json!({ + "provider_actual_service_tier": "default" + })), ..UsageEventData::default() }, }) @@ -210,6 +218,14 @@ mod tests { .and_then(serde_json::Value::as_str), Some("priority") ); + assert_eq!( + record + .request_metadata + .as_ref() + .and_then(|value| value.get("provider_actual_service_tier")) + .and_then(serde_json::Value::as_str), + Some("default") + ); assert_eq!(record.finalized_at_unix_secs, Some(1_700_000_000)); } diff --git a/crates/aether-usage-runtime/src/report.rs b/crates/aether-usage-runtime/src/report.rs index 82997d223..9f7fca67b 100644 --- a/crates/aether-usage-runtime/src/report.rs +++ b/crates/aether-usage-runtime/src/report.rs @@ -109,6 +109,9 @@ pub fn infer_internal_finalize_signature(payload: &GatewaySyncReportRequest) -> if report_kind.starts_with("openai_image_") { return Some("openai:image".to_string()); } + if report_kind.starts_with("openai_search_") { + return Some("openai:search".to_string()); + } if report_kind.starts_with("openai_cli_") { return Some("openai:responses".to_string()); } @@ -150,6 +153,11 @@ pub fn resolve_internal_finalize_route(signature: &str) -> Option Some(InternalFinalizeRoute { + public_path: "/v1/alpha/search", + route_family: "openai", + route_kind: "search", + }), "openai:image" => Some(InternalFinalizeRoute { public_path: "/v1/images/generations", route_family: "openai", @@ -246,6 +254,7 @@ pub fn is_local_ai_sync_report_kind(report_kind: &str) -> bool { | "openai_responses_compact_sync_error" | "openai_cli_sync_success" | "openai_image_sync_success" + | "openai_search_sync_success" | "openai_image_sync_error" | "openai_embedding_sync_success" | "openai_embedding_sync_error" @@ -872,6 +881,7 @@ mod tests { "openai_responses_compact_sync_error" )); assert!(is_local_ai_sync_report_kind("openai_image_sync_success")); + assert!(is_local_ai_sync_report_kind("openai_search_sync_success")); assert!(is_local_ai_sync_report_kind("openai_image_sync_error")); assert!(is_local_ai_sync_report_kind( "openai_embedding_sync_success" @@ -1082,12 +1092,27 @@ mod tests { Some("openai:responses:compact".to_string()) ); + let from_search_report_kind = + sample_sync_report_with_context("openai_search_sync_finalize", json!({})); + assert_eq!( + infer_internal_finalize_signature(&from_search_report_kind), + Some("openai:search".to_string()) + ); + let unknown = sample_sync_report("unknown_sync_finalize", 200); assert_eq!(infer_internal_finalize_signature(&unknown), None); } #[test] fn resolves_internal_finalize_route_for_supported_signatures() { + assert_eq!( + resolve_internal_finalize_route("openai:search"), + Some(InternalFinalizeRoute { + public_path: "/v1/alpha/search", + route_family: "openai", + route_kind: "search", + }) + ); assert_eq!( resolve_internal_finalize_route("openai:responses:compact"), Some(InternalFinalizeRoute { diff --git a/crates/aether-usage-runtime/src/request_metadata.rs b/crates/aether-usage-runtime/src/request_metadata.rs index 14d6d3c63..4c752e4b1 100644 --- a/crates/aether-usage-runtime/src/request_metadata.rs +++ b/crates/aether-usage-runtime/src/request_metadata.rs @@ -4,7 +4,10 @@ use aether_ai_formats::api::{ use aether_ai_formats::UPSTREAM_IS_STREAM_KEY; use aether_contracts::ExecutionPlan; use aether_data_contracts::repository::usage::{ + extract_provider_actual_service_tier_from_response, extract_provider_reasoning_effort_from_body, extract_provider_service_tier_from_body, + normalize_provider_service_tier, resolve_provider_cache_ttl_minutes, + PROVIDER_ACTUAL_SERVICE_TIER_METADATA_KEY, PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY, PROVIDER_REASONING_EFFORT_METADATA_KEY, PROVIDER_SERVICE_TIER_METADATA_KEY, }; use serde_json::{json, Map, Value}; @@ -75,12 +78,25 @@ pub(crate) fn sanitize_usage_request_metadata_ref(value: Option<&Value>) -> Opti pub(crate) fn attach_provider_request_body_metadata( metadata: Option, + provider_api_format: Option<&str>, + provider_model: Option<&str>, + source_model: Option<&str>, provider_request_body: Option<&Value>, ) -> Option { let provider_body_is_object = provider_request_body.and_then(Value::as_object).is_some(); let reasoning_effort = extract_provider_reasoning_effort_from_body(provider_request_body); let service_tier = extract_provider_service_tier_from_body(provider_request_body); - if !provider_body_is_object && reasoning_effort.is_none() && service_tier.is_none() { + let cache_ttl_minutes = resolve_provider_cache_ttl_minutes( + provider_api_format, + provider_model, + source_model, + provider_request_body, + ); + if !provider_body_is_object + && reasoning_effort.is_none() + && service_tier.is_none() + && cache_ttl_minutes.is_none() + { return metadata; } let mut object = match metadata { @@ -90,6 +106,7 @@ pub(crate) fn attach_provider_request_body_metadata( if provider_body_is_object { object.remove(PROVIDER_REASONING_EFFORT_METADATA_KEY); object.remove(PROVIDER_SERVICE_TIER_METADATA_KEY); + object.remove(PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY); } if let Some(reasoning_effort) = reasoning_effort { object.insert( @@ -103,6 +120,50 @@ pub(crate) fn attach_provider_request_body_metadata( Value::String(service_tier), ); } + if let Some(cache_ttl_minutes) = cache_ttl_minutes { + object.insert( + PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY.to_string(), + Value::Number(cache_ttl_minutes.into()), + ); + } + (!object.is_empty()).then_some(Value::Object(object)) +} + +pub(crate) fn attach_provider_response_body_metadata( + metadata: Option, + provider_response_body: Option<&Value>, +) -> Option { + if metadata + .as_ref() + .and_then(Value::as_object) + .and_then(|object| object.get(PROVIDER_ACTUAL_SERVICE_TIER_METADATA_KEY)) + .and_then(Value::as_str) + .and_then(normalize_provider_service_tier) + .is_some() + { + return metadata; + } + let actual_service_tier = + extract_provider_actual_service_tier_from_response(provider_response_body); + attach_provider_actual_service_tier_metadata(metadata, actual_service_tier.as_deref()) +} + +pub(crate) fn attach_provider_actual_service_tier_metadata( + metadata: Option, + actual_service_tier: Option<&str>, +) -> Option { + let Some(actual_service_tier) = actual_service_tier.and_then(normalize_provider_service_tier) + else { + return metadata; + }; + let mut object = match metadata { + Some(Value::Object(object)) => object, + _ => Map::new(), + }; + object.insert( + PROVIDER_ACTUAL_SERVICE_TIER_METADATA_KEY.to_string(), + Value::String(actual_service_tier), + ); (!object.is_empty()).then_some(Value::Object(object)) } @@ -120,6 +181,8 @@ fn copy_allowed_metadata_fields(source: &Map, target: &mut Map, target: &mut Map remove_non_empty_string(&mut source, target, "request_path_and_query"); remove_non_empty_string(&mut source, target, PROVIDER_REASONING_EFFORT_METADATA_KEY); remove_non_empty_string(&mut source, target, PROVIDER_SERVICE_TIER_METADATA_KEY); + remove_non_empty_string( + &mut source, + target, + PROVIDER_ACTUAL_SERVICE_TIER_METADATA_KEY, + ); + remove_number(&mut source, target, PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY); remove_number(&mut source, target, "provider_request_body_base64_bytes"); remove_number(&mut source, target, "provider_response_body_base64_bytes"); remove_number(&mut source, target, "client_response_body_base64_bytes"); @@ -449,7 +518,8 @@ mod tests { use std::collections::BTreeMap; use super::{ - attach_provider_request_body_metadata, build_usage_request_metadata_seed, + attach_provider_actual_service_tier_metadata, attach_provider_request_body_metadata, + attach_provider_response_body_metadata, build_usage_request_metadata_seed, merge_usage_request_metadata, merge_usage_request_metadata_owned, sanitize_usage_request_metadata, sanitize_usage_request_metadata_ref, MAX_USAGE_REQUEST_METADATA_BYTES, MAX_USAGE_REQUEST_METADATA_DEPTH, @@ -770,8 +840,11 @@ mod tests { let updated = attach_provider_request_body_metadata( metadata.clone(), + Some("openai:responses"), + Some("gpt-5.6-sol"), + Some("gpt-5.6-sol"), Some(&json!({ - "model": "gpt-5", + "model": "gpt-5.6-sol", "reasoning": { "effort": "low" }, "service_tier": "standard" })), @@ -783,12 +856,16 @@ mod tests { json!({ "trace_id": "trace-1", "provider_reasoning_effort": "low", - "provider_service_tier": "standard" + "provider_service_tier": "standard", + "provider_cache_ttl_minutes": 30 }) ); let cleared = attach_provider_request_body_metadata( metadata, + Some("openai:responses"), + Some("gpt-5"), + Some("gpt-5"), Some(&json!({ "model": "gpt-5" })), @@ -803,6 +880,59 @@ mod tests { ); } + #[test] + fn provider_response_metadata_preserves_terminal_actual_service_tier() { + let metadata = attach_provider_response_body_metadata( + Some(json!({"provider_service_tier": "priority"})), + Some(&json!({ + "chunks": [ + {"service_tier": "priority"}, + {"service_tier": "Default", "usage": {"total_tokens": 12}} + ] + })), + ) + .expect("requested and actual provider tiers should be preserved"); + + assert_eq!( + metadata, + json!({ + "provider_service_tier": "priority", + "provider_actual_service_tier": "default" + }) + ); + } + + #[test] + fn terminal_summary_tier_uses_the_same_normalized_metadata_field() { + let metadata = attach_provider_actual_service_tier_metadata( + Some(json!({"trace_id": "trace-1"})), + Some(" Flex "), + ) + .expect("terminal summary tier should be retained"); + + assert_eq!( + metadata, + json!({ + "trace_id": "trace-1", + "provider_actual_service_tier": "flex" + }) + ); + } + + #[test] + fn terminal_summary_tier_precedes_truncated_response_capture() { + let metadata = attach_provider_response_body_metadata( + Some(json!({"provider_actual_service_tier": "default"})), + Some(&json!({"chunks": [{"service_tier": "priority"}]})), + ) + .expect("terminal summary tier should remain"); + + assert_eq!( + metadata.get("provider_actual_service_tier"), + Some(&Value::String("default".to_string())) + ); + } + #[test] fn owned_merge_matches_filtered_merge_for_trusted_objects() { let base = Some(json!({ @@ -811,7 +941,8 @@ mod tests { })); let override_value = Some(json!({ "billing_snapshot_status": "complete", - "trace_id": "trace-2" + "trace_id": "trace-2", + "provider_actual_service_tier": "default" })); assert_eq!( diff --git a/crates/aether-usage-runtime/src/runtime.rs b/crates/aether-usage-runtime/src/runtime.rs index e342ab12b..38745b131 100644 --- a/crates/aether-usage-runtime/src/runtime.rs +++ b/crates/aether-usage-runtime/src/runtime.rs @@ -13,6 +13,7 @@ use tokio::sync::mpsc; use tracing::{info, warn}; use crate::executor::spawn_on_usage_background_runtime; +use crate::request_metadata::attach_provider_response_body_metadata; use crate::worker::{ build_usage_queue_worker_with_record_gate, UsageWorkerControl, UsageWorkerObservation, }; @@ -921,6 +922,7 @@ impl UsageRuntime { where T: UsageRuntimeAccess, { + preserve_provider_response_facts(event); match self.cached_body_capture_policy(data).await { Ok(policy) => apply_usage_body_capture_policy_to_event(policy, event), Err(err) => { @@ -1200,6 +1202,12 @@ impl UsageRuntime { } } +fn preserve_provider_response_facts(event: &mut UsageEvent) { + let metadata = event.data.request_metadata.take(); + event.data.request_metadata = + attach_provider_response_body_metadata(metadata, event.data.response_body.as_ref()); +} + impl UsageQueueHealthSnapshot { fn apply_stream_stats(&mut self, stats: RuntimeQueueStats) { self.stream_length = stats.stream_length; @@ -2318,8 +2326,9 @@ mod tests { use tokio::time::{sleep, Duration}; use super::{ - UsageBillingEventEnricher, UsageBodyCapturePolicy, UsageRequestRecordLevel, - UsageRuntimeAccess, UsageWorkerObservation, UsageWorkerSupervisorState, + preserve_provider_response_facts, UsageBillingEventEnricher, UsageBodyCapturePolicy, + UsageRequestRecordLevel, UsageRuntimeAccess, UsageWorkerObservation, + UsageWorkerSupervisorState, }; use crate::worker::ManualProxyNodeCounter; use crate::{ @@ -4010,16 +4019,21 @@ mod tests { provider_request_body_ref: Some( "usage://request/req-basic-1/provider_request_body".to_string(), ), - response_body: Some(json!({"error":{"message":"bad gateway"}})), + response_body: Some(json!({ + "error":{"message":"bad gateway"}, + "service_tier": "Default" + })), response_body_ref: Some("usage://request/req-basic-1/response_body".to_string()), client_response_body: Some(json!({"detail":"bad gateway"})), client_response_body_ref: Some( "usage://request/req-basic-1/client_response_body".to_string(), ), + request_metadata: Some(json!({"provider_service_tier": "priority"})), ..UsageEventData::default() }, ); + preserve_provider_response_facts(&mut event); apply_usage_body_capture_policy_to_event( UsageBodyCapturePolicy { record_level: UsageRequestRecordLevel::Basic, @@ -4038,5 +4052,14 @@ mod tests { assert!(event.data.response_body_ref.is_none()); assert!(event.data.client_response_body.is_none()); assert!(event.data.client_response_body_ref.is_none()); + assert_eq!( + event + .data + .request_metadata + .as_ref() + .and_then(|metadata| metadata.get("provider_actual_service_tier")) + .and_then(serde_json::Value::as_str), + Some("default") + ); } } diff --git a/crates/aether-usage-runtime/src/write.rs b/crates/aether-usage-runtime/src/write.rs index 6b729ae18..aa6e572a5 100644 --- a/crates/aether-usage-runtime/src/write.rs +++ b/crates/aether-usage-runtime/src/write.rs @@ -13,7 +13,8 @@ use crate::body_capture::{ RuntimeBodyCaptureMetadataInput, }; use crate::request_metadata::{ - attach_provider_request_body_metadata, build_usage_request_metadata_seed, + attach_provider_actual_service_tier_metadata, attach_provider_request_body_metadata, + attach_provider_response_body_metadata, build_usage_request_metadata_seed, merge_usage_request_metadata, merge_usage_request_metadata_owned, sanitize_usage_request_metadata, sanitize_usage_request_metadata_ref, }; @@ -163,6 +164,7 @@ pub struct StreamTerminalUsagePayloadSeed { pub client_response: Option, pub client_response_body_state: Option, pub standardized_usage: Option, + pub provider_actual_service_tier: Option, pub observed_stream_finish: Option, pub terminal_error_message: Option, pub capture_metadata: Option, @@ -689,8 +691,13 @@ fn build_terminal_usage_event_from_seed_impl( } else { merge_usage_request_metadata(request_metadata, audit_payload) }; - let request_metadata = - attach_provider_request_body_metadata(request_metadata, provider_request.as_ref()); + let request_metadata = attach_provider_request_body_metadata( + request_metadata, + Some(provider_contract.as_str()), + target_model.as_deref().or(Some(model.as_str())), + Some(model.as_str()), + provider_request.as_ref(), + ); let mut data = UsageEventData { user_id, @@ -965,6 +972,10 @@ pub fn build_stream_terminal_usage_payload_seed( .terminal_summary .as_ref() .and_then(|summary| summary.standardized_usage.clone()), + provider_actual_service_tier: payload + .terminal_summary + .as_ref() + .and_then(|summary| summary.provider_actual_service_tier.clone()), observed_stream_finish, terminal_error_message, capture_metadata: build_payload_body_capture_metadata( @@ -1004,6 +1015,10 @@ pub fn build_sync_terminal_usage_seed( status_code, provider_response_full.as_ref(), ); + let request_metadata = attach_provider_response_body_metadata( + context_seed.request_metadata, + provider_response_full.as_ref(), + ); TerminalUsageSeed { terminal_state, @@ -1046,7 +1061,7 @@ pub fn build_sync_terminal_usage_seed( provider_response: provider_response_full, client_response_headers, client_response, - request_metadata: context_seed.request_metadata, + request_metadata, audit_payload: capture_metadata, standardized_usage, } @@ -1088,6 +1103,7 @@ pub fn build_stream_terminal_usage_seed( mut client_response, mut client_response_body_state, standardized_usage, + provider_actual_service_tier, observed_stream_finish, terminal_error_message, capture_metadata, @@ -1177,6 +1193,12 @@ pub fn build_stream_terminal_usage_seed( missing_observed_finish, terminal_error_message.is_some(), ); + let request_metadata = attach_provider_actual_service_tier_metadata( + context_seed.request_metadata, + provider_actual_service_tier.as_deref(), + ); + let request_metadata = + attach_provider_response_body_metadata(request_metadata, provider_response_full.as_ref()); TerminalUsageSeed { terminal_state, @@ -1219,7 +1241,7 @@ pub fn build_stream_terminal_usage_seed( provider_response: provider_response_full, client_response_headers, client_response, - request_metadata: context_seed.request_metadata, + request_metadata, audit_payload: capture_metadata, standardized_usage, } @@ -4381,6 +4403,7 @@ mod tests { finish_reason: None, response_id: Some("resp_cancel_summary_1".to_string()), model: Some("gpt-5.4".to_string()), + provider_actual_service_tier: None, observed_finish: true, unknown_event_count: 0, parser_error: None, @@ -4587,6 +4610,7 @@ mod tests { finish_reason: Some("stop".to_string()), response_id: Some("resp_summary_1".to_string()), model: Some("gpt-5.4".to_string()), + provider_actual_service_tier: Some("Default".to_string()), observed_finish: true, unknown_event_count: 0, parser_error: None, @@ -4605,6 +4629,15 @@ mod tests { assert_eq!(event.data.cache_read_input_tokens, Some(3)); assert!(event.data.response_body.is_none()); assert!(event.data.client_response_body.is_none()); + assert_eq!( + event + .data + .request_metadata + .as_ref() + .and_then(|metadata| metadata.get("provider_actual_service_tier")) + .and_then(Value::as_str), + Some("default") + ); } #[test] @@ -5041,6 +5074,7 @@ mod tests { finish_reason: Some("stop".to_string()), response_id: Some("resp_image_estimate_1".to_string()), model: Some("gpt-image-2".to_string()), + provider_actual_service_tier: None, observed_finish: true, unknown_event_count: 0, parser_error: None, @@ -5295,6 +5329,7 @@ mod tests { finish_reason: None, response_id: Some("resp_123".to_string()), model: Some("gpt-5.5".to_string()), + provider_actual_service_tier: None, observed_finish: true, unknown_event_count: 0, parser_error: None, diff --git a/docs/api/format-enum-mapping.md b/docs/api/format-enum-mapping.md index 1d493efab..519dac292 100644 --- a/docs/api/format-enum-mapping.md +++ b/docs/api/format-enum-mapping.md @@ -1,7 +1,5 @@ # Format Enum Mapping -Last audited: 2026-06-03 - Status values used below: - `native`: same semantic value exists in the target format. @@ -16,19 +14,19 @@ Provider-specific types: `OpenAiChatReasoningEffort` for Chat `reasoning_effort` | Source field | Source value | Target field | Target value | Status | | --- | --- | --- | --- | --- | | Chat `reasoning_effort` | `none` | Responses `reasoning.effort` | `none` | native | -| Chat `reasoning_effort` | `minimal` | Responses `reasoning.effort` | `minimal` | native | +| Chat `reasoning_effort` | `minimal` | Responses `reasoning.effort` | `minimal` | native when the target model supports it; blocked for GPT-5.6 | | Chat `reasoning_effort` | `low` | Responses `reasoning.effort` | `low` | native | | Chat `reasoning_effort` | `medium` | Responses `reasoning.effort` | `medium` | native | | Chat `reasoning_effort` | `high` | Responses `reasoning.effort` | `high` | native | | Chat `reasoning_effort` | `xhigh` | Responses `reasoning.effort` | `xhigh` | native | -| Chat `reasoning_effort` | `max` | Responses `reasoning.effort` | none | blocked, invalid OpenAI enum | +| Chat `reasoning_effort` | `max` | Responses `reasoning.effort` | `max` | native for GPT-5.6; blocked for models that do not publish `max` | | Responses `reasoning.effort` | `none` | Chat `reasoning_effort` | `none` | native | -| Responses `reasoning.effort` | `minimal` | Chat `reasoning_effort` | `minimal` | native | +| Responses `reasoning.effort` | `minimal` | Chat `reasoning_effort` | `minimal` | native when the target model supports it; blocked for GPT-5.6 | | Responses `reasoning.effort` | `low` | Chat `reasoning_effort` | `low` | native | | Responses `reasoning.effort` | `medium` | Chat `reasoning_effort` | `medium` | native | | Responses `reasoning.effort` | `high` | Chat `reasoning_effort` | `high` | native | | Responses `reasoning.effort` | `xhigh` | Chat `reasoning_effort` | `xhigh` | native | -| Responses `reasoning.effort` | `max` | Chat `reasoning_effort` | none | blocked, invalid Responses enum | +| Responses `reasoning.effort` | `max` | Chat `reasoning_effort` | `max` | native for GPT-5.6; blocked for models that do not publish `max` | | Responses `reasoning.summary` | any | Chat | none | blocked | | Responses `reasoning.budget_tokens` | any | Chat | none | blocked | @@ -42,7 +40,9 @@ Internal model directive values: | `medium` | `medium` | `medium` | `medium` | `medium` | Budget maps to `2048`. | | `high` | `high` | `high` | `high` | `high` | Budget maps to `4096`. | | `xhigh` | `xhigh` | `xhigh` | `xhigh` | `high` | Budget maps to `8192`. | -| `max` | `xhigh` | `xhigh` | `max` | `high` | Internal directive only; not accepted as raw OpenAI input. | +| `max` | `max` for GPT-5.6 | `max` for GPT-5.6 | `max` | `high` | OpenAI emission is capability-gated by the resolved model. | + +GPT-5.6 (`gpt-5.6`, `gpt-5.6-sol`, `gpt-5.6-terra`, and `gpt-5.6-luna`) publishes `none`, `low`, `medium`, `high`, `xhigh`, and `max`, and does not support `minimal`. Additional non-empty effort values advertised by a model are preserved verbatim across OpenAI Chat and Responses conversion. `ultra` is a Codex client preset that resolves to `max` before transmission and is not an OpenAI wire effort. Known effort capabilities are validated against the resolved provider model, so aliases mapped to GPT-5.6 receive the GPT-5.6 contract while concrete model families keep their published constraints. ## Tool Choice diff --git a/frontend/src/api/__tests__/models-dev-pricing.spec.ts b/frontend/src/api/__tests__/models-dev-pricing.spec.ts new file mode 100644 index 000000000..0cd7e1588 --- /dev/null +++ b/frontend/src/api/__tests__/models-dev-pricing.spec.ts @@ -0,0 +1,176 @@ +import { describe, expect, it } from 'vitest' + +import { + buildModelsDevTieredPricing, + resolveModelsDevTieredPricing, +} from '@/api/models-dev-pricing' + +describe('buildModelsDevTieredPricing', () => { + it('maps context bands and cache prices without flattening them', () => { + expect(buildModelsDevTieredPricing({ + input: 5, + output: 30, + cache_read: 0.5, + cache_write: 6.25, + tiers: [{ + input: 10, + output: 45, + cache_read: 1, + cache_write: 12.5, + tier: { type: 'context', size: 272_000 }, + }], + })).toEqual({ + tiers: [ + { + up_to: 271_999, + input_price_per_1m: 5, + output_price_per_1m: 30, + cache_creation_price_per_1m: 6.25, + cache_read_price_per_1m: 0.5, + }, + { + up_to: null, + input_price_per_1m: 10, + output_price_per_1m: 45, + cache_creation_price_per_1m: 12.5, + cache_read_price_per_1m: 1, + }, + ], + }) + }) + + it('sorts multiple context boundaries into contiguous Aether bands', () => { + const cost = { + input: 1, + output: 2, + tiers: [ + { input: 5, output: 6, tier: { type: 'context' as const, size: 200_000 } }, + { input: 3, output: 4, tier: { type: 'context' as const, size: 100_000 } }, + ], + } + + expect(buildModelsDevTieredPricing(cost)?.tiers).toEqual([ + { up_to: 99_999, input_price_per_1m: 1, output_price_per_1m: 2 }, + { up_to: 199_999, input_price_per_1m: 3, output_price_per_1m: 4 }, + { up_to: null, input_price_per_1m: 5, output_price_per_1m: 6 }, + ]) + expect(cost.tiers.map(tier => tier.tier.size)).toEqual([200_000, 100_000]) + }) + + it('keeps flat token pricing as one unbounded band', () => { + expect(buildModelsDevTieredPricing({ input: 0, output: 0.1 })).toEqual({ + tiers: [{ up_to: null, input_price_per_1m: 0, output_price_per_1m: 0.1 }], + }) + }) + + it('omits an empty base band when context pricing starts at zero', () => { + expect(buildModelsDevTieredPricing({ + input: 1, + output: 2, + tiers: [ + { input: 3, output: 4, tier: { type: 'context', size: 0 } }, + { input: 5, output: 6, tier: { type: 'context', size: 100_000 } }, + ], + })?.tiers).toEqual([ + { up_to: 99_999, input_price_per_1m: 3, output_price_per_1m: 4 }, + { up_to: null, input_price_per_1m: 5, output_price_per_1m: 6 }, + ]) + }) + + it.each([ + { input: -1, output: 2 }, + { input: 1, output: Number.POSITIVE_INFINITY }, + { + input: 1, + output: 2, + tiers: [{ input: 3, output: 4, tier: { type: 'context', size: Number.MAX_SAFE_INTEGER + 1 } }], + }, + { + input: 1, + output: 2, + tiers: [{ input: 3, output: 4, tier: { type: 'context', size: -1 } }], + }, + { + input: 1, + output: 2, + tiers: [ + { input: 3, output: 4, tier: { type: 'context', size: 100 } }, + { input: 5, output: 6, tier: { type: 'context', size: 100 } }, + ], + }, + ])('fails closed for malformed structured pricing', (cost) => { + expect(buildModelsDevTieredPricing(cost)).toBeNull() + }) +}) + +describe('resolveModelsDevTieredPricing', () => { + it.each([ + { + modelId: 'gpt-5.6-sol', + standard: [5, 30, 6.25, 0.5], + longContext: [10, 45, 12.5, 1], + }, + { + modelId: 'gpt-5.6-terra', + standard: [2.5, 15, 3.125, 0.25], + longContext: [5, 22.5, 6.25, 0.5], + }, + { + modelId: 'gpt-5.6-luna', + standard: [1, 6, 1.25, 0.1], + longContext: [2, 9, 2.5, 0.2], + }, + ])('uses the complete OpenAI catalog for $modelId', ({ modelId, standard, longContext }) => { + const tier = ( + upTo: number | null, + prices: number[], + multiplier: number, + ) => ({ + up_to: upTo, + input_price_per_1m: prices[0] * multiplier, + output_price_per_1m: prices[1] * multiplier, + cache_creation_price_per_1m: prices[2] * multiplier, + cache_read_price_per_1m: prices[3] * multiplier, + }) + + expect(resolveModelsDevTieredPricing('openai', modelId, { input: 999, output: 999 })) + .toEqual({ + tiers: [ + tier(272_000, standard, 1), + tier(null, longContext, 1), + ], + processing_tiers: { + flex: { + tiers: [ + tier(272_000, standard, 0.5), + tier(null, longContext, 0.5), + ], + }, + priority: { + tiers: [tier(272_000, standard, 2)], + }, + }, + }) + }) + + it('keeps the models.dev lower-bound conversion for models outside the catalog', () => { + expect(resolveModelsDevTieredPricing('openai', 'other-model', { + input: 1, + output: 2, + tiers: [{ input: 3, output: 4, tier: { type: 'context', size: 272_000 } }], + })?.tiers.map(tier => tier.up_to)).toEqual([271_999, null]) + }) + + it.each([ + ['other-provider', 'gpt-5.6-sol'], + ['openai', 'GPT-5.6-SOL'], + ['openai', 'gpt-5.6-sol-latest'], + ['openai', '__proto__'], + ['openai', 'constructor'], + ])('matches provider and model identities exactly for %s/%s', (providerId, modelId) => { + expect(resolveModelsDevTieredPricing(providerId, modelId, { input: 1, output: 2 })) + .toEqual({ + tiers: [{ up_to: null, input_price_per_1m: 1, output_price_per_1m: 2 }], + }) + }) +}) diff --git a/frontend/src/api/authoritative-model-pricing.ts b/frontend/src/api/authoritative-model-pricing.ts new file mode 100644 index 000000000..448c82c76 --- /dev/null +++ b/frontend/src/api/authoritative-model-pricing.ts @@ -0,0 +1,71 @@ +import type { PricingTier, TieredPricingConfig } from './endpoints/types' + +interface TokenPrices { + input: number + output: number + cacheCreation: number + cacheRead: number +} + +interface ContextPricing { + standard: TokenPrices + longContext: TokenPrices +} + +const OPENAI_GPT_56_PRICING = new Map([ + ['gpt-5.6-sol', { + standard: { input: 5, output: 30, cacheCreation: 6.25, cacheRead: 0.5 }, + longContext: { input: 10, output: 45, cacheCreation: 12.5, cacheRead: 1 }, + }], + ['gpt-5.6-terra', { + standard: { input: 2.5, output: 15, cacheCreation: 3.125, cacheRead: 0.25 }, + longContext: { input: 5, output: 22.5, cacheCreation: 6.25, cacheRead: 0.5 }, + }], + ['gpt-5.6-luna', { + standard: { input: 1, output: 6, cacheCreation: 1.25, cacheRead: 0.1 }, + longContext: { input: 2, output: 9, cacheCreation: 2.5, cacheRead: 0.2 }, + }], +]) + +const STANDARD_CONTEXT_LIMIT = 272_000 + +function pricingTier( + upTo: number | null, + prices: TokenPrices, + multiplier = 1, +): PricingTier { + return { + up_to: upTo, + input_price_per_1m: prices.input * multiplier, + output_price_per_1m: prices.output * multiplier, + cache_creation_price_per_1m: prices.cacheCreation * multiplier, + cache_read_price_per_1m: prices.cacheRead * multiplier, + } +} + +function contextPricingTiers(pricing: ContextPricing, multiplier: number): PricingTier[] { + return [ + pricingTier(STANDARD_CONTEXT_LIMIT, pricing.standard, multiplier), + pricingTier(null, pricing.longContext, multiplier), + ] +} + +export function getAuthoritativeModelPricing( + providerId: string, + modelId: string, +): TieredPricingConfig | null { + if (providerId !== 'openai') return null + + const pricing = OPENAI_GPT_56_PRICING.get(modelId) + if (!pricing) return null + + return { + tiers: contextPricingTiers(pricing, 1), + processing_tiers: { + flex: { tiers: contextPricingTiers(pricing, 0.5) }, + priority: { + tiers: [pricingTier(STANDARD_CONTEXT_LIMIT, pricing.standard, 2)], + }, + }, + } +} diff --git a/frontend/src/api/dashboard.ts b/frontend/src/api/dashboard.ts index d83108aeb..dd4b063e8 100644 --- a/frontend/src/api/dashboard.ts +++ b/frontend/src/api/dashboard.ts @@ -154,6 +154,18 @@ export interface RequestSchedulingFailure { no_upstream_attempt?: boolean | null } +export interface RequestSettlementPricingSnapshot { + requested_processing_tier?: string | null + actual_processing_tier?: string | null + billing_processing_tier?: string | null + [key: string]: unknown +} + +export interface RequestSettlementSnapshot { + pricing_snapshot?: RequestSettlementPricingSnapshot | null + [key: string]: unknown +} + export interface RequestDetail { id: string // UUID request_id: string @@ -175,6 +187,7 @@ export interface RequestDetail { target_model?: string | null // 映射后的目标模型名 reasoning_effort?: string | null service_tier?: string | null + actual_service_tier?: string | null tokens: { input: number output: number @@ -262,6 +275,7 @@ export interface RequestDetail { cache_creation_price_per_1m?: number cache_read_price_per_1m?: number price_per_request?: number + settlement_snapshot?: RequestSettlementSnapshot | null } | null // 阶梯计费信息 tiered_pricing?: { diff --git a/frontend/src/api/endpoints/types/__tests__/api-format.spec.ts b/frontend/src/api/endpoints/types/__tests__/api-format.spec.ts index 44122b355..70b002c6c 100644 --- a/frontend/src/api/endpoints/types/__tests__/api-format.spec.ts +++ b/frontend/src/api/endpoints/types/__tests__/api-format.spec.ts @@ -2,6 +2,7 @@ import { describe, expect, it } from 'vitest' import { API_FORMATS, + apiFormatPermissionCovers, formatApiFormat, formatApiFormatShort, groupApiFormats, @@ -14,6 +15,8 @@ describe('api format display helpers', () => { expect(normalizeApiFormatAlias('CLAUDE_MESSAGES')).toBe(API_FORMATS.CLAUDE_MESSAGES) expect(normalizeApiFormatAlias('OPENAI_RESPONSES')).toBe(API_FORMATS.OPENAI_RESPONSES) expect(normalizeApiFormatAlias('OPENAI_RESPONSES_COMPACT')).toBe(API_FORMATS.OPENAI_RESPONSES_COMPACT) + expect(normalizeApiFormatAlias('OPENAI_SEARCH')).toBe(API_FORMATS.OPENAI_SEARCH) + expect(normalizeApiFormatAlias('SEARCH')).toBe(API_FORMATS.OPENAI_SEARCH) expect(normalizeApiFormatAlias('GEMINI_GENERATE_CONTENT')).toBe(API_FORMATS.GEMINI_GENERATE_CONTENT) expect(normalizeApiFormatAlias('OPENAI_EMBEDDING')).toBe(API_FORMATS.OPENAI_EMBEDDING) expect(normalizeApiFormatAlias('OPENAI_RERANK')).toBe(API_FORMATS.OPENAI_RERANK) @@ -33,6 +36,25 @@ describe('api format display helpers', () => { expect(formatApiFormatShort(API_FORMATS.JINA_RERANK)).toBe('JR') }) + it('formats OpenAI Search as a first-class api format', () => { + expect(formatApiFormat(API_FORMATS.OPENAI_SEARCH)).toBe('OpenAI Search') + expect(formatApiFormatShort(API_FORMATS.OPENAI_SEARCH)).toBe('OS') + expect(sortApiFormats([ + API_FORMATS.OPENAI_EMBEDDING, + API_FORMATS.OPENAI_SEARCH, + API_FORMATS.OPENAI_RESPONSES, + ])).toEqual([ + API_FORMATS.OPENAI_RESPONSES, + API_FORMATS.OPENAI_SEARCH, + API_FORMATS.OPENAI_EMBEDDING, + ]) + }) + + it('applies Responses to Search permissions in one direction', () => { + expect(apiFormatPermissionCovers('OPENAI_RESPONSES', 'openai:search')).toBe(true) + expect(apiFormatPermissionCovers('openai:search', 'openai:responses')).toBe(false) + }) + it('formats embedding api format ids distinctly from chat formats', () => { expect(formatApiFormat(API_FORMATS.GEMINI_INTERACTIONS)).toBe('Gemini Interactions') expect(formatApiFormatShort(API_FORMATS.GEMINI_INTERACTIONS)).toBe('GI') diff --git a/frontend/src/api/endpoints/types/api-format.ts b/frontend/src/api/endpoints/types/api-format.ts index bae97e933..636355b0c 100644 --- a/frontend/src/api/endpoints/types/api-format.ts +++ b/frontend/src/api/endpoints/types/api-format.ts @@ -6,6 +6,7 @@ export const API_FORMATS = { OPENAI: 'openai:chat', OPENAI_RESPONSES: 'openai:responses', OPENAI_RESPONSES_COMPACT: 'openai:responses:compact', + OPENAI_SEARCH: 'openai:search', OPENAI_IMAGE: 'openai:image', OPENAI_VIDEO: 'openai:video', OPENAI_EMBEDDING: 'openai:embedding', @@ -30,6 +31,7 @@ export const API_FORMAT_LABELS: Record = { [API_FORMATS.OPENAI]: 'OpenAI Chat', [API_FORMATS.OPENAI_RESPONSES]: 'OpenAI Responses', [API_FORMATS.OPENAI_RESPONSES_COMPACT]: 'OpenAI Responses Compact', + [API_FORMATS.OPENAI_SEARCH]: 'OpenAI Search', [API_FORMATS.OPENAI_IMAGE]: 'OpenAI Image', [API_FORMATS.OPENAI_VIDEO]: 'OpenAI Video', [API_FORMATS.OPENAI_EMBEDDING]: 'OpenAI Embedding', @@ -48,6 +50,7 @@ export const API_FORMAT_LABELS: Record = { OPENAI: 'OpenAI Chat', OPENAI_RESPONSES: 'OpenAI Responses', OPENAI_RESPONSES_COMPACT: 'OpenAI Responses Compact', + OPENAI_SEARCH: 'OpenAI Search', OPENAI_IMAGE: 'OpenAI Image', OPENAI_VIDEO: 'OpenAI Video', OPENAI_EMBEDDING: 'OpenAI Embedding', @@ -69,6 +72,7 @@ export const API_FORMAT_SHORT: Record = { [API_FORMATS.OPENAI]: 'O', [API_FORMATS.OPENAI_RESPONSES]: 'OR', [API_FORMATS.OPENAI_RESPONSES_COMPACT]: 'ORC', + [API_FORMATS.OPENAI_SEARCH]: 'OS', [API_FORMATS.OPENAI_IMAGE]: 'OI', [API_FORMATS.OPENAI_VIDEO]: 'OV', [API_FORMATS.OPENAI_EMBEDDING]: 'OE', @@ -86,6 +90,7 @@ export const API_FORMAT_SHORT: Record = { OPENAI: 'O', OPENAI_RESPONSES: 'OR', OPENAI_RESPONSES_COMPACT: 'ORC', + OPENAI_SEARCH: 'OS', OPENAI_IMAGE: 'OI', OPENAI_VIDEO: 'OV', OPENAI_EMBEDDING: 'OE', @@ -109,6 +114,7 @@ export const API_FORMAT_ORDER: string[] = [ API_FORMATS.OPENAI, API_FORMATS.OPENAI_RESPONSES, API_FORMATS.OPENAI_RESPONSES_COMPACT, + API_FORMATS.OPENAI_SEARCH, API_FORMATS.OPENAI_EMBEDDING, API_FORMATS.OPENAI_RERANK, API_FORMATS.OPENAI_IMAGE, @@ -140,6 +146,7 @@ export const API_FORMAT_KIND_LABELS: Record = { chat: 'Chat', responses: 'Responses', 'responses:compact': 'Responses Compact', + search: 'Search', messages: 'Messages', generate_content: 'Generate Content', interactions: 'Interactions', @@ -174,6 +181,9 @@ export function normalizeApiFormatAlias(format: string | null | undefined): stri return API_FORMATS.OPENAI_RESPONSES case 'OPENAI_RESPONSES_COMPACT': return API_FORMATS.OPENAI_RESPONSES_COMPACT + case 'OPENAI_SEARCH': + case 'SEARCH': + return API_FORMATS.OPENAI_SEARCH case 'OPENAI_IMAGE': return API_FORMATS.OPENAI_IMAGE case 'OPENAI_VIDEO': @@ -216,6 +226,18 @@ export function normalizeApiFormatAlias(format: string | null | undefined): stri } } +export function apiFormatPermissionCovers( + allowedFormat: string | null | undefined, + requestedFormat: string | null | undefined, +): boolean { + const allowed = normalizeApiFormatAlias(allowedFormat) + const requested = normalizeApiFormatAlias(requestedFormat) + return Boolean(allowed) + && Boolean(requested) + && (allowed === requested + || (allowed === API_FORMATS.OPENAI_RESPONSES && requested === API_FORMATS.OPENAI_SEARCH)) +} + // 工具函数：按 family 分组并排序 API 格式数组 export interface ApiFormatGroup { family: string @@ -272,14 +294,16 @@ export function formatApiFormatShort(format: string | null | undefined): string // 工具函数：按标准顺序排序 API 格式数组 export function sortApiFormats(formats: string[]): string[] { - return [...formats].sort((a, b) => { - const aIdx = API_FORMAT_ORDER.indexOf(normalizeApiFormatAlias(a)) - const bIdx = API_FORMAT_ORDER.indexOf(normalizeApiFormatAlias(b)) - if (aIdx === -1 && bIdx === -1) return 0 - if (aIdx === -1) return 1 - if (bIdx === -1) return -1 - return aIdx - bIdx - }) + return [...formats].sort(compareApiFormats) +} + +export function compareApiFormats(a: string, b: string): number { + const aIdx = API_FORMAT_ORDER.indexOf(normalizeApiFormatAlias(a)) + const bIdx = API_FORMAT_ORDER.indexOf(normalizeApiFormatAlias(b)) + if (aIdx === -1 && bIdx === -1) return 0 + if (aIdx === -1) return 1 + if (bIdx === -1) return -1 + return aIdx - bIdx } // openai family 格式只支持 bearer（Authorization header），不允许覆盖认证方式 diff --git a/frontend/src/api/endpoints/types/model.ts b/frontend/src/api/endpoints/types/model.ts index b97f1a0d0..ef9fe8f8b 100644 --- a/frontend/src/api/endpoints/types/model.ts +++ b/frontend/src/api/endpoints/types/model.ts @@ -6,6 +6,7 @@ import type { ProviderModelMapping } from './provider' export interface CacheTTLPricing { ttl_minutes: number cache_creation_price_per_1m: number + [key: string]: unknown } /** 单个价格阶梯配置 */ @@ -16,22 +17,39 @@ export interface PricingTier { cache_creation_price_per_1m?: number cache_read_price_per_1m?: number cache_ttl_pricing?: CacheTTLPricing[] + [key: string]: unknown } export type ImageOutputQuality = 'low' | 'medium' | 'high' +export interface ImageOutputQualityPricing extends Partial> { + [quality: string]: unknown +} + export interface ImageOutputPriceRange { up_to_pixels: number | null - prices: Partial> + prices: ImageOutputQualityPricing label?: string | null + [key: string]: unknown +} + +/** 按处理层级覆盖的费率配置。允许图像或未来计费字段独立扩展。 */ +export interface ProcessingTierPricingConfig { + tiers?: PricingTier[] + image_output_prices?: Record | null + image_output_price_default?: number | null + image_output_price_ranges?: ImageOutputPriceRange[] | null + [key: string]: unknown } /** 阶梯计费配置 */ export interface TieredPricingConfig { tiers: PricingTier[] - image_output_prices?: Record> | null + image_output_prices?: Record | null image_output_price_default?: number | null image_output_price_ranges?: ImageOutputPriceRange[] | null + processing_tiers?: Record | null + [key: string]: unknown } export interface Model { @@ -270,6 +288,8 @@ export interface UpstreamModel { id: string owned_by?: string display_name?: string + visibility?: string + supported_in_api?: boolean api_formats: string[] // 该模型支持的所有 API 格式（后端保证返回数组） model_test_capabilities?: ModelTestCapabilities | null } diff --git a/frontend/src/api/me.ts b/frontend/src/api/me.ts index 531de4e5e..39efc1a28 100644 --- a/frontend/src/api/me.ts +++ b/frontend/src/api/me.ts @@ -56,6 +56,7 @@ export interface UsageRecordDetail { model: string reasoning_effort?: string | null service_tier?: string | null + actual_service_tier?: string | null input_tokens: number effective_input_tokens?: number output_tokens: number @@ -369,6 +370,7 @@ export const meApi = { target_model?: string | null reasoning_effort?: string | null service_tier?: string | null + actual_service_tier?: string | null }> }> { const params = ids ? { ids } : {} diff --git a/frontend/src/api/models-dev-pricing.ts b/frontend/src/api/models-dev-pricing.ts new file mode 100644 index 000000000..d875f6b39 --- /dev/null +++ b/frontend/src/api/models-dev-pricing.ts @@ -0,0 +1,102 @@ +import type { PricingTier, TieredPricingConfig } from './endpoints/types' +import { getAuthoritativeModelPricing } from './authoritative-model-pricing' + +export interface ModelsDevTokenCost { + input: number + output: number + reasoning?: number + cache_read?: number + cache_write?: number + input_audio?: number + output_audio?: number +} + +export interface ModelsDevCostTier extends ModelsDevTokenCost { + tier: { + type: 'context' + size: number + } +} + +export interface ModelsDevCost extends ModelsDevTokenCost { + tiers?: ModelsDevCostTier[] +} + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null && !Array.isArray(value) +} + +function isPrice(value: unknown): value is number { + return typeof value === 'number' && Number.isFinite(value) && value >= 0 +} + +function parseTokenPrices(value: unknown): Omit | null { + if (!isRecord(value) || !isPrice(value.input) || !isPrice(value.output)) return null + if (value.cache_write !== undefined && !isPrice(value.cache_write)) return null + if (value.cache_read !== undefined && !isPrice(value.cache_read)) return null + + return { + input_price_per_1m: value.input, + output_price_per_1m: value.output, + ...(value.cache_write === undefined + ? {} + : { cache_creation_price_per_1m: value.cache_write }), + ...(value.cache_read === undefined + ? {} + : { cache_read_price_per_1m: value.cache_read }), + } +} + +function parseContextTier(value: unknown): { size: number; prices: Omit } | null { + if (!isRecord(value) || !isRecord(value.tier)) return null + if ( + value.tier.type !== 'context' + || typeof value.tier.size !== 'number' + || !Number.isSafeInteger(value.tier.size) + || value.tier.size < 0 + ) { + return null + } + const prices = parseTokenPrices(value) + return prices ? { size: value.tier.size, prices } : null +} + +export function buildModelsDevTieredPricing(cost: unknown): TieredPricingConfig | null { + const basePrices = parseTokenPrices(cost) + if (!basePrices || !isRecord(cost)) return null + + const rawTiers = cost.tiers + if (rawTiers !== undefined && !Array.isArray(rawTiers)) return null + const contextTiers = (rawTiers ?? []).map(parseContextTier) + if (contextTiers.some(tier => tier === null)) return null + + const sortedTiers = contextTiers + .filter((tier): tier is NonNullable => tier !== null) + .sort((a, b) => a.size - b.size) + if (sortedTiers.some((tier, index) => index > 0 && tier.size === sortedTiers[index - 1].size)) { + return null + } + + const tiers: PricingTier[] = [] + if (sortedTiers[0]?.size !== 0) { + tiers.push({ + ...basePrices, + up_to: sortedTiers[0] ? sortedTiers[0].size - 1 : null, + }) + } + tiers.push(...sortedTiers.map((tier, index) => ({ + ...tier.prices, + up_to: sortedTiers[index + 1] ? sortedTiers[index + 1].size - 1 : null, + }))) + + return { tiers } +} + +export function resolveModelsDevTieredPricing( + providerId: string, + modelId: string, + cost: unknown, +): TieredPricingConfig | null { + return getAuthoritativeModelPricing(providerId, modelId) + ?? buildModelsDevTieredPricing(cost) +} diff --git a/frontend/src/api/models-dev.ts b/frontend/src/api/models-dev.ts index f5ef5c949..f73d22723 100644 --- a/frontend/src/api/models-dev.ts +++ b/frontend/src/api/models-dev.ts @@ -4,19 +4,19 @@ */ import api from './client' +import { + resolveModelsDevTieredPricing, + type ModelsDevCost, +} from './models-dev-pricing' +import type { TieredPricingConfig } from './endpoints/types' + +export type { ModelsDevCost, ModelsDevCostTier, ModelsDevTokenCost } from './models-dev-pricing' // 缓存配置 const CACHE_KEY = 'models_dev_cache' const CACHE_DURATION = 15 * 60 * 1000 // 15 分钟 // Models.dev API 数据结构 -export interface ModelsDevCost { - input?: number - output?: number - reasoning?: number - cache_read?: number -} - export interface ModelsDevLimit { context?: number output?: number @@ -64,6 +64,7 @@ export interface ModelsDevModelItem { family?: string inputPrice?: number outputPrice?: number + tieredPricing?: TieredPricingConfig contextLimit?: number outputLimit?: number supportsVision?: boolean @@ -165,14 +166,17 @@ export async function getModelsDevList(officialOnly: boolean = true): Promise }> { diff --git a/frontend/src/features/models/components/GlobalModelFormDialog.vue b/frontend/src/features/models/components/GlobalModelFormDialog.vue index 5abb078d7..08d06be60 100644 --- a/frontend/src/features/models/components/GlobalModelFormDialog.vue +++ b/frontend/src/features/models/components/GlobalModelFormDialog.vue @@ -602,6 +602,7 @@ import { buildGlobalModelCreatePayload, buildGlobalModelUpdatePayload, } from './global-model-form-helpers' +import { tieredPricingHasImageOutputPricing } from '../utils/tiered-pricing' const props = defineProps<{ open: boolean @@ -1104,17 +1105,9 @@ function selectModel(model: ModelsDevModelItem) { } loadVideoPricingFromConfig() - if (model.inputPrice !== undefined || model.outputPrice !== undefined) { - tieredPricing.value = { - tiers: [{ - up_to: null, - input_price_per_1m: model.inputPrice || 0, - output_price_per_1m: model.outputPrice || 0, - }] - } - } else { - tieredPricing.value = null - } + tieredPricing.value = model.tieredPricing + ? structuredClone(model.tieredPricing) + : null presetPanelCollapsed.value = true scrollToBasicInformation() @@ -1204,6 +1197,12 @@ async function handleSubmit() { return } + const pricingValidationError = tieredPricingEditorRef.value?.getValidationError() + if (pricingValidationError) { + showError(pricingValidationError, '价格配置错误') + return + } + const finalTieredPricing = tieredPricingEditorRef.value?.getFinalPricing() ?? tieredPricing.value if (!finalTieredPricing?.tiers?.length) { @@ -1244,28 +1243,4 @@ async function handleSubmit() { } } -function tieredPricingHasImageOutputPricing(pricing: TieredPricingConfig | null | undefined): boolean { - if (!pricing) return false - if (toFinitePrice(pricing.image_output_price_default) !== null) return true - if (Object.values(pricing.image_output_prices || {}).some((prices) => { - if (!prices || typeof prices !== 'object') return false - return Object.values(prices).some((price) => toFinitePrice(price) !== null) - })) return true - return (pricing.image_output_price_ranges || []).some((range) => { - if (!range || typeof range !== 'object') return false - const prices = range.prices && typeof range.prices === 'object' - ? range.prices - : range as Record - return Object.values(prices).some((price) => toFinitePrice(price) !== null) - }) -} - -function toFinitePrice(value: unknown): number | null { - if (typeof value === 'number' && Number.isFinite(value)) return value - if (typeof value === 'string' && value.trim()) { - const parsed = Number(value) - return Number.isFinite(parsed) ? parsed : null - } - return null -} diff --git a/frontend/src/features/models/components/ModelDetailDrawer.vue b/frontend/src/features/models/components/ModelDetailDrawer.vue index 8939837e9..7f8ac50bc 100644 --- a/frontend/src/features/models/components/ModelDetailDrawer.vue +++ b/frontend/src/features/models/components/ModelDetailDrawer.vue @@ -137,6 +137,7 @@

+ @@ -145,6 +146,8 @@ 默认定价 + +
+
+
+
+
+ 处理层级定价 +
+

+ {{ activeEntry.label }} +

+
+
+ +
+
+ +
+ + + + + Token 区间 + + + 输入 ($/M) + + + 输出 ($/M) + + + 缓存创建 + + + 缓存读取 + + + 1h 创建 + + + + + + + {{ formatTokenRange(activeTokenTiers, index) }} + + + {{ formatPrice(tier.input_price_per_1m) }} + + + {{ formatPrice(tier.output_price_per_1m) }} + + + {{ formatPrice(tier.cache_creation_price_per_1m) }} + + + {{ formatPrice(tier.cache_read_price_per_1m) }} + + + {{ formatPrice(cacheCreationPriceForTtl(tier, 60)) }} + + + +
+
+ +
+
+ 图片输出 + 默认 {{ formatPrice(activeImageDefaultPrice) }}/张 +
+ +
+ + + + + 分辨率 + + + {{ quality }} + + + + + + + {{ formatImageSize(row.size) }} + + + {{ formatPrice(row.prices[quality]) }} + + + +
+
+ +
+ + + + + 像素区间 + + + {{ quality }} + + + + + + + {{ row.label || formatPixelRange(activeImageRangeRows, index) }} + + + {{ formatPrice(row.prices[quality]) }} + + + +
+
+
+
+ + + diff --git a/frontend/src/features/models/components/TieredPricingEditor.vue b/frontend/src/features/models/components/TieredPricingEditor.vue index 40b3d5ca5..23caee0ac 100644 --- a/frontend/src/features/models/components/TieredPricingEditor.vue +++ b/frontend/src/features/models/components/TieredPricingEditor.vue @@ -1,24 +1,100 @@