Merge PR #669: align GPT-5.6 and Codex request protocols

This commit is contained in:
elky
2026-07-12 21:50:20 +08:00
313 changed files with 29823 additions and 5178 deletions
+7
View File
@@ -383,6 +383,13 @@ jobs:
AETHER_TEST_POSTGRES_URL: postgres://aether:aether@127.0.0.1:5432/aether_test 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 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 - name: Run Postgres core export smoke test
env: env:
RUSTC_WRAPPER: sccache RUSTC_WRAPPER: sccache
@@ -68,6 +68,7 @@ pub(crate) use aether_ai_formats::api::{
OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND, OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
}; };
pub(crate) use aether_ai_formats::protocol::stream::CanonicalUsage as StreamingCanonicalUsage; 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( pub(crate) fn parse_direct_request_body(
parts: &http::request::Parts, parts: &http::request::Parts,
@@ -31,6 +31,7 @@ fn test_decision() -> GatewayControlDecision {
auth_context: None, auth_context: None,
admin_principal: None, admin_principal: None,
local_auth_rejection: 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.report_kind, "openai_responses_compact_sync_success");
assert_eq!( assert_eq!(
report.client_body_json.expect("client body should exist")["object"], 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 .background_report
.expect("compact tool-call should downgrade to success report"); .expect("compact tool-call should downgrade to success report");
let client_body = report.client_body_json.expect("client body should exist"); 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"); 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, auth_context: None,
admin_principal: None, admin_principal: None,
local_auth_rejection: None, local_auth_rejection: None,
model_directive_policy: Default::default(),
}, },
&payload, &payload,
) )
@@ -1908,6 +1910,7 @@ fn local_finalize_handles_gemini_cli_cross_format_sync_response_from_claude_cli(
auth_context: None, auth_context: None,
admin_principal: None, admin_principal: None,
local_auth_rejection: None, local_auth_rejection: None,
model_directive_policy: Default::default(),
}, },
&payload, &payload,
) )
+14 -11
View File
@@ -13,7 +13,9 @@ use crate::{usage::GatewaySyncReportRequest, AppState, GatewayError};
pub(crate) use self::adaptation::{ pub(crate) use self::adaptation::{
maybe_build_provider_private_stream_normalizer, ProviderPrivateStreamNormalizer, 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::common::LocalCoreSyncFinalizeOutcome;
pub(crate) use self::finalize::internal::{ pub(crate) use self::finalize::internal::{
maybe_bridge_standard_sync_json_to_stream, maybe_build_stream_response_rewriter, 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_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_family_sync_plan_and_reports, build_standard_stream_plan_from_decision,
build_standard_sync_plan_from_decision, candidate_auth_channel_skip_reason, build_standard_sync_plan_from_decision, candidate_auth_channel_skip_reason,
extract_pool_sticky_session_token, maybe_build_stream_decision_payload, codex_model_capabilities_for_transport, extract_pool_sticky_session_token,
maybe_build_stream_plan_payload, maybe_build_sync_decision_payload, maybe_build_stream_decision_payload, maybe_build_stream_plan_payload,
maybe_build_sync_plan_payload, planner_is_matching_stream_request, provider_key_pool_score_id, maybe_build_sync_decision_payload, maybe_build_sync_plan_payload,
provider_key_pool_score_scope, read_candidate_transport_snapshot, planner_is_matching_stream_request, provider_key_pool_score_id, provider_key_pool_score_scope,
record_local_runtime_candidate_skip_reason, read_candidate_transport_snapshot, record_local_runtime_candidate_skip_reason,
set_local_openai_chat_execution_exhausted_diagnostic, resolve_upstream_is_stream_for_provider, set_local_openai_chat_execution_exhausted_diagnostic,
set_local_openai_image_execution_exhausted_diagnostic, CandidateFailureDiagnostic, set_local_openai_image_execution_exhausted_diagnostic, validate_final_openai_provider_request,
CandidateFailureDiagnosticKind, EligibleLocalExecutionCandidate, GatewayAuthApiKeySnapshot, CandidateFailureDiagnostic, CandidateFailureDiagnosticKind, EligibleLocalExecutionCandidate,
GatewayProviderTransportSnapshot, LocalExecutionAttemptSource, LocalExecutionCandidateKind, GatewayAuthApiKeySnapshot, GatewayProviderTransportSnapshot, LocalExecutionAttemptSource,
LocalResolvedOAuthRequestAuth, PlannerAppState, SkippedLocalExecutionCandidate, LocalExecutionCandidateKind, LocalResolvedOAuthRequestAuth, PlannerAppState,
SkippedLocalExecutionCandidate,
}; };
pub(crate) use self::pure::*; pub(crate) use self::pure::*;
pub(crate) use self::transport::{ pub(crate) use self::transport::{
@@ -702,6 +702,7 @@ pub(crate) async fn build_lazy_requested_model_execution_candidate_attempt_sourc
G, G,
>( >(
state: PlannerAppState<'a>, state: PlannerAppState<'a>,
model_directive_policy: &crate::system_features::ModelDirectivePolicySnapshot,
trace_id: &str, trace_id: &str,
client_api_format: &str, client_api_format: &str,
requested_model: &str, requested_model: &str,
@@ -730,6 +731,7 @@ where
let record_runtime_miss_diagnostic = persistence_policy.skipped.record_runtime_miss_diagnostic; let record_runtime_miss_diagnostic = persistence_policy.skipped.record_runtime_miss_diagnostic;
let page_cursor = LocalCandidatePreselectionPageCursor::new( let page_cursor = LocalCandidatePreselectionPageCursor::new(
state, state,
model_directive_policy,
client_api_format, client_api_format,
requested_model, requested_model,
require_streaming, require_streaming,
@@ -1151,6 +1153,9 @@ async fn resolve_priority_candidate_page_with_cache(
.page_cursor .page_cursor
.resolved_page_cache_use_api_format_alias_match(), .resolved_page_cache_use_api_format_alias_match(),
cursor.client_session_affinity.as_ref(), cursor.client_session_affinity.as_ref(),
cursor
.page_cursor
.resolved_page_cache_model_directive_policy_hash(),
cursor.resolution_mode, cursor.resolution_mode,
); );
let page_candidates_for_fallback = page_candidates.clone(); let page_candidates_for_fallback = page_candidates.clone();
@@ -1949,6 +1954,7 @@ mod tests {
global_model_id: "global-model-1".to_string(), global_model_id: "global-model-1".to_string(),
global_model_name: "gpt-5".to_string(), global_model_name: "gpt-5".to_string(),
selected_provider_model_name: "gpt-5".to_string(), selected_provider_model_name: "gpt-5".to_string(),
supports_streaming: true,
mapping_matched_model: None, mapping_matched_model: None,
} }
} }
@@ -2177,8 +2183,11 @@ mod tests {
async fn resolved_candidate_page_cache_requires_fixed_order_or_explicit_affinity() { async fn resolved_candidate_page_cache_requires_fixed_order_or_explicit_affinity() {
let app = AppState::new().expect("state should build"); let app = AppState::new().expect("state should build");
let auth_snapshot = sample_auth_snapshot(); let auth_snapshot = sample_auth_snapshot();
let model_directive_policy =
crate::system_features::ModelDirectivePolicySnapshot::default();
let mut page_cursor = LocalCandidatePreselectionPageCursor::new( let mut page_cursor = LocalCandidatePreselectionPageCursor::new(
PlannerAppState::new(&app), PlannerAppState::new(&app),
&model_directive_policy,
"openai:chat", "openai:chat",
"gpt-5", "gpt-5",
true, true,
@@ -2233,6 +2242,7 @@ mod tests {
let mut page_cursor = LocalCandidatePreselectionPageCursor::new( let mut page_cursor = LocalCandidatePreselectionPageCursor::new(
PlannerAppState::new(&app), PlannerAppState::new(&app),
&model_directive_policy,
"openai:chat", "openai:chat",
"gpt-5", "gpt-5",
true, true,
@@ -2269,6 +2279,7 @@ mod tests {
); );
let mut page_cursor = LocalCandidatePreselectionPageCursor::new( let mut page_cursor = LocalCandidatePreselectionPageCursor::new(
PlannerAppState::new(&fixed_order_app), PlannerAppState::new(&fixed_order_app),
&model_directive_policy,
"openai:chat", "openai:chat",
"gpt-5", "gpt-5",
true, true,
@@ -134,6 +134,7 @@ mod tests {
global_model_id: "global-1".to_string(), global_model_id: "global-1".to_string(),
global_model_name: "gpt-5.4".to_string(), global_model_name: "gpt-5.4".to_string(),
selected_provider_model_name: "gpt-5.4".to_string(), selected_provider_model_name: "gpt-5.4".to_string(),
supports_streaming: true,
mapping_matched_model: None, mapping_matched_model: None,
} }
} }
@@ -169,6 +169,7 @@ mod tests {
global_model_id: "global-model-1".to_string(), global_model_id: "global-model-1".to_string(),
global_model_name: "gpt-test".to_string(), global_model_name: "gpt-test".to_string(),
selected_provider_model_name: "gpt-test-upstream".to_string(), selected_provider_model_name: "gpt-test-upstream".to_string(),
supports_streaming: true,
mapping_matched_model: None, mapping_matched_model: None,
} }
} }
@@ -348,6 +348,7 @@ mod tests {
global_model_id: "global-model-1".to_string(), global_model_id: "global-model-1".to_string(),
global_model_name: "gpt-4.1".to_string(), global_model_name: "gpt-4.1".to_string(),
selected_provider_model_name: "gpt-4.1".to_string(), selected_provider_model_name: "gpt-4.1".to_string(),
supports_streaming: true,
mapping_matched_model: None, mapping_matched_model: None,
} }
} }
@@ -624,6 +625,7 @@ mod tests {
global_model_id: "global-model-1".to_string(), global_model_id: "global-model-1".to_string(),
global_model_name: "gpt-4.1".to_string(), global_model_name: "gpt-4.1".to_string(),
selected_provider_model_name: "gpt-4.1".to_string(), selected_provider_model_name: "gpt-4.1".to_string(),
supports_streaming: true,
mapping_matched_model: None, mapping_matched_model: None,
} }
} }
@@ -643,6 +643,7 @@ mod tests {
global_model_id: "global-model-1".to_string(), global_model_id: "global-model-1".to_string(),
global_model_name: "claude-sonnet".to_string(), global_model_name: "claude-sonnet".to_string(),
selected_provider_model_name: "claude-sonnet".to_string(), selected_provider_model_name: "claude-sonnet".to_string(),
supports_streaming: true,
mapping_matched_model: None, mapping_matched_model: None,
} }
} }
@@ -64,10 +64,25 @@ struct GatewayLocalCandidatePreselectionPort<'a> {
use_api_format_alias_match: bool, use_api_format_alias_match: bool,
key_mode: LocalCandidatePreselectionKeyMode, key_mode: LocalCandidatePreselectionKeyMode,
candidate_api_formats: Vec<String>, candidate_api_formats: Vec<String>,
model_directive_enabled_api_formats: BTreeSet<String>, model_directive_routing_models: BTreeMap<String, String>,
ranking_seed: u64, 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] #[async_trait]
impl AiCandidatePreselectionPort for GatewayLocalCandidatePreselectionPort<'_> { impl AiCandidatePreselectionPort for GatewayLocalCandidatePreselectionPort<'_> {
type Candidate = SchedulerMinimalCandidateSelectionCandidate; type Candidate = SchedulerMinimalCandidateSelectionCandidate;
@@ -99,12 +114,13 @@ impl AiCandidatePreselectionPort for GatewayLocalCandidatePreselectionPort<'_> {
.state .state
.list_selectable_candidates_with_skip_reasons( .list_selectable_candidates_with_skip_reasons(
candidate_api_format, candidate_api_format,
self.requested_model, self.routing_model(candidate_api_format),
self.require_streaming, self.require_streaming,
self.required_capabilities, self.required_capabilities,
auth_snapshot, auth_snapshot,
self.client_session_affinity, self.client_session_affinity,
self.ranking_seed, self.ranking_seed,
false,
) )
.await?; .await?;
@@ -123,16 +139,13 @@ impl AiCandidatePreselectionPort for GatewayLocalCandidatePreselectionPort<'_> {
candidate_api_format: &str, candidate_api_format: &str,
matches_client_format: bool, matches_client_format: bool,
) -> 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) routing_policy_allows_provider(self.routing_policy, candidate)
&& (matches_client_format && (matches_client_format
|| auth_snapshot_allows_cross_format_candidate( || auth_snapshot_allows_cross_format_candidate(
self.auth_snapshot, self.auth_snapshot,
self.requested_model, self.requested_model,
self.model_directive_base_model(candidate_api_format),
candidate, candidate,
enable_model_directives,
)) ))
} }
@@ -142,16 +155,13 @@ impl AiCandidatePreselectionPort for GatewayLocalCandidatePreselectionPort<'_> {
candidate_api_format: &str, candidate_api_format: &str,
matches_client_format: bool, matches_client_format: bool,
) -> 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) routing_policy_allows_provider(self.routing_policy, &skipped_candidate.candidate)
&& (matches_client_format && (matches_client_format
|| auth_snapshot_allows_cross_format_candidate( || auth_snapshot_allows_cross_format_candidate(
self.auth_snapshot, self.auth_snapshot,
self.requested_model, self.requested_model,
self.model_directive_base_model(candidate_api_format),
&skipped_candidate.candidate, &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<String, String> {
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)] #[allow(clippy::too_many_arguments)]
pub(crate) async fn preselect_local_execution_candidates_with_serving( pub(crate) async fn preselect_local_execution_candidates_with_serving(
state: PlannerAppState<'_>, state: PlannerAppState<'_>,
model_directive_policy: &crate::system_features::ModelDirectivePolicySnapshot,
client_api_format: &str, client_api_format: &str,
requested_model: &str, requested_model: &str,
require_streaming: bool, require_streaming: bool,
@@ -190,6 +218,7 @@ pub(crate) async fn preselect_local_execution_candidates_with_serving(
.collect::<Vec<_>>(); .collect::<Vec<_>>();
preselect_local_execution_candidates_for_api_formats_with_serving( preselect_local_execution_candidates_for_api_formats_with_serving(
state, state,
model_directive_policy,
client_api_format, client_api_format,
requested_model, requested_model,
require_streaming, require_streaming,
@@ -207,6 +236,7 @@ pub(crate) async fn preselect_local_execution_candidates_with_serving(
#[allow(clippy::too_many_arguments)] #[allow(clippy::too_many_arguments)]
pub(crate) async fn preselect_local_execution_candidates_for_api_formats_with_serving( pub(crate) async fn preselect_local_execution_candidates_for_api_formats_with_serving(
state: PlannerAppState<'_>, state: PlannerAppState<'_>,
model_directive_policy: &crate::system_features::ModelDirectivePolicySnapshot,
client_api_format: &str, client_api_format: &str,
requested_model: &str, requested_model: &str,
require_streaming: bool, require_streaming: bool,
@@ -224,19 +254,11 @@ pub(crate) async fn preselect_local_execution_candidates_for_api_formats_with_se
>, >,
GatewayError, GatewayError,
> { > {
let mut model_directive_enabled_api_formats = BTreeSet::new(); let model_directive_routing_models = resolve_model_directive_routing_models(
for api_format in &candidate_api_formats { model_directive_policy,
if crate::system_features::reasoning_model_directive_enabled_for_api_format_and_model( &candidate_api_formats,
state.app(), requested_model,
api_format, );
Some(requested_model),
)
.await
{
model_directive_enabled_api_formats
.insert(crate::ai_serving::normalize_api_format_alias(api_format));
}
}
let port = GatewayLocalCandidatePreselectionPort { let port = GatewayLocalCandidatePreselectionPort {
state, state,
client_api_format, 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, use_api_format_alias_match,
key_mode, key_mode,
candidate_api_formats, candidate_api_formats,
model_directive_enabled_api_formats, model_directive_routing_models,
ranking_seed: request_distribution_seed(), ranking_seed: request_distribution_seed(),
}; };
@@ -271,7 +293,8 @@ pub(crate) struct LocalCandidatePreselectionPageCursor<'a> {
key_mode: LocalCandidatePreselectionKeyMode, key_mode: LocalCandidatePreselectionKeyMode,
allow_priority_page_cache: bool, allow_priority_page_cache: bool,
candidate_api_formats: Vec<String>, candidate_api_formats: Vec<String>,
model_directive_enabled_api_formats: BTreeSet<String>, model_directive_routing_models: BTreeMap<String, String>,
model_directive_policy_cache_key: String,
ordering_config: SchedulerOrderingConfig, ordering_config: SchedulerOrderingConfig,
ranking_seed: u64, ranking_seed: u64,
priority_page_emitted: bool, priority_page_emitted: bool,
@@ -294,9 +317,23 @@ pub(crate) struct LocalCandidatePreselectionPageCursor<'a> {
} }
impl<'a> 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)] #[allow(clippy::too_many_arguments)]
pub(crate) async fn new( pub(crate) async fn new(
state: PlannerAppState<'a>, state: PlannerAppState<'a>,
model_directive_policy: &crate::system_features::ModelDirectivePolicySnapshot,
client_api_format: &str, client_api_format: &str,
requested_model: &str, requested_model: &str,
require_streaming: bool, require_streaming: bool,
@@ -315,19 +352,11 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
.into_iter() .into_iter()
.map(str::to_string) .map(str::to_string)
.collect::<Vec<_>>(); .collect::<Vec<_>>();
let mut model_directive_enabled_api_formats = BTreeSet::new(); let model_directive_routing_models = resolve_model_directive_routing_models(
for api_format in &candidate_api_formats { model_directive_policy,
if crate::system_features::reasoning_model_directive_enabled_for_api_format_and_model( &candidate_api_formats,
state.app(), requested_model,
api_format, );
Some(requested_model),
)
.await
{
model_directive_enabled_api_formats
.insert(crate::ai_serving::normalize_api_format_alias(api_format));
}
}
let ordering_config = let ordering_config =
super::candidate_ranking::scheduler_ordering_config_for_routing_policy( super::candidate_ranking::scheduler_ordering_config_for_routing_policy(
@@ -351,7 +380,8 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
key_mode, key_mode,
allow_priority_page_cache, allow_priority_page_cache,
candidate_api_formats, 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, ordering_config,
ranking_seed: request_distribution_seed(), ranking_seed: request_distribution_seed(),
priority_page_emitted: false, priority_page_emitted: false,
@@ -426,6 +456,10 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
self.use_api_format_alias_match 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 { pub(crate) fn should_cache_current_priority_resolved_page(&self) -> bool {
if !(self.priority_page_emitted if !(self.priority_page_emitted
&& self.format_index == 0 && self.format_index == 0
@@ -491,6 +525,7 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
self.key_mode.cache_key_name(), self.key_mode.cache_key_name(),
self.use_api_format_alias_match, self.use_api_format_alias_match,
self.client_session_affinity.as_ref(), self.client_session_affinity.as_ref(),
&self.model_directive_policy_cache_key,
); );
let cache = self.state.app().candidate_page_cache.clone(); let cache = self.state.app().candidate_page_cache.clone();
let ttl = candidate_page_cache_ttl_from_env(); let ttl = candidate_page_cache_ttl_from_env();
@@ -752,11 +787,8 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
if normalized_api_format.is_empty() { if normalized_api_format.is_empty() {
return Ok(None); return Ok(None);
} }
let enable_model_directives = self.model_directive_enabled_api_formats.contains( let routing_model = self.routing_model(candidate_api_format).to_string();
&crate::ai_serving::normalize_api_format_alias(candidate_api_format), let requested_names = requested_model_candidate_names(&routing_model, false);
);
let requested_names =
requested_model_candidate_names(&self.requested_model, enable_model_directives);
let scanned = *self let scanned = *self
.scanned_rows_by_format .scanned_rows_by_format
.get(&normalized_api_format) .get(&normalized_api_format)
@@ -772,11 +804,7 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
.or_insert(0); .or_insert(0);
let Some(requested_name) = requested_names.get(requested_name_index) else { let Some(requested_name) = requested_names.get(requested_name_index) else {
return self return self
.next_fallback_page_for_api_format( .next_fallback_page_for_api_format(candidate_api_format, &normalized_api_format)
candidate_api_format,
&normalized_api_format,
enable_model_directives,
)
.await; .await;
}; };
if requested_name.trim().is_empty() { if requested_name.trim().is_empty() {
@@ -803,9 +831,9 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
.read_requested_model_rows_fast_path_page_cached( .read_requested_model_rows_fast_path_page_cached(
&normalized_api_format, &normalized_api_format,
requested_name, requested_name,
&routing_model,
offset, offset,
limit, limit,
enable_model_directives,
) )
.await?; .await?;
self.scanned_rows_by_format.insert( self.scanned_rows_by_format.insert(
@@ -824,7 +852,6 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
.next_fallback_page_for_api_format( .next_fallback_page_for_api_format(
candidate_api_format, candidate_api_format,
&normalized_api_format, &normalized_api_format,
enable_model_directives,
) )
.await; .await;
} }
@@ -835,7 +862,6 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
.build_page_outcome_from_rows( .build_page_outcome_from_rows(
candidate_api_format, candidate_api_format,
&normalized_api_format, &normalized_api_format,
enable_model_directives,
page.rows, page.rows,
) )
.await? .await?
@@ -849,17 +875,17 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
&self, &self,
normalized_api_format: &str, normalized_api_format: &str,
requested_name: &str, requested_name: &str,
routing_model: &str,
offset: u32, offset: u32,
limit: u32, limit: u32,
enable_model_directives: bool,
) -> Result<RequestedModelCandidateRowsPage, GatewayError> { ) -> Result<RequestedModelCandidateRowsPage, GatewayError> {
let key = CandidateRowPageCacheKey::new( let key = CandidateRowPageCacheKey::new(
normalized_api_format, normalized_api_format,
&self.requested_model, routing_model,
requested_name, requested_name,
offset, offset,
limit, limit,
enable_model_directives, false,
); );
let cache = self.state.app().candidate_row_page_cache.clone(); let cache = self.state.app().candidate_row_page_cache.clone();
let ttl = candidate_page_cache_ttl_from_env(); 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( let page = read_requested_model_rows_fast_path_page(
self.state.app().data.as_ref(), self.state.app().data.as_ref(),
normalized_api_format, normalized_api_format,
&self.requested_model, routing_model,
requested_name, requested_name,
offset, offset,
limit, limit,
enable_model_directives, false,
) )
.await .await
.map_err(|err| GatewayError::Internal(err.to_string()))?; .map_err(|err| GatewayError::Internal(err.to_string()))?;
@@ -913,7 +939,6 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
&mut self, &mut self,
candidate_api_format: &str, candidate_api_format: &str,
normalized_api_format: &str, normalized_api_format: &str,
enable_model_directives: bool,
) -> Result< ) -> Result<
Option< Option<
AiCandidatePreselectionOutcome< AiCandidatePreselectionOutcome<
@@ -930,6 +955,7 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
return Ok(None); return Ok(None);
} }
let routing_model = self.routing_model(candidate_api_format).to_string();
let rows = self let rows = self
.state .state
.app() .app()
@@ -941,27 +967,21 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
.filter(|row| { .filter(|row| {
row_supports_requested_model_with_model_directives( row_supports_requested_model_with_model_directives(
row, row,
&self.requested_model, &routing_model,
normalized_api_format, normalized_api_format,
enable_model_directives, false,
) )
}) })
.collect::<Vec<_>>(); .collect::<Vec<_>>();
self.build_page_outcome_from_rows( self.build_page_outcome_from_rows(candidate_api_format, normalized_api_format, rows)
candidate_api_format, .await
normalized_api_format,
enable_model_directives,
rows,
)
.await
} }
async fn build_page_outcome_from_rows( async fn build_page_outcome_from_rows(
&mut self, &mut self,
candidate_api_format: &str, candidate_api_format: &str,
normalized_api_format: &str, normalized_api_format: &str,
enable_model_directives: bool,
rows: Vec<StoredMinimalCandidateSelectionRow>, rows: Vec<StoredMinimalCandidateSelectionRow>,
) -> Result< ) -> Result<
Option< Option<
@@ -984,15 +1004,16 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
if rows.is_empty() { if rows.is_empty() {
return Ok(None); return Ok(None);
} }
let routing_model = self.routing_model(candidate_api_format).to_string();
let resolved_global_model_name = let resolved_global_model_name =
if let Some(value) = self.resolved_global_model_names.get(normalized_api_format) { if let Some(value) = self.resolved_global_model_names.get(normalized_api_format) {
value.clone() value.clone()
} else { } else {
let Some(value) = resolve_requested_global_model_name_with_model_directives( let Some(value) = resolve_requested_global_model_name_with_model_directives(
&rows, &rows,
&self.requested_model, &routing_model,
normalized_api_format, normalized_api_format,
enable_model_directives, false,
) else { ) else {
return Ok(None); return Ok(None);
}; };
@@ -1016,22 +1037,18 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
EnumerateMinimalCandidateSelectionInput { EnumerateMinimalCandidateSelectionInput {
rows, rows,
normalized_api_format, normalized_api_format,
requested_model_name: &self.requested_model, requested_model_name: &routing_model,
resolved_global_model_name: resolved_global_model_name.as_str(), resolved_global_model_name: resolved_global_model_name.as_str(),
require_streaming: self.require_streaming, require_streaming: self.require_streaming,
required_capabilities: self.required_capabilities.as_ref(), required_capabilities: self.required_capabilities.as_ref(),
auth_constraints: auth_constraints.as_ref(), auth_constraints: auth_constraints.as_ref(),
}, },
enable_model_directives, false,
) )
.map_err(|err| GatewayError::Internal(err.to_string()))?; .map_err(|err| GatewayError::Internal(err.to_string()))?;
let mut candidates = Vec::new(); let mut candidates = Vec::new();
for candidate in enumerated_candidates { for candidate in enumerated_candidates {
if !self.candidate_allowed_for_page( if !self.candidate_allowed_for_page(&candidate, candidate_api_format) {
&candidate,
candidate_api_format,
enable_model_directives,
) {
continue; continue;
} }
if !self if !self
@@ -1065,11 +1082,7 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
.into_iter() .into_iter()
.map(skipped_local_execution_candidate_from_scheduler_skip) .map(skipped_local_execution_candidate_from_scheduler_skip)
.filter(|skipped_candidate| { .filter(|skipped_candidate| {
self.skipped_candidate_allowed_for_page( self.skipped_candidate_allowed_for_page(skipped_candidate, candidate_api_format)
skipped_candidate,
candidate_api_format,
enable_model_directives,
)
}) })
.collect::<Vec<_>>(); .collect::<Vec<_>>();
@@ -1083,7 +1096,6 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
&self, &self,
candidate: &SchedulerMinimalCandidateSelectionCandidate, candidate: &SchedulerMinimalCandidateSelectionCandidate,
candidate_api_format: &str, candidate_api_format: &str,
enable_model_directives: bool,
) -> bool { ) -> bool {
routing_policy_allows_provider(self.routing_policy.as_ref(), candidate) routing_policy_allows_provider(self.routing_policy.as_ref(), candidate)
&& (matches_client_api_format( && (matches_client_api_format(
@@ -1093,8 +1105,8 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
) || auth_snapshot_allows_cross_format_candidate( ) || auth_snapshot_allows_cross_format_candidate(
&self.auth_snapshot, &self.auth_snapshot,
&self.requested_model, &self.requested_model,
self.model_directive_base_model(candidate_api_format),
candidate, candidate,
enable_model_directives,
)) ))
} }
@@ -1102,7 +1114,6 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
&self, &self,
skipped_candidate: &SkippedLocalExecutionCandidate, skipped_candidate: &SkippedLocalExecutionCandidate,
candidate_api_format: &str, candidate_api_format: &str,
enable_model_directives: bool,
) -> bool { ) -> bool {
routing_policy_allows_provider(self.routing_policy.as_ref(), &skipped_candidate.candidate) routing_policy_allows_provider(self.routing_policy.as_ref(), &skipped_candidate.candidate)
&& (matches_client_api_format( && (matches_client_api_format(
@@ -1112,8 +1123,8 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
) || auth_snapshot_allows_cross_format_candidate( ) || auth_snapshot_allows_cross_format_candidate(
&self.auth_snapshot, &self.auth_snapshot,
&self.requested_model, &self.requested_model,
self.model_directive_base_model(candidate_api_format),
&skipped_candidate.candidate, &skipped_candidate.candidate,
enable_model_directives,
)) ))
} }
} }
@@ -1199,8 +1210,8 @@ fn matches_client_api_format(
pub(crate) fn auth_snapshot_allows_cross_format_candidate( pub(crate) fn auth_snapshot_allows_cross_format_candidate(
auth_snapshot: &GatewayAuthApiKeySnapshot, auth_snapshot: &GatewayAuthApiKeySnapshot,
requested_model: &str, requested_model: &str,
requested_base_model: Option<&str>,
candidate: &SchedulerMinimalCandidateSelectionCandidate, candidate: &SchedulerMinimalCandidateSelectionCandidate,
enable_model_directives: bool,
) -> bool { ) -> bool {
if let Some(allowed_providers) = auth_snapshot.effective_allowed_providers() { if let Some(allowed_providers) = auth_snapshot.effective_allowed_providers() {
let provider_allowed = allowed_providers.iter().any(|value| { 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() { 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| { let model_allowed = allowed_models.iter().any(|value| {
value == requested_model value == requested_model
|| value == &candidate.global_model_name || value == &candidate.global_model_name
|| requested_base_model || requested_base_model.is_some_and(|base_model| value == base_model)
.as_ref()
.is_some_and(|base_model| value == base_model)
}); });
if !model_allowed { if !model_allowed {
return false; return false;
@@ -1303,8 +1309,11 @@ mod tests {
.expect("gateway state should build") .expect("gateway state should build")
.with_data_state_for_tests(data_state); .with_data_state_for_tests(data_state);
let auth_snapshot = unrestricted_auth_snapshot(); let auth_snapshot = unrestricted_auth_snapshot();
let model_directive_policy =
crate::system_features::ModelDirectivePolicySnapshot::load(&app).await;
let mut cursor = LocalCandidatePreselectionPageCursor::new( let mut cursor = LocalCandidatePreselectionPageCursor::new(
PlannerAppState::new(&app), PlannerAppState::new(&app),
&model_directive_policy,
"openai:chat", "openai:chat",
"gpt-5", "gpt-5",
true, true,
@@ -1545,8 +1554,11 @@ mod tests {
.expect("gateway state should build") .expect("gateway state should build")
.with_data_state_for_tests(data_state); .with_data_state_for_tests(data_state);
let auth_snapshot = unrestricted_auth_snapshot(); let auth_snapshot = unrestricted_auth_snapshot();
let model_directive_policy =
crate::system_features::ModelDirectivePolicySnapshot::load(&app).await;
let mut cursor = LocalCandidatePreselectionPageCursor::new( let mut cursor = LocalCandidatePreselectionPageCursor::new(
PlannerAppState::new(&app), PlannerAppState::new(&app),
&model_directive_policy,
"claude:messages", "claude:messages",
"gpt-5.5-xhigh", "gpt-5.5-xhigh",
false, 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<dyn MinimalCandidateSelectionReadRepository> =
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] #[tokio::test]
async fn claude_request_uses_cross_format_key_when_same_provider_messages_key_lacks_model() { async fn claude_request_uses_cross_format_key_when_same_provider_messages_key_lacks_model() {
let repository: Arc<dyn MinimalCandidateSelectionReadRepository> = let repository: Arc<dyn MinimalCandidateSelectionReadRepository> =
@@ -1605,8 +1688,11 @@ mod tests {
.expect("gateway state should build") .expect("gateway state should build")
.with_data_state_for_tests(data_state); .with_data_state_for_tests(data_state);
let auth_snapshot = unrestricted_auth_snapshot(); let auth_snapshot = unrestricted_auth_snapshot();
let model_directive_policy =
crate::system_features::ModelDirectivePolicySnapshot::load(&app).await;
let mut cursor = LocalCandidatePreselectionPageCursor::new( let mut cursor = LocalCandidatePreselectionPageCursor::new(
PlannerAppState::new(&app), PlannerAppState::new(&app),
&model_directive_policy,
"claude:messages", "claude:messages",
"deepseek-v4-pro", "deepseek-v4-pro",
false, false,
@@ -1680,8 +1766,11 @@ mod tests {
.expect("gateway state should build") .expect("gateway state should build")
.with_data_state_for_tests(data_state); .with_data_state_for_tests(data_state);
let auth_snapshot = unrestricted_auth_snapshot(); let auth_snapshot = unrestricted_auth_snapshot();
let model_directive_policy =
crate::system_features::ModelDirectivePolicySnapshot::load(&app).await;
let mut cursor = LocalCandidatePreselectionPageCursor::new( let mut cursor = LocalCandidatePreselectionPageCursor::new(
PlannerAppState::new(&app), PlannerAppState::new(&app),
&model_directive_policy,
"claude:messages", "claude:messages",
"gpt-5", "gpt-5",
false, false,
@@ -1,11 +1,12 @@
use axum::body::Bytes; use axum::body::Bytes;
use crate::ai_serving::is_json_request;
use crate::ai_serving::{ use crate::ai_serving::{
endpoint_config_forces_upstream_stream_policy as endpoint_config_forces_upstream_stream_policy_impl, 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, enforce_request_body_stream_field as enforce_request_body_stream_field_impl,
force_upstream_streaming_for_provider as force_upstream_streaming_for_provider_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, 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, resolve_format_upstream_is_stream_for_provider as resolve_upstream_is_stream_for_provider_impl,
}; };
pub(crate) use crate::ai_serving::{ pub(crate) use crate::ai_serving::{
CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND, CLAUDE_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_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_RERANK_SYNC_PLAN_KIND, OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND,
OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND, OPENAI_RESPONSES_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_RESPONSES_SYNC_PLAN_KIND, OPENAI_SEARCH_SYNC_PLAN_KIND,
OPENAI_VIDEO_CONTENT_PLAN_KIND, OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_VIDEO_CONTENT_PLAN_KIND,
OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND, OPENAI_VIDEO_REMIX_SYNC_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; 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, client_is_stream: bool,
hard_requires_streaming: bool, hard_requires_streaming: bool,
) -> bool { ) -> bool {
let hard_requires_streaming = hard_requires_streaming resolve_upstream_is_stream_for_provider_impl(
|| force_upstream_streaming_for_provider(provider_type, provider_api_format);
resolve_upstream_is_stream_from_endpoint_config_impl(
endpoint_config, endpoint_config,
provider_type,
provider_api_format,
client_is_stream, client_is_stream,
hard_requires_streaming, hard_requires_streaming,
) )
@@ -178,6 +180,27 @@ mod tests {
true, true,
false, 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] #[test]
@@ -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_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_RERANK_SYNC_PLAN_KIND, OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND,
OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND, OPENAI_RESPONSES_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_RESPONSES_SYNC_PLAN_KIND, OPENAI_SEARCH_SYNC_PLAN_KIND,
OPENAI_VIDEO_CONTENT_PLAN_KIND, OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_VIDEO_CONTENT_PLAN_KIND,
OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND, OPENAI_VIDEO_REMIX_SYNC_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::{ use crate::ai_serving::planner::plan_builders::{
build_gemini_stream_plan_from_decision, build_gemini_sync_plan_from_decision, 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 => { OPENAI_RESPONSES_SYNC_PLAN_KIND => {
build_openai_responses_sync_plan_from_decision(parts, body_json, payload, false)? 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 => { OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND => {
build_openai_responses_sync_plan_from_decision(parts, body_json, payload, true)? build_openai_responses_sync_plan_from_decision(parts, body_json, payload, true)?
} }
@@ -14,8 +14,11 @@ use serde_json::{json, Value};
use tracing::warn; use tracing::warn;
use crate::ai_serving::planner::common::extract_standard_requested_model; use crate::ai_serving::planner::common::extract_standard_requested_model;
use crate::ai_serving::{ExecutionRuntimeAuthContext, GatewayAuthApiKeySnapshot, PlannerAppState}; use crate::ai_serving::{
use crate::client_session_affinity::client_session_affinity_from_request; 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::clock::current_unix_secs;
use crate::routing::{ use crate::routing::{
apply_routing_mutation_plan, build_routing_trace_seed, resolve_gateway_routing_policy, 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}; use crate::{AiExecutionDecision, AppState, GatewayError};
const ROUTING_GROUP_SELECTION_CACHE_TTL: Duration = Duration::from_secs(30); 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)] #[derive(Debug, Clone)]
pub(crate) struct ResolvedLocalDecisionAuthInput { pub(crate) struct ResolvedLocalDecisionAuthInput {
pub(crate) auth_context: ExecutionRuntimeAuthContext, pub(crate) auth_context: ExecutionRuntimeAuthContext,
pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot, pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot,
pub(crate) required_capabilities: Option<serde_json::Value>, pub(crate) required_capabilities: Option<serde_json::Value>,
pub(crate) model_directive_policy: crate::system_features::ModelDirectivePolicySnapshot,
} }
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
@@ -46,6 +52,7 @@ pub(crate) struct LocalRequestedModelDecisionInput {
pub(crate) routing_policy: Option<ResolvedRoutingPolicy>, pub(crate) routing_policy: Option<ResolvedRoutingPolicy>,
pub(crate) routing_trace_seed: Option<RoutingDecisionTrace>, pub(crate) routing_trace_seed: Option<RoutingDecisionTrace>,
pub(crate) routing_context: Option<LocalRoutingRequestContext>, pub(crate) routing_context: Option<LocalRoutingRequestContext>,
pub(crate) model_directive_policy: crate::system_features::ModelDirectivePolicySnapshot,
} }
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
@@ -86,14 +93,52 @@ impl LocalRequestedModelDecisionInput {
pub(crate) fn apply_provider_request_routing_policy_to_decision( pub(crate) fn apply_provider_request_routing_policy_to_decision(
input: &LocalRequestedModelDecisionInput, input: &LocalRequestedModelDecisionInput,
decision: &mut AiExecutionDecision, decision: &mut AiExecutionDecision,
transport: Option<&GatewayProviderTransportSnapshot>,
) -> Result<(), GatewayError> { ) -> 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 { let Some(context) = input.routing_context.as_ref() else {
return Ok(()); return Ok(());
}; };
let provider_api_format = decision let provider_body_rules = decision
.provider_api_format .report_context
.as_deref() .as_ref()
.unwrap_or(context.client_api_format.as_str()); .and_then(|context| context.get("body_rules"))
.cloned();
let resolved_model = decision let resolved_model = decision
.mapped_model .mapped_model
.as_deref() .as_deref()
@@ -104,6 +149,21 @@ pub(crate) fn apply_provider_request_routing_policy_to_decision(
.clone() .clone()
.unwrap_or(serde_json::Value::Null); .unwrap_or(serde_json::Value::Null);
let mut provider_headers = btree_headers_to_header_map(&decision.provider_request_headers)?; 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::<Vec<_>>();
let provider_headers_json = headers_to_routing_value(&provider_headers); let provider_headers_json = headers_to_routing_value(&provider_headers);
let policy = resolve_gateway_routing_policy(GatewayRoutingPolicyInput { let policy = resolve_gateway_routing_policy(GatewayRoutingPolicyInput {
group_id: context.group_id.as_deref(), 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(), selection_source: context.selection_source.as_str(),
requested_model: input.requested_model.as_str(), requested_model: input.requested_model.as_str(),
resolved_model, resolved_model,
api_format: provider_api_format, api_format: provider_api_format.as_str(),
user_id: Some(input.auth_context.user_id.as_str()), user_id: Some(input.auth_context.user_id.as_str()),
api_key_id: Some(input.auth_context.api_key_id.as_str()), api_key_id: Some(input.auth_context.api_key_id.as_str()),
headers: &provider_headers_json, headers: &provider_headers_json,
@@ -134,6 +194,80 @@ pub(crate) fn apply_provider_request_routing_policy_to_decision(
&mut provider_headers, &mut provider_headers,
&policy.mutation_plan, &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); decision.provider_request_headers = header_map_to_btree_headers(&provider_headers);
if original_provider_request_body.is_some() { if original_provider_request_body.is_some() {
decision.provider_request_body = Some(provider_request_body); 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> { struct GatewayAuthenticatedDecisionInputPort<'a> {
state: PlannerAppState<'a>, state: PlannerAppState<'a>,
now_unix_secs: u64, now_unix_secs: u64,
model_directive_policy: &'a crate::system_features::ModelDirectivePolicySnapshot,
model_directive_base_model: Option<String>,
} }
#[async_trait] #[async_trait]
@@ -181,6 +317,7 @@ impl AiAuthenticatedDecisionInputPort for GatewayAuthenticatedDecisionInputPort<
&auth_context.api_key_id, &auth_context.api_key_id,
requested_model, requested_model,
explicit_required_capabilities, explicit_required_capabilities,
self.model_directive_base_model.as_deref(),
) )
.await) .await)
} }
@@ -195,6 +332,7 @@ impl AiAuthenticatedDecisionInputPort for GatewayAuthenticatedDecisionInputPort<
auth_context, auth_context,
auth_snapshot, auth_snapshot,
required_capabilities, 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_policy: None,
routing_trace_seed: None, routing_trace_seed: None,
routing_context: 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 let Some((group_id, group_version, group_config_json, selection_source)) = selected_group
else { else {
input.client_session_affinity = input.client_session_affinity = client_session_affinity_from_api_request(
client_session_affinity_from_request(&parts.headers, Some(body_json)); client_api_format,
&parts.headers,
Some(body_json),
);
input.routing_policy = None; input.routing_policy = None;
input.routing_trace_seed = None; input.routing_trace_seed = None;
input.routing_context = 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 { 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) input.required_capabilities = PlannerAppState::new(state)
.resolve_request_candidate_required_capabilities( .resolve_request_candidate_required_capabilities(
&input.auth_context.user_id, &input.auth_context.user_id,
&input.auth_context.api_key_id, &input.auth_context.api_key_id,
Some(input.requested_model.as_str()), Some(input.requested_model.as_str()),
input.required_capabilities.as_ref(), input.required_capabilities.as_ref(),
model_directive_resolution.base_model(),
) )
.await; .await;
} }
let effective_headers_json = headers_to_routing_value(&effective_headers); let effective_headers_json = headers_to_routing_value(&effective_headers);
input.client_session_affinity = input.client_session_affinity = client_session_affinity_from_api_request(
client_session_affinity_from_request(&effective_headers, Some(&effective_body_json)); client_api_format,
&effective_headers,
Some(&effective_body_json),
);
let final_policy_resolve_started_at = std::time::Instant::now(); let final_policy_resolve_started_at = std::time::Instant::now();
let mut final_policy = resolve_gateway_routing_policy(GatewayRoutingPolicyInput { let mut final_policy = resolve_gateway_routing_policy(GatewayRoutingPolicyInput {
group_id: group_id.as_deref(), 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, static_policy_resolve_started_at.elapsed().as_millis() as u64,
); );
input.client_session_affinity = input.client_session_affinity = client_session_affinity_from_api_request(
client_session_affinity_from_request(&parts.headers, Some(body_json)); client_api_format,
&parts.headers,
Some(body_json),
);
input.routing_trace_seed = Some(build_routing_trace_seed(&policy, client_api_format)); input.routing_trace_seed = Some(build_routing_trace_seed(&policy, client_api_format));
input.routing_policy = Some(policy); input.routing_policy = Some(policy);
input.routing_context = None; input.routing_context = None;
@@ -475,11 +627,22 @@ pub(crate) async fn resolve_local_authenticated_decision_input(
state: &AppState, state: &AppState,
auth_context: ExecutionRuntimeAuthContext, auth_context: ExecutionRuntimeAuthContext,
requested_model: Option<&str>, requested_model: Option<&str>,
requested_model_api_format: Option<&str>,
explicit_required_capabilities: Option<&serde_json::Value>, explicit_required_capabilities: Option<&serde_json::Value>,
model_directive_policy: &crate::system_features::ModelDirectivePolicySnapshot,
) -> Result<Option<ResolvedLocalDecisionAuthInput>, GatewayError> { ) -> Result<Option<ResolvedLocalDecisionAuthInput>, 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 { let port = GatewayAuthenticatedDecisionInputPort {
state: PlannerAppState::new(state), state: PlannerAppState::new(state),
now_unix_secs: current_unix_secs(), now_unix_secs: current_unix_secs(),
model_directive_policy,
model_directive_base_model,
}; };
run_ai_authenticated_decision_input( run_ai_authenticated_decision_input(
@@ -730,6 +893,10 @@ fn ensure_report_context_routing_trace(
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
use aether_provider_transport::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
GatewayProviderTransportProvider,
};
fn sample_auth_context() -> ExecutionRuntimeAuthContext { fn sample_auth_context() -> ExecutionRuntimeAuthContext {
ExecutionRuntimeAuthContext { ExecutionRuntimeAuthContext {
@@ -782,6 +949,7 @@ mod tests {
client_session_affinity: None, client_session_affinity: None,
routing_policy: None, routing_policy: None,
routing_trace_seed: None, routing_trace_seed: None,
model_directive_policy: Default::default(),
routing_context: Some(LocalRoutingRequestContext { routing_context: Some(LocalRoutingRequestContext {
group_id: Some("group-1".to_string()), group_id: Some("group-1".to_string()),
group_version: Some(3), group_version: Some(3),
@@ -830,6 +998,7 @@ mod tests {
request_id: Some("trace-1".to_string()), request_id: Some("trace-1".to_string()),
candidate_id: Some("candidate-1".to_string()), candidate_id: Some("candidate-1".to_string()),
provider_name: Some("provider".to_string()), provider_name: Some("provider".to_string()),
provider_type: Some("openai".to_string()),
provider_id: Some("provider-1".to_string()), provider_id: Some("provider-1".to_string()),
endpoint_id: Some("endpoint-1".to_string()), endpoint_id: Some("endpoint-1".to_string()),
key_id: Some("key-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!({ let config = json!({
"allowed_models": ["gpt-5"], "allowed_models": allowed_models,
"rules": [{ "rules": [{
"id": "provider-patch", "id": "provider-patch",
"priority": 1, "priority": 1,
@@ -904,6 +1144,7 @@ mod tests {
client_session_affinity: None, client_session_affinity: None,
routing_policy: None, routing_policy: None,
routing_trace_seed: None, routing_trace_seed: None,
model_directive_policy: Default::default(),
routing_context: Some(LocalRoutingRequestContext { routing_context: Some(LocalRoutingRequestContext {
group_id: Some("stale".to_string()), group_id: Some("stale".to_string()),
group_version: Some(1), group_version: Some(1),
@@ -971,6 +1212,7 @@ mod tests {
routing_policy: None, routing_policy: None,
routing_trace_seed: None, routing_trace_seed: None,
routing_context: None, routing_context: None,
model_directive_policy: Default::default(),
}; };
let group_config_json = json!({ let group_config_json = json!({
"rules": [{ "rules": [{
@@ -1006,7 +1248,7 @@ mod tests {
let input = sample_decision_input(); let input = sample_decision_input();
let mut decision = sample_decision(); 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"); .expect("provider routing mutation should apply");
assert_eq!( 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] #[test]
fn provider_request_routing_policy_rejects_body_patch_without_json_body() { fn provider_request_routing_policy_rejects_body_patch_without_json_body() {
let input = sample_decision_input(); let input = sample_decision_input();
@@ -1042,7 +1457,7 @@ mod tests {
decision.provider_request_body = None; decision.provider_request_body = None;
decision.provider_request_body_base64 = Some("AA==".to_string()); 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"); .expect_err("provider body patch should reject binary upstream bodies");
match error { match error {
@@ -1067,6 +1482,7 @@ mod tests {
let mut input = sample_decision_input(); let mut input = sample_decision_input();
set_provider_request_rules( set_provider_request_rules(
&mut input, &mut input,
&["gpt-5"],
json!([{ json!([{
"type": "patch_headers", "type": "patch_headers",
"patch": [{ "patch": [{
@@ -1080,7 +1496,7 @@ mod tests {
decision.provider_request_body = None; decision.provider_request_body = None;
decision.provider_request_body_base64 = Some("AA==".to_string()); 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"); .expect("header-only provider routing mutation should apply without JSON body");
assert_eq!(decision.provider_request_body, None); assert_eq!(decision.provider_request_body, None);
@@ -1112,7 +1528,7 @@ mod tests {
"priority_slot": 3 "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"); .expect("provider routing mutation should seed pool trace");
let routing_trace = &decision.report_context.as_ref().unwrap()["routing_trace"]; let routing_trace = &decision.report_context.as_ref().unwrap()["routing_trace"];
@@ -34,6 +34,7 @@ pub(crate) use self::candidate_resolution::{
candidate_auth_channel_skip_reason, read_candidate_transport_snapshot, candidate_auth_channel_skip_reason, read_candidate_transport_snapshot,
EligibleLocalExecutionCandidate, LocalExecutionCandidateKind, SkippedLocalExecutionCandidate, EligibleLocalExecutionCandidate, LocalExecutionCandidateKind, SkippedLocalExecutionCandidate,
}; };
pub(crate) use self::common::resolve_upstream_is_stream_for_provider;
pub(crate) use self::passthrough::{ pub(crate) use self::passthrough::{
build_local_same_format_stream_attempt_source, build_local_same_format_stream_plan_and_reports, 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, 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::{ pub(crate) use self::pool_scores::{
build_provider_key_pool_score_upsert, provider_key_pool_score_id, provider_key_pool_score_scope, 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::route::is_matching_stream_request as planner_is_matching_stream_request;
pub(crate) use self::runtime_miss::{ pub(crate) use self::runtime_miss::{
apply_local_runtime_candidate_terminal_reason, record_local_runtime_candidate_skip_reason, 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_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_attempt_source as build_standard_family_sync_attempt_source,
build_local_sync_plan_and_reports as build_standard_family_sync_plan_and_reports, 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::{ pub(crate) use self::state::{
GatewayAuthApiKeySnapshot, GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth, GatewayAuthApiKeySnapshot, GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth,
@@ -71,6 +71,7 @@ pub(crate) fn build_passthrough_stream_plan_from_decision(
.content_type .content_type
.take() .take()
.or_else(|| provider_request_headers.get("content-type").cloned()); .or_else(|| provider_request_headers.get("content-type").cloned());
let stream = payload.upstream_is_stream;
let plan = build_ai_execution_plan_from_decision( let plan = build_ai_execution_plan_from_decision(
&mut payload, &mut payload,
AiExecutionPlanFromDecisionParts { AiExecutionPlanFromDecisionParts {
@@ -84,7 +85,7 @@ pub(crate) fn build_passthrough_stream_plan_from_decision(
body_bytes_b64: None, body_bytes_b64: None,
body_ref: None, body_ref: None,
}, },
stream: true, stream,
}, },
); );
@@ -24,7 +24,7 @@ use crate::ai_serving::{
ai_local_execution_contract_for_formats, extract_pool_sticky_session_token, ai_local_execution_contract_for_formats, extract_pool_sticky_session_token,
resolve_local_decision_execution_runtime_auth_context, GatewayControlDecision, PlannerAppState, 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::clock::current_unix_secs;
use crate::{AppState, GatewayError}; use crate::{AppState, GatewayError};
@@ -60,7 +60,9 @@ pub(crate) async fn resolve_local_same_format_provider_decision_input(
state, state,
auth_context, auth_context,
Some(requested_model.as_str()), Some(requested_model.as_str()),
decision.auth_endpoint_signature.as_deref(),
None, None,
&decision.model_directive_policy,
) )
.await .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); let mut input = build_local_requested_model_decision_input(resolved_input, requested_model);
input.request_auth_channel = decision.request_auth_channel.clone(); 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( if let Err(err) = attach_routing_policy_to_local_requested_model_input(
state, state,
parts, parts,
@@ -115,15 +121,22 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
input.required_capabilities.as_ref(), input.required_capabilities.as_ref(),
LocalCandidatePersistencePolicyKind::SameFormatProviderDecision, 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 let (candidates, preselection_skipped) = planner_state
.list_selectable_candidates_with_skip_reasons( .list_selectable_candidates_with_skip_reasons(
spec_metadata.api_format, spec_metadata.api_format,
&input.requested_model, routing_model,
spec_metadata.require_streaming, spec_metadata.require_streaming,
input.required_capabilities.as_ref(), input.required_capabilities.as_ref(),
Some(&input.auth_snapshot), Some(&input.auth_snapshot),
input.client_session_affinity.as_ref(), input.client_session_affinity.as_ref(),
current_unix_secs(), current_unix_secs(),
false,
) )
.await?; .await?;
let outcome = materialize_local_execution_candidates_with_serving( 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(), input.required_capabilities.as_ref(),
LocalCandidatePersistencePolicyKind::SameFormatProviderDecision, 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 let (candidates, preselection_skipped) = planner_state
.list_selectable_candidates_with_skip_reasons( .list_selectable_candidates_with_skip_reasons(
spec_metadata.api_format, spec_metadata.api_format,
&input.requested_model, routing_model,
spec_metadata.require_streaming, spec_metadata.require_streaming,
input.required_capabilities.as_ref(), input.required_capabilities.as_ref(),
Some(&input.auth_snapshot), Some(&input.auth_snapshot),
input.client_session_affinity.as_ref(), input.client_session_affinity.as_ref(),
current_unix_secs(), current_unix_secs(),
false,
) )
.await?; .await?;
@@ -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::spec_metadata::local_same_format_provider_spec_metadata;
use crate::ai_serving::planner::CandidateFailureDiagnostic; use crate::ai_serving::planner::CandidateFailureDiagnostic;
use crate::ai_serving::planner::{ 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, AiExecutionDecisionResponseParts,
}; };
use crate::ai_serving::transport::{ use crate::ai_serving::transport::{
@@ -184,7 +184,7 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
compatibility_edits: _, compatibility_edits: _,
request_redacted: _, request_redacted: _,
} = resolved; } = 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 { let mut decision = build_ai_execution_decision_response(AiExecutionDecisionResponseParts {
decision_is_stream: spec_metadata.require_streaming, 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(), request_id: trace_id.to_string(),
candidate_id: candidate_id.to_string(), candidate_id: candidate_id.to_string(),
provider_name: transport.provider.name.clone(), provider_name: transport.provider.name.clone(),
provider_type: transport.provider.provider_type.clone(),
provider_id: candidate.provider_id.clone(), provider_id: candidate.provider_id.clone(),
endpoint_id: candidate.endpoint_id.clone(), endpoint_id: candidate.endpoint_id.clone(),
key_id: candidate.key_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: Some(provider_request_body),
provider_request_body_base64: None, provider_request_body_base64: None,
content_type: Some("application/json".to_string()), content_type: Some("application/json".to_string()),
content_encoding: None, content_encoding: request_encoding.content_encoding,
request_gzip, request_gzip: request_encoding.request_gzip,
proxy, proxy,
transport_profile, transport_profile,
timeouts: resolve_transport_execution_timeouts(&transport), 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), report_context: Some(report_context),
auth_context: input.auth_context.clone(), 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)) Ok(Some(decision))
} }
@@ -40,7 +40,9 @@ use super::{
LocalSameFormatProviderCandidateAttempt, LocalSameFormatProviderDecisionInput, LocalSameFormatProviderCandidateAttempt, LocalSameFormatProviderDecisionInput,
LocalSameFormatProviderSpec, 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( pub(crate) fn resolve_same_format_provider_transport_unsupported_reason_for_trace(
transport: &GatewayProviderTransportSnapshot, transport: &GatewayProviderTransportSnapshot,
@@ -51,6 +53,7 @@ pub(crate) fn resolve_same_format_provider_transport_unsupported_reason_for_trac
"openai:chat" => "openai:chat", "openai:chat" => "openai:chat",
"openai:responses" => "openai:responses", "openai:responses" => "openai:responses",
"openai:responses:compact" => "openai:responses:compact", "openai:responses:compact" => "openai:responses:compact",
"openai:search" => "openai:search",
"openai:embedding" => "openai:embedding", "openai:embedding" => "openai:embedding",
"openai:rerank" => "openai:rerank", "openai:rerank" => "openai:rerank",
"claude:messages" => "claude:messages", "claude:messages" => "claude:messages",
@@ -137,13 +140,26 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
else { else {
return Ok(None); return Ok(None);
}; };
let enable_model_directives = let model_directive_resolution = input
crate::system_features::reasoning_model_directive_enabled_for_api_format_and_model( .model_directive_policy
state, .resolve_reasoning(spec.api_format, Some(&input.requested_model));
spec.api_format, let model_directive_mapping =
Some(&input.requested_model), match model_directive_resolution.mapping_patch_for_mapped_model(&prepared.mapped_model) {
) Ok(mapping) => mapping,
.await; 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 effective_headers = input.effective_headers(&parts.headers);
let redaction = resolve_provider_chat_pii_redaction( let redaction = resolve_provider_chat_pii_redaction(
state, state,
@@ -169,7 +185,7 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
prepared.force_body_stream_field, prepared.force_body_stream_field,
prepared.kiro_auth.as_ref(), prepared.kiro_auth.as_ref(),
prepared.is_claude_code, prepared.is_claude_code,
enable_model_directives, false,
) )
else { else {
mark_skipped_local_same_format_provider_candidate_with_extra_data( 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 base_provider_request_body = base_provider_request.body;
let mut compatibility_edits = base_provider_request.compatibility_edits; let mut compatibility_edits = base_provider_request.compatibility_edits;
if let Some(mapping) = if let Some(mapping) = model_directive_mapping.as_ref() {
crate::system_features::reasoning_model_directive_mapping_for_api_format_and_model(
state,
spec.api_format,
Some(&input.requested_model),
)
.await
{
let before_mapping = base_provider_request_body.clone(); let before_mapping = base_provider_request_body.clone();
crate::ai_serving::apply_model_directive_mapping_patch( crate::ai_serving::apply_model_directive_mapping_patch(
&mut base_provider_request_body, &mut base_provider_request_body,
&mapping, mapping,
); );
if before_mapping != base_provider_request_body { if before_mapping != base_provider_request_body {
compatibility_edits.push(SameFormatProviderCompatibilityEdit { 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 antigravity_auth = if prepared.is_antigravity {
let mut antigravity_support = classify_local_antigravity_request_support( let mut antigravity_support = classify_local_antigravity_request_support(
&transport, &transport,
@@ -467,6 +518,27 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
.await; .await;
return Ok(None); 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( request_identity_response_encoding_when_redacted(
&mut provider_request_headers, &mut provider_request_headers,
redaction.redacted, redaction.redacted,
@@ -1,23 +1,50 @@
use aether_ai_serving::AiRequestGzipPolicy; use aether_ai_serving::AiRequestGzipPolicy;
use serde_json::Value; 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; use super::state::GatewayProviderTransportSnapshot;
const DEFAULT_CODEX_REQUEST_GZIP_MIN_BYTES: usize = 64 * 1024; #[derive(Debug, Clone, Default, PartialEq, Eq)]
pub(crate) struct TransportRequestEncodingPolicy {
pub(crate) fn resolve_transport_request_gzip_policy( pub content_encoding: Option<String>,
transport: &GatewayProviderTransportSnapshot, pub request_gzip: Option<AiRequestGzipPolicy>,
) -> Option<AiRequestGzipPolicy> {
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))
} }
fn default_transport_request_gzip_policy( pub(crate) fn resolve_transport_request_encoding_policy(
transport: &GatewayProviderTransportSnapshot, transport: &GatewayProviderTransportSnapshot,
) -> Option<AiRequestGzipPolicy> { ) -> 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<String> {
if !transport if !transport
.provider .provider
.provider_type .provider_type
@@ -26,19 +53,24 @@ fn default_transport_request_gzip_policy(
{ {
return None; 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; return None;
} }
Some(AiRequestGzipPolicy { Some("zstd".to_string())
enabled: Some(true),
min_bytes: Some(DEFAULT_CODEX_REQUEST_GZIP_MIN_BYTES),
})
} }
fn is_codex_request_gzip_endpoint_api_format(api_format: &str) -> bool { fn is_codex_request_compression_api_format(api_format: &str) -> bool {
is_openai_responses_family_format(api_format) normalize_api_format_alias(api_format) == "openai:responses"
|| api_format.trim().eq_ignore_ascii_case("openai:image")
} }
fn transport_request_gzip_policy_from_config( fn transport_request_gzip_policy_from_config(
@@ -216,6 +248,16 @@ mod tests {
} }
} }
fn resolved_gzip_policy(
transport: &GatewayProviderTransportSnapshot,
) -> Option<AiRequestGzipPolicy> {
resolve_transport_request_encoding_policy(transport).request_gzip
}
fn resolved_content_encoding(transport: &GatewayProviderTransportSnapshot) -> Option<String> {
resolve_transport_request_encoding_policy(transport).content_encoding
}
#[test] #[test]
fn endpoint_request_gzip_policy_overrides_provider_policy() { fn endpoint_request_gzip_policy_overrides_provider_policy() {
let transport = sample_transport( let transport = sample_transport(
@@ -226,7 +268,7 @@ mod tests {
); );
assert_eq!( assert_eq!(
resolve_transport_request_gzip_policy(&transport), resolved_gzip_policy(&transport),
Some(AiRequestGzipPolicy { Some(AiRequestGzipPolicy {
enabled: Some(true), enabled: Some(true),
min_bytes: Some(1024), min_bytes: Some(1024),
@@ -244,7 +286,7 @@ mod tests {
); );
assert_eq!( assert_eq!(
resolve_transport_request_gzip_policy(&transport), resolved_gzip_policy(&transport),
Some(AiRequestGzipPolicy { Some(AiRequestGzipPolicy {
enabled: Some(false), enabled: Some(false),
min_bytes: None, min_bytes: None,
@@ -265,7 +307,7 @@ mod tests {
); );
assert_eq!( assert_eq!(
resolve_transport_request_gzip_policy(&transport), resolved_gzip_policy(&transport),
Some(AiRequestGzipPolicy { Some(AiRequestGzipPolicy {
enabled: Some(true), enabled: Some(true),
min_bytes: Some(4096), min_bytes: Some(4096),
@@ -283,7 +325,7 @@ mod tests {
); );
assert_eq!( assert_eq!(
resolve_transport_request_gzip_policy(&transport), resolved_gzip_policy(&transport),
Some(AiRequestGzipPolicy { Some(AiRequestGzipPolicy {
enabled: Some(true), enabled: Some(true),
min_bytes: Some(1), min_bytes: Some(1),
@@ -292,35 +334,73 @@ mod tests {
} }
#[test] #[test]
fn codex_responses_endpoint_gets_default_request_gzip_policy() { fn codex_responses_endpoint_uses_zstd_without_a_size_threshold() {
let transport = sample_transport("codex", "openai:responses", None, None); let mut transport = sample_transport("codex", "openai:responses", None, None);
transport.key.auth_type = "oauth".to_string();
assert_eq!( assert_eq!(
resolve_transport_request_gzip_policy(&transport), resolved_content_encoding(&transport).as_deref(),
Some(AiRequestGzipPolicy { Some("zstd")
enabled: Some(true),
min_bytes: Some(DEFAULT_CODEX_REQUEST_GZIP_MIN_BYTES),
})
); );
assert_eq!(resolved_gzip_policy(&transport), None);
} }
#[test] #[test]
fn codex_image_endpoint_gets_default_request_gzip_policy() { fn codex_responses_api_key_auth_does_not_enable_default_compression() {
let transport = sample_transport("codex", "openai:image", None, None); 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!( assert_eq!(
resolve_transport_request_gzip_policy(&transport), resolved_content_encoding(&transport).as_deref(),
Some(AiRequestGzipPolicy { Some("zstd")
enabled: Some(true),
min_bytes: Some(DEFAULT_CODEX_REQUEST_GZIP_MIN_BYTES),
})
); );
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] #[test]
fn non_codex_endpoint_does_not_get_default_request_gzip_policy() { fn non_codex_endpoint_does_not_get_default_request_gzip_policy() {
let transport = sample_transport("openai", "openai:responses", None, None); 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);
} }
} }
@@ -77,6 +77,7 @@ mod tests {
auth_endpoint_signature: None, auth_endpoint_signature: None,
execution_runtime_candidate: true, execution_runtime_candidate: true,
local_auth_rejection: None, local_auth_rejection: None,
model_directive_policy: Default::default(),
} }
} }
@@ -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::spec_metadata::local_gemini_files_spec_metadata;
use crate::ai_serving::planner::{ 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, AiExecutionDecisionResponseParts,
}; };
use crate::ai_serving::transport::{ use crate::ai_serving::transport::{
@@ -124,7 +124,7 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
upstream_url, upstream_url,
file_name: _, file_name: _,
} = resolved; } = 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 { let mut decision = build_ai_execution_decision_response(AiExecutionDecisionResponseParts {
decision_is_stream: spec_metadata.require_streaming, 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(), request_id: trace_id.to_string(),
candidate_id: candidate_id.clone(), candidate_id: candidate_id.clone(),
provider_name: transport.provider.name.clone(), provider_name: transport.provider.name.clone(),
provider_type: transport.provider.provider_type.clone(),
provider_id: candidate.provider_id.clone(), provider_id: candidate.provider_id.clone(),
endpoint_id: candidate.endpoint_id.clone(), endpoint_id: candidate.endpoint_id.clone(),
key_id: candidate.key_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) .map(str::trim)
.filter(|value| !value.is_empty()) .filter(|value| !value.is_empty())
.map(ToOwned::to_owned), .map(ToOwned::to_owned),
content_encoding: None, content_encoding: request_encoding.content_encoding,
request_gzip, request_gzip: request_encoding.request_gzip,
proxy, proxy,
transport_profile, transport_profile,
timeouts: resolve_transport_execution_timeouts(&transport), 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), report_context: Some(report_context),
auth_context: input.auth_context.clone(), 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)) Ok(Some(decision))
} }
@@ -53,7 +53,9 @@ pub(super) async fn resolve_local_gemini_files_decision_input(
state, state,
auth_context, auth_context,
None, None,
decision.auth_endpoint_signature.as_deref(),
Some(&explicit_required_capabilities), Some(&explicit_required_capabilities),
&decision.model_directive_policy,
) )
.await .await
{ {
@@ -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::spec_metadata::local_openai_image_spec_metadata;
use crate::ai_serving::planner::{ 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, AiExecutionDecisionResponseParts,
}; };
use crate::ai_serving::transport::{ 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), serde_json::Value::Bool(true),
); );
} }
let upstream_is_stream = resolved let upstream_is_stream = resolved.upstream_is_stream;
.provider_request_body
.get("stream")
.and_then(serde_json::Value::as_bool)
.unwrap_or(spec_metadata.require_streaming);
let effective_headers = input.effective_headers(&parts.headers); let effective_headers = input.effective_headers(&parts.headers);
let report_context = append_execution_contract_fields_to_value( let report_context = append_execution_contract_fields_to_value(
build_local_execution_report_context(LocalExecutionReportContextParts { 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, spec_metadata.api_format,
provider_api_format.as_str(), 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 { let mut decision = build_ai_execution_decision_response(AiExecutionDecisionResponseParts {
decision_is_stream: spec_metadata.require_streaming, 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(), request_id: trace_id.to_string(),
candidate_id: candidate_id.clone(), candidate_id: candidate_id.clone(),
provider_name: transport.provider.name.clone(), provider_name: transport.provider.name.clone(),
provider_type: transport.provider.provider_type.clone(),
provider_id: candidate.provider_id.clone(), provider_id: candidate.provider_id.clone(),
endpoint_id: candidate.endpoint_id.clone(), endpoint_id: candidate.endpoint_id.clone(),
key_id: candidate.key_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: Some(resolved.provider_request_body),
provider_request_body_base64: None, provider_request_body_base64: None,
content_type: Some("application/json".to_string()), content_type: Some("application/json".to_string()),
content_encoding: None, content_encoding: request_encoding.content_encoding,
request_gzip, request_gzip: request_encoding.request_gzip,
proxy, proxy,
transport_profile, transport_profile,
timeouts: resolve_transport_execution_timeouts(&transport), 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), report_context: Some(report_context),
auth_context: input.auth_context.clone(), 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)) Ok(Some(decision))
} }
@@ -16,8 +16,8 @@ use crate::ai_serving::transport::{
ProviderOpenAiImageHeadersInput, StandardProviderRequestHeadersInput, GROK_CHAT_PATH, ProviderOpenAiImageHeadersInput, StandardProviderRequestHeadersInput, GROK_CHAT_PATH,
}; };
use crate::ai_serving::{ use crate::ai_serving::{
apply_codex_openai_responses_special_body_edits, apply_codex_openai_responses_special_headers, apply_codex_openai_special_headers, build_chatgpt_web_image_request_body,
build_chatgpt_web_image_request_body, build_codex_openai_image_api_provider_request_body,
build_gemini_image_request_body_from_openai_image_request, build_gemini_image_request_body_from_openai_image_request,
build_openai_image_api_provider_request_body, build_openai_image_provider_request_body, build_openai_image_api_provider_request_body, build_openai_image_provider_request_body,
default_model_for_openai_image_operation, normalize_openai_image_request, 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) upstream_url: String,
pub(super) input_summary: Value, pub(super) input_summary: Value,
pub(super) transport_profile: Option<ResolvedTransportProfile>, pub(super) transport_profile: Option<ResolvedTransportProfile>,
pub(super) upstream_is_stream: bool,
} }
pub(super) async fn resolve_local_openai_image_candidate_payload_parts( 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, parts,
body_json, body_json,
body_base64, 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 { let Some(normalized_request) = normalized_request else {
mark_skipped_local_openai_image_candidate_with_failure_diagnostic( 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 { } else {
build_openai_image_upstream_url(transport, Some(parts.uri.path()), parts.uri.query()) build_openai_image_upstream_url(transport, Some(parts.uri.path()), parts.uri.query())
}; };
let mut provider_request_body = if is_chatgpt_web { let upstream_is_stream =
match build_chatgpt_web_image_request_body(parts, body_json, body_base64) { crate::ai_serving::planner::common::resolve_upstream_is_stream_for_provider(
Ok(body) => body, transport.endpoint.config.as_ref(),
Err(err) => err.to_error_json(), transport.provider.provider_type.as_str(),
} spec_metadata.api_format,
} else if is_codex || is_grok { spec_metadata.require_streaming && candidate.supports_streaming,
build_openai_image_provider_request_body(&normalized_request) 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 { } else {
build_openai_image_api_provider_request_body( build_openai_image_api_provider_request_body(
&normalized_request, &normalized_request,
Some(prepared_candidate.mapped_model.as_str()), Some(prepared_candidate.mapped_model.as_str()),
upstream_is_stream,
) )
}; };
if !is_chatgpt_web { let Some(provider_request_body) = provider_request_body else {
apply_codex_openai_responses_special_body_edits( mark_skipped_local_openai_image_candidate_with_failure_diagnostic(
&mut provider_request_body, state,
transport.provider.provider_type.as_str(), input,
spec_metadata.api_format, trace_id,
transport.endpoint.body_rules.as_ref(), candidate,
Some(candidate.key_id.as_str()), 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 { let Some(mut provider_request_headers) = (if is_grok {
build_grok_browser_headers(GrokHeaderInput { build_grok_browser_headers(GrokHeaderInput {
transport, transport,
@@ -214,10 +245,12 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
headers: effective_headers, headers: effective_headers,
auth_header: &auth_header, auth_header: &auth_header,
auth_value: &auth_value, auth_value: &auth_value,
accept: if is_codex || is_chatgpt_web { accept: if is_codex {
"text/event-stream" None
} else if upstream_is_stream {
Some("text/event-stream")
} else { } else {
"application/json" Some("application/json")
}, },
header_rules: transport.endpoint.header_rules.as_ref(), header_rules: transport.endpoint.header_rules.as_ref(),
provider_request_body: &provider_request_body, 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()); provider_request_headers.insert("x-aether-chatgpt-web-image".to_string(), "1".to_string());
} else if is_grok { } else if is_grok {
} else { } else {
apply_codex_openai_responses_special_headers( apply_codex_openai_special_headers(
&mut provider_request_headers, &mut provider_request_headers,
&provider_request_body, &provider_request_body,
effective_headers, effective_headers,
@@ -287,6 +320,7 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
upstream_url, upstream_url,
input_summary, input_summary,
transport_profile, transport_profile,
upstream_is_stream,
}) })
} }
@@ -403,7 +437,14 @@ async fn resolve_local_openai_image_to_gemini_candidate_payload_parts(
return None; 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( let Some(upstream_url) = crate::ai_serving::planner::standard::build_standard_upstream_url(
parts, parts,
transport, transport,
@@ -474,6 +515,7 @@ async fn resolve_local_openai_image_to_gemini_candidate_payload_parts(
upstream_url, upstream_url,
input_summary: converted.summary_json, input_summary: converted.summary_json,
transport_profile: None, transport_profile: None,
upstream_is_stream,
}) })
} }
@@ -58,7 +58,9 @@ pub(super) async fn resolve_local_openai_image_decision_input(
state, state,
auth_context, auth_context,
Some(requested_model.as_str()), Some(requested_model.as_str()),
decision.auth_endpoint_signature.as_deref(),
None, None,
&decision.model_directive_policy,
) )
.await .await
{ {
@@ -124,6 +126,7 @@ pub(super) async fn list_local_openai_image_candidate_attempts(
matches_client_format.then_some(&input.auth_snapshot), matches_client_format.then_some(&input.auth_snapshot),
input.client_session_affinity.as_ref(), input.client_session_affinity.as_ref(),
current_unix_secs(), current_unix_secs(),
false,
) )
.await .await
{ {
@@ -144,8 +147,8 @@ pub(super) async fn list_local_openai_image_candidate_attempts(
auth_snapshot_allows_cross_format_candidate( auth_snapshot_allows_cross_format_candidate(
&input.auth_snapshot, &input.auth_snapshot,
&input.requested_model, &input.requested_model,
None,
candidate, 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), matches_client_format.then_some(&input.auth_snapshot),
input.client_session_affinity.as_ref(), input.client_session_affinity.as_ref(),
current_unix_secs(), current_unix_secs(),
false,
) )
.await .await
{ {
@@ -206,16 +210,16 @@ pub(super) async fn build_local_openai_image_candidate_attempt_source<'a>(
auth_snapshot_allows_cross_format_candidate( auth_snapshot_allows_cross_format_candidate(
&input.auth_snapshot, &input.auth_snapshot,
&input.requested_model, &input.requested_model,
None,
candidate, candidate,
false,
) )
}); });
format_skipped.retain(|candidate| { format_skipped.retain(|candidate| {
auth_snapshot_allows_cross_format_candidate( auth_snapshot_allows_cross_format_candidate(
&input.auth_snapshot, &input.auth_snapshot,
&input.requested_model, &input.requested_model,
None,
&candidate.candidate, &candidate.candidate,
false,
) )
}); });
} }
@@ -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::spec_metadata::local_video_create_spec_metadata;
use crate::ai_serving::planner::{ 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, AiExecutionDecisionResponseParts,
}; };
use crate::ai_serving::transport::{ 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, provider_request_body,
upstream_url, upstream_url,
} = resolved; } = 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 { let mut decision = build_ai_execution_decision_response(AiExecutionDecisionResponseParts {
decision_is_stream: false, 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(), request_id: trace_id.to_string(),
candidate_id: candidate_id.clone(), candidate_id: candidate_id.clone(),
provider_name: transport.provider.name.clone(), provider_name: transport.provider.name.clone(),
provider_type: transport.provider.provider_type.clone(),
provider_id: candidate.provider_id.clone(), provider_id: candidate.provider_id.clone(),
endpoint_id: candidate.endpoint_id.clone(), endpoint_id: candidate.endpoint_id.clone(),
key_id: candidate.key_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) .map(str::trim)
.filter(|value| !value.is_empty()) .filter(|value| !value.is_empty())
.map(ToOwned::to_owned), .map(ToOwned::to_owned),
content_encoding: None, content_encoding: request_encoding.content_encoding,
request_gzip, request_gzip: request_encoding.request_gzip,
proxy, proxy,
transport_profile, transport_profile,
timeouts: resolve_transport_execution_timeouts(&transport), 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), report_context: Some(report_context),
auth_context: input.auth_context.clone(), 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)) Ok(Some(decision))
} }
@@ -62,7 +62,9 @@ pub(super) async fn resolve_local_video_create_decision_input(
state, state,
auth_context, auth_context,
Some(requested_model.as_str()), Some(requested_model.as_str()),
decision.auth_endpoint_signature.as_deref(),
None, None,
&decision.model_directive_policy,
) )
.await .await
{ {
@@ -130,6 +132,7 @@ pub(super) async fn list_local_video_create_candidate_attempts(
Some(&input.auth_snapshot), Some(&input.auth_snapshot),
input.client_session_affinity.as_ref(), input.client_session_affinity.as_ref(),
current_unix_secs(), current_unix_secs(),
false,
) )
.await .await
{ {
@@ -186,6 +189,7 @@ pub(super) async fn build_local_video_create_candidate_attempt_source<'a>(
Some(&input.auth_snapshot), Some(&input.auth_snapshot),
input.client_session_affinity.as_ref(), input.client_session_affinity.as_ref(),
current_unix_secs(), current_unix_secs(),
false,
) )
.await .await
{ {
@@ -3,5 +3,42 @@
mod tests; mod tests;
pub(crate) use crate::ai_serving::{ 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<crate::ai_serving::CodexResponsesModelCapabilities> {
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<crate::ai_serving::CodexResponsesModelCapabilities> {
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,
),
)
}
@@ -1,16 +1,56 @@
use std::collections::BTreeMap; use std::collections::BTreeMap;
use super::{ 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 crate::ai_serving::planner::standard::build_local_openai_responses_request_body;
use http::{HeaderMap, HeaderValue}; use http::{HeaderMap, HeaderValue};
use serde_json::json; 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] #[test]
fn applies_codex_defaults_when_body_rules_do_not_handle_fields() { fn applies_codex_defaults_when_body_rules_do_not_handle_fields() {
let mut body = json!({ let mut body = json!({
"model": "gpt-5", "model": "gpt-5.4",
"max_output_tokens": 128, "max_output_tokens": 128,
"temperature": 0.3, "temperature": 0.3,
"top_p": 0.9, "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("top_p").is_none());
assert!(body.get("metadata").is_none()); assert!(body.get("metadata").is_none());
assert_eq!(body["store"], false); 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["include"], json!(["reasoning.encrypted_content"]));
assert_eq!(body["parallel_tool_calls"], true); 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] #[test]
@@ -80,7 +121,7 @@ fn strips_store_for_compact_even_when_body_rules_handle_it() {
{"action":"set","path":"top_p","value":0.5} {"action":"set","path":"top_p","value":0.5}
]); ]);
let mut body = json!({ let mut body = json!({
"model": "gpt-5", "model": "gpt-5.4",
"max_output_tokens": 128, "max_output_tokens": 128,
"metadata": {"client": "desktop", "mode": "custom"}, "metadata": {"client": "desktop", "mode": "custom"},
"store": true, "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("max_output_tokens").is_none());
assert!(body.get("store").is_none()); assert!(body.get("store").is_none());
assert_eq!(body["instructions"], "Keep custom"); assert_eq!(body["instructions"], "Keep custom");
assert_eq!(body["metadata"]["mode"], "custom"); assert!(body.get("metadata").is_none());
assert_eq!(body["top_p"], 0.5); 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] #[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!({ let mut body = json!({
"model": "gpt-5", "model": "gpt-5",
"input": "hello", "input": "hello",
@@ -118,10 +176,7 @@ fn injects_stable_prompt_cache_key_for_codex_requests() {
Some("key-123"), Some("key-123"),
); );
assert_eq!( assert!(body.get("prompt_cache_key").is_none());
body["prompt_cache_key"],
"53363264-dbb0-5f9d-b9c7-3e92c45c5bdf"
);
} }
#[test] #[test]
@@ -144,57 +199,96 @@ fn keeps_existing_prompt_cache_key_for_codex_requests() {
} }
#[test] #[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 mut headers = BTreeMap::new();
let body = json!({ let body = json!({
"model": "gpt-5", "model": "gpt-5",
"prompt_cache_key": "172c39e6-c0a0-5a70-8b63-e0f8e0d185a3", "prompt_cache_key": "172c39e6-c0a0-5a70-8b63-e0f8e0d185a3",
}); });
apply_codex_openai_responses_special_headers( apply_codex_openai_special_headers(
&mut headers, &mut headers,
&body, &body,
&HeaderMap::new(), &HeaderMap::new(),
"codex", "codex",
"openai:responses", "openai:responses",
Some("trace-codex-123"), 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!( assert_eq!(
headers.get("chatgpt-account-id"), headers.get("chatgpt-account-id"),
Some(&"acc-123".to_string()) Some(&"acc-123".to_string())
); );
assert_eq!(
headers.get("x-client-request-id"),
Some(&"trace-codex-123".to_string())
);
assert_eq!( assert_eq!(
headers.get("user-agent"), headers.get("user-agent"),
Some( Some(&"codex_cli_rs/0.144.1".to_string())
&"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())
); );
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] #[test]
fn respects_existing_codex_request_and_session_headers() { fn preserves_client_context_headers_and_enforces_codex_provider_identity() {
let mut headers = BTreeMap::new(); let mut headers = BTreeMap::new();
headers.insert( headers.insert(
"x-client-request-id".to_string(), "x-client-request-id".to_string(),
"kept-by-rule-request".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!({ let body = json!({
"model": "gpt-5", "model": "gpt-5",
"prompt_cache_key": "172c39e6-c0a0-5a70-8b63-e0f8e0d185a3", "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"), HeaderValue::from_static("user-specified-request"),
); );
original_headers.insert( original_headers.insert(
"session_id", "session-id",
HeaderValue::from_static("user-specified-session"), HeaderValue::from_static("user-specified-session"),
); );
original_headers.insert( original_headers.insert(
"conversation_id", "thread-id",
HeaderValue::from_static("user-specified-conversation"), HeaderValue::from_static("user-specified-thread"),
); );
original_headers.insert( original_headers.insert(
"user-agent", "user-agent",
@@ -220,64 +314,92 @@ fn respects_existing_codex_request_and_session_headers() {
"originator", "originator",
HeaderValue::from_static("user-specified-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, &mut headers,
&body, &body,
&original_headers, &original_headers,
"codex", "codex",
"openai:responses", "openai:responses",
Some("trace-codex-123"), Some("trace-codex-123"),
Some(r#"{"account_id":"acc-123"}"#), Some(r#"{"account_id":"acc-123","is_fedramp":true}"#),
); );
assert_eq!( assert_eq!(
headers.get("x-client-request-id"), headers.get("x-client-request-id"),
Some(&"kept-by-rule-request".to_string()) Some(&"kept-by-rule-request".to_string())
); );
assert!(!headers.contains_key("user-agent")); assert_eq!(
assert!(!headers.contains_key("originator")); headers.get("user-agent"),
assert_eq!(headers.get("session_id"), Some(&"kept-by-rule".to_string())); Some(&"codex_cli_rs/0.144.1".to_string())
assert!(!headers.contains_key("conversation_id")); );
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] #[test]
fn skips_conversation_id_for_compact_codex_requests() { fn compact_does_not_derive_session_headers_from_body() {
let mut headers = BTreeMap::new(); let mut headers = BTreeMap::new();
let body = json!({ let body = json!({
"model": "gpt-5", "model": "gpt-5",
"prompt_cache_key": "172c39e6-c0a0-5a70-8b63-e0f8e0d185a3", "prompt_cache_key": "172c39e6-c0a0-5a70-8b63-e0f8e0d185a3",
}); });
apply_codex_openai_responses_special_headers( apply_codex_openai_special_headers(
&mut headers, &mut headers,
&body, &body,
&HeaderMap::new(), &HeaderMap::new(),
"codex", "codex",
"openai:responses:compact", "openai:responses:compact",
Some("trace-codex-compact-123"), Some("trace-codex-compact-123"),
Some(r#"{"account_id":"acc-123"}"#), Some(r#"{"account_id":"acc-123","is_fedramp":true}"#),
); );
assert_eq!( assert_eq!(
headers.get("chatgpt-account-id"), headers.get("chatgpt-account-id"),
Some(&"acc-123".to_string()) Some(&"acc-123".to_string())
); );
assert_eq!( assert_eq!(headers.get("x-client-request-id"), None);
headers.get("x-client-request-id"),
Some(&"trace-codex-compact-123".to_string())
);
assert_eq!( assert_eq!(
headers.get("user-agent"), headers.get("user-agent"),
Some( Some(&"codex_cli_rs/0.144.1".to_string())
&"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("originator"), Some(&"codex_cli_rs".to_string()));
assert_eq!( assert!(!headers.contains_key("version"));
headers.get("session_id"), assert_eq!(headers.get("x-openai-fedramp"), Some(&"true".to_string()));
Some(&"ab5ecce4f0d110fe".to_string()) assert_eq!(headers.get("session-id"), None);
); assert_eq!(headers.get("thread-id"), None);
assert!(!headers.contains_key("conversation_id"));
} }
@@ -60,7 +60,9 @@ pub(super) async fn resolve_local_standard_decision_input(
state, state,
auth_context, auth_context,
Some(requested_model.as_str()), Some(requested_model.as_str()),
decision.auth_endpoint_signature.as_deref(),
None, None,
&decision.model_directive_policy,
) )
.await .await
{ {
@@ -119,6 +121,7 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
); );
let preselection = preselect_local_execution_candidates_with_serving( let preselection = preselect_local_execution_candidates_with_serving(
planner_state, planner_state,
&input.model_directive_policy,
spec_metadata.api_format, spec_metadata.api_format,
&input.requested_model, &input.requested_model,
false, false,
@@ -243,6 +246,7 @@ pub(super) async fn build_local_standard_candidate_attempt_source<'a>(
let (source, candidate_count) = let (source, candidate_count) =
build_lazy_requested_model_execution_candidate_attempt_source_with_serving( build_lazy_requested_model_execution_candidate_attempt_source_with_serving(
planner_state, planner_state,
&input.model_directive_policy,
trace_id, trace_id,
spec_metadata.api_format, spec_metadata.api_format,
&input.requested_model, &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( let image_preselection = preselect_local_execution_candidates_for_api_formats_with_serving(
planner_state, planner_state,
&input.model_directive_policy,
spec_metadata.api_format, spec_metadata.api_format,
&input.requested_model, &input.requested_model,
spec_metadata.require_streaming, spec_metadata.require_streaming,
@@ -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::spec_metadata::local_standard_spec_metadata;
use crate::ai_serving::planner::CandidateFailureDiagnostic; use crate::ai_serving::planner::CandidateFailureDiagnostic;
use crate::ai_serving::planner::{ 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, AiExecutionDecisionResponseParts,
}; };
use crate::ai_serving::transport::{ use crate::ai_serving::transport::{
@@ -176,7 +176,7 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
transport_profile: _, transport_profile: _,
request_redacted: _, request_redacted: _,
} = resolved; } = 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 { let mut decision = build_ai_execution_decision_response(AiExecutionDecisionResponseParts {
decision_is_stream: spec_metadata.require_streaming, 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(), request_id: trace_id.to_string(),
candidate_id: candidate_id.to_string(), candidate_id: candidate_id.to_string(),
provider_name: candidate.provider_name.clone(), provider_name: candidate.provider_name.clone(),
provider_type: transport.provider.provider_type.clone(),
provider_id: candidate.provider_id.clone(), provider_id: candidate.provider_id.clone(),
endpoint_id: candidate.endpoint_id.clone(), endpoint_id: candidate.endpoint_id.clone(),
key_id: candidate.key_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: Some(provider_request_body),
provider_request_body_base64: None, provider_request_body_base64: None,
content_type: Some("application/json".to_string()), content_type: Some("application/json".to_string()),
content_encoding: None, content_encoding: request_encoding.content_encoding,
request_gzip, request_gzip: request_encoding.request_gzip,
proxy, proxy,
transport_profile, transport_profile,
timeouts, timeouts,
@@ -213,7 +214,11 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
report_context: Some(report_context), report_context: Some(report_context),
auth_context: input.auth_context.clone(), 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)) Ok(Some(decision))
} }
@@ -372,6 +377,7 @@ mod tests {
routing_policy: None, routing_policy: None,
routing_trace_seed: None, routing_trace_seed: None,
routing_context: None, routing_context: None,
model_directive_policy: Default::default(),
} }
} }
@@ -475,6 +481,7 @@ mod tests {
} else { } else {
"gpt-4o-upstream".to_string() "gpt-4o-upstream".to_string()
}, },
supports_streaming: true,
mapping_matched_model: None, mapping_matched_model: None,
} }
} }
@@ -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::spec_metadata::local_standard_spec_metadata;
use crate::ai_serving::planner::standard::{ use crate::ai_serving::planner::standard::{
apply_codex_openai_responses_special_headers, apply_deepseek_tool_call_thinking_compat, apply_codex_openai_special_headers, apply_deepseek_tool_call_thinking_compat,
is_deepseek_provider, request_body_build_failure_extra_data, codex_model_capabilities_for_transport, is_deepseek_provider,
request_conversion_failure_extra_data, request_body_build_failure_extra_data, request_conversion_failure_extra_data,
}; };
use crate::ai_serving::transport::kiro::{ use crate::ai_serving::transport::kiro::{
build_kiro_provider_headers, build_kiro_provider_request_body, build_kiro_provider_headers, build_kiro_provider_request_body,
@@ -44,7 +44,9 @@ use crate::ai_serving::transport::{
}; };
use crate::ai_serving::{ use crate::ai_serving::{
build_openai_image_request_body_from_gemini_image_request, gemini_request_is_image_generation, 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, CandidateFailureDiagnostic, GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth,
OpenAiImageOperation,
}; };
use crate::{AppState, GatewayError}; use crate::{AppState, GatewayError};
@@ -313,7 +315,13 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
{ {
return Ok( return Ok(
resolve_local_gemini_image_to_openai_image_candidate_payload_parts( 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, .await,
); );
@@ -555,13 +563,27 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
); );
let force_body_stream_field = let force_body_stream_field =
endpoint_config_forces_body_stream_field(transport.endpoint.config.as_ref()); endpoint_config_forces_body_stream_field(transport.endpoint.config.as_ref());
let enable_model_directives = let model_directive_resolution = input
crate::system_features::reasoning_model_directive_enabled_for_api_format_and_model( .model_directive_policy
state, .resolve_reasoning(provider_api_format, Some(&input.requested_model));
provider_api_format, let model_directive_mapping = match model_directive_resolution
Some(&input.requested_model), .mapping_patch_for_mapped_model(&prepared_candidate.mapped_model)
) {
.await; 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( let redaction = resolve_provider_chat_pii_redaction(
state, state,
parts, 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(input.auth_context.api_key_id.as_str()),
Some(effective_headers), Some(effective_headers),
enable_model_directives, false,
) { ) {
Some(body) => body, Some(body) => body,
None => { None => {
@@ -655,18 +677,8 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
provider_api_format, provider_api_format,
Some(body_json), Some(body_json),
); );
if let Some(mapping) = if let Some(mapping) = model_directive_mapping.as_ref() {
crate::system_features::reasoning_model_directive_mapping_for_api_format_and_model( crate::ai_serving::apply_model_directive_mapping_patch(&mut provider_request_body, mapping);
state,
provider_api_format,
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`; // Directive mapping is a deep-merge patch and may overwrite/add `stream`;
// re-enforce stream-field policy afterward. // re-enforce stream-field policy afterward.
enforce_provider_body_stream_policy( 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() { if let Some(kiro_auth) = kiro_auth.as_ref() {
return Ok(build_kiro_cross_format_payload_parts( return Ok(build_kiro_cross_format_payload_parts(
state, state,
@@ -752,8 +822,6 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
.await); .await);
} }
let normalized_provider_api_format =
crate::ai_serving::normalize_api_format_alias(provider_api_format);
if normalized_provider_api_format == "gemini:generate_content" if normalized_provider_api_format == "gemini:generate_content"
&& is_gemini_cli_provider_transport(transport) && is_gemini_cli_provider_transport(transport)
{ {
@@ -838,7 +906,7 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
return Ok(None); return Ok(None);
}; };
let mut provider_request_headers = resolved_headers.headers; let mut provider_request_headers = resolved_headers.headers;
apply_codex_openai_responses_special_headers( apply_codex_openai_special_headers(
&mut provider_request_headers, &mut provider_request_headers,
&provider_request_body, &provider_request_body,
effective_headers, effective_headers,
@@ -988,7 +1056,7 @@ async fn build_gemini_cli_cross_format_payload_parts(
}; };
let mut provider_request_headers = resolved.headers.headers; let mut provider_request_headers = resolved.headers.headers;
apply_codex_openai_responses_special_headers( apply_codex_openai_special_headers(
&mut provider_request_headers, &mut provider_request_headers,
&resolved.body, &resolved.body,
effective_headers, effective_headers,
@@ -1146,6 +1214,7 @@ async fn resolve_local_gemini_image_to_openai_image_candidate_payload_parts(
body_json: &serde_json::Value, body_json: &serde_json::Value,
input: &LocalStandardDecisionInput, input: &LocalStandardDecisionInput,
attempt: &LocalStandardCandidateAttempt, attempt: &LocalStandardCandidateAttempt,
client_requires_streaming: bool,
) -> Option<LocalStandardCandidatePayloadParts> { ) -> Option<LocalStandardCandidatePayloadParts> {
let client_api_format = "gemini:generate_content"; let client_api_format = "gemini:generate_content";
let provider_api_format = "openai:image"; let provider_api_format = "openai:image";
@@ -1221,9 +1290,44 @@ async fn resolve_local_gemini_image_to_openai_image_candidate_payload_parts(
return None; return None;
}; };
let upstream_is_stream = true; let upstream_is_stream = resolve_upstream_is_stream_for_provider(
let upstream_url = transport.endpoint.config.as_ref(),
build_openai_image_upstream_url(transport, Some("/v1/images/generations"), None); 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 effective_headers = input.effective_headers(&parts.headers);
let Some(mut provider_request_headers) = let Some(mut provider_request_headers) =
build_openai_image_headers(ProviderOpenAiImageHeadersInput { build_openai_image_headers(ProviderOpenAiImageHeadersInput {
@@ -1231,9 +1335,15 @@ async fn resolve_local_gemini_image_to_openai_image_candidate_payload_parts(
headers: effective_headers, headers: effective_headers,
auth_header: &prepared_candidate.auth_header, auth_header: &prepared_candidate.auth_header,
auth_value: &prepared_candidate.auth_value, 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(), 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, original_request_body: body_json,
}) })
else { else {
@@ -1254,9 +1364,9 @@ async fn resolve_local_gemini_image_to_openai_image_candidate_payload_parts(
.await; .await;
return None; return None;
}; };
apply_codex_openai_responses_special_headers( apply_codex_openai_special_headers(
&mut provider_request_headers, &mut provider_request_headers,
&converted.body_json, &provider_request_body,
effective_headers, effective_headers,
transport.provider.provider_type.as_str(), transport.provider.provider_type.as_str(),
provider_api_format, 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, auth_value: prepared_candidate.auth_value,
mapped_model: converted.mapped_model, mapped_model: converted.mapped_model,
provider_api_format: provider_api_format.to_string(), provider_api_format: provider_api_format.to_string(),
provider_request_body: converted.body_json, provider_request_body,
provider_request_headers, provider_request_headers,
upstream_url, upstream_url,
upstream_is_stream, upstream_is_stream,
@@ -15,7 +15,8 @@ mod normalize;
mod openai; mod openai;
pub(crate) use self::codex::{ 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::deepseek::{apply_deepseek_tool_call_thinking_compat, is_deepseek_provider};
pub(crate) use self::family::{ pub(crate) use self::family::{
@@ -25,9 +26,11 @@ pub(crate) use self::family::{
pub(crate) use self::normalize::{ pub(crate) use self::normalize::{
build_cross_format_openai_chat_request_body, build_cross_format_openai_chat_upstream_url, 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,
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_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_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::{ pub(crate) use self::openai::{
build_local_openai_chat_stream_attempt_source_for_kind, build_local_openai_chat_stream_attempt_source_for_kind,
@@ -297,7 +300,7 @@ mod tests {
let converted = build_standard_request_body( let converted = build_standard_request_body(
&request, &request,
"claude:messages", "claude:messages",
"gpt-5", "gpt-5.4",
"codex", "codex",
"openai:responses", "openai:responses",
"/v1/messages", "/v1/messages",
@@ -309,7 +312,7 @@ mod tests {
assert!(converted.get("metadata").is_none()); assert!(converted.get("metadata").is_none());
assert_eq!(converted["store"], false); 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["include"], json!(["reasoning.encrypted_content"]));
assert_eq!(converted["parallel_tool_calls"], true); assert_eq!(converted["parallel_tool_calls"], true);
assert_eq!(converted["reasoning"]["effort"], "medium"); assert_eq!(converted["reasoning"]["effort"], "medium");
@@ -12,9 +12,38 @@ pub(crate) use self::chat::{
}; };
pub(crate) use self::responses::{ pub(crate) use self::responses::{
build_cross_format_openai_responses_request_body, 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_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, build_local_openai_responses_upstream_url,
}; };
pub(super) use crate::ai_serving::planner::common::{ pub(super) use crate::ai_serving::planner::common::{
enforce_provider_body_stream_policy, request_requires_body_stream_field, 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()
}
@@ -9,7 +9,10 @@ use crate::ai_serving::{
GatewayProviderTransportSnapshot, 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( pub(crate) fn build_local_openai_chat_request_body(
body_json: &Value, body_json: &Value,
@@ -39,6 +42,12 @@ pub(crate) fn build_local_openai_chat_request_body(
upstream_is_stream, upstream_is_stream,
request_requires_body_stream_field(body_json, force_body_stream_field), 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) Some(provider_request_body)
} }
@@ -92,6 +101,12 @@ pub(crate) fn build_cross_format_openai_chat_request_body(
upstream_is_stream, upstream_is_stream,
request_requires_body_stream_field(body_json, force_body_stream_field), 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) Some(provider_request_body)
} }
@@ -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::transport::apply_standard_provider_request_body_rules_with_request_headers;
use crate::ai_serving::{ use crate::ai_serving::{
apply_codex_openai_responses_special_body_edits,
apply_openai_responses_compact_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_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, build_local_openai_responses_request_body_with_model_directives as surface_build_local_openai_responses_request_body,
GatewayProviderTransportSnapshot, 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( pub(crate) fn build_local_openai_responses_request_body(
body_json: &Value, body_json: &Value,
@@ -19,9 +21,35 @@ pub(crate) fn build_local_openai_responses_request_body(
provider_type: &str, provider_type: &str,
provider_api_format: &str, provider_api_format: &str,
body_rules: Option<&Value>, body_rules: Option<&Value>,
user_api_key_id: Option<&str>, _user_api_key_id: Option<&str>,
request_headers: &http::HeaderMap, request_headers: &http::HeaderMap,
enable_model_directives: bool, enable_model_directives: bool,
) -> Option<Value> {
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<Value> { ) -> Option<Value> {
let provider_request_body = surface_build_local_openai_responses_request_body( let provider_request_body = surface_build_local_openai_responses_request_body(
body_json, body_json,
@@ -36,12 +64,18 @@ pub(crate) fn build_local_openai_responses_request_body(
body_json, body_json,
request_headers, 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, &mut provider_request_body,
provider_type, provider_type,
provider_api_format, provider_api_format,
mapped_model,
source_model,
model_capabilities,
body_rules, body_rules,
user_api_key_id,
); );
apply_openai_responses_compact_special_body_edits( apply_openai_responses_compact_special_body_edits(
&mut provider_request_body, &mut provider_request_body,
@@ -53,6 +87,12 @@ pub(crate) fn build_local_openai_responses_request_body(
require_streaming, require_streaming,
request_requires_body_stream_field(body_json, force_body_stream_field), 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) Some(provider_request_body)
} }
@@ -65,9 +105,37 @@ pub(crate) fn build_cross_format_openai_responses_request_body(
force_body_stream_field: bool, force_body_stream_field: bool,
provider_type: &str, provider_type: &str,
body_rules: Option<&Value>, body_rules: Option<&Value>,
user_api_key_id: Option<&str>, _user_api_key_id: Option<&str>,
request_headers: &http::HeaderMap, request_headers: &http::HeaderMap,
enable_model_directives: bool, enable_model_directives: bool,
) -> Option<Value> {
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<Value> { ) -> Option<Value> {
let provider_request_body = surface_build_cross_format_openai_responses_request_body( let provider_request_body = surface_build_cross_format_openai_responses_request_body(
body_json, body_json,
@@ -84,12 +152,18 @@ pub(crate) fn build_cross_format_openai_responses_request_body(
body_json, body_json,
request_headers, 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, &mut provider_request_body,
provider_type, provider_type,
provider_api_format, provider_api_format,
mapped_model,
source_model,
model_capabilities,
body_rules, body_rules,
user_api_key_id,
); );
apply_openai_responses_compact_special_body_edits( apply_openai_responses_compact_special_body_edits(
&mut provider_request_body, &mut provider_request_body,
@@ -101,6 +175,12 @@ pub(crate) fn build_cross_format_openai_responses_request_body(
upstream_is_stream, upstream_is_stream,
request_requires_body_stream_field(body_json, force_body_stream_field), 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) Some(provider_request_body)
} }
@@ -6,8 +6,8 @@ use http::Request;
use serde_json::{json, Value}; use serde_json::{json, Value};
use super::{ use super::{
build_cross_format_openai_responses_request_body, build_local_openai_responses_request_body, build_cross_format_openai_responses_request_body, build_local_openai_chat_request_body,
build_local_openai_responses_upstream_url, build_local_openai_responses_request_body, build_local_openai_responses_upstream_url,
}; };
fn object_keys(value: &Value) -> Vec<&str> { fn object_keys(value: &Value) -> Vec<&str> {
@@ -146,12 +146,10 @@ fn local_openai_responses_wrapper_preserves_body_order_after_edits() {
"reasoning", "reasoning",
"tool_choice", "tool_choice",
"parallel_tool_calls", "parallel_tool_calls",
"instructions",
"prompt_cache_key",
] ]
); );
assert_eq!(provider_request_body["parallel_tool_calls"], json!(true)); 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] #[test]
@@ -181,18 +179,49 @@ fn local_openai_responses_compact_wrapper_strips_store_for_same_format_requests(
} }
#[test] #[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!({ let body_json = json!({
"model": "gpt-5.4", "model": "gpt-5.6-sol",
"input": [], "input": [{
"type": "message",
"role": "user",
"content": [{"type": "input_text", "text": "hello"}]
}],
"instructions": "Work carefully",
"client_metadata": {"origin": "codex"},
"include": ["reasoning.encrypted_content"], "include": ["reasoning.encrypted_content"],
"store": true, "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, &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,
false, false,
"codex", "codex",
@@ -202,22 +231,44 @@ fn local_openai_responses_compact_wrapper_strips_include_for_codex_requests() {
&http::HeaderMap::new(), &http::HeaderMap::new(),
false, false,
) )
.expect("local codex compact body should build"); .expect("local Codex Compact body should build");
assert!(provider_request_body.get("include").is_none()); for field in [
assert!(provider_request_body.get("store").is_none()); "client_metadata",
assert!(provider_request_body.get("stream").is_none()); "include",
assert_eq!(provider_request_body["instructions"], ""); "store",
assert_eq!( "stream",
provider_request_body["prompt_cache_key"], "stream_options",
"3d2e2842-74cb-55dd-803a-b8940b3500c2" "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] #[test]
fn local_openai_responses_wrapper_applies_model_directive_before_body_rules() { fn local_openai_responses_wrapper_applies_model_directive_before_body_rules() {
let body_json = json!({ let body_json = json!({
"model": "gpt-5.4-max", "model": "gpt-5.6-sol-max",
"input": "hello", "input": "hello",
"reasoning": {"effort": "low", "summary": "auto"} "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( let provider_request_body = build_local_openai_responses_request_body(
&body_json, &body_json,
"gpt-5.4", "gpt-5.6-sol",
false, false,
false, false,
"openai", "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); 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] #[test]
fn local_openai_responses_upstream_url_preserves_codex_base_path() { fn local_openai_responses_upstream_url_preserves_codex_base_path() {
let request = Request::builder() let request = Request::builder()
@@ -371,7 +548,7 @@ fn applies_codex_defaults_unless_body_rules_handle_the_field() {
} }
#[test] #[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!({ let body_json = json!({
"model": "claude-sonnet-4-5", "model": "claude-sonnet-4-5",
"messages": [{ "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"); .expect("claude cli to codex request should build");
assert_eq!( assert!(provider_request_body.get("prompt_cache_key").is_none());
provider_request_body["prompt_cache_key"],
"4ee6ea6e-3ac6-5a18-8cb8-1f8b956419e5"
);
} }
#[test] #[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!({ let body_json = json!({
"model": "gpt-5", "model": "gpt-5",
"messages": [{ "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"); .expect("openai chat to codex request should build");
assert_eq!( assert!(provider_request_body.get("prompt_cache_key").is_none());
provider_request_body["prompt_cache_key"],
"4ee6ea6e-3ac6-5a18-8cb8-1f8b956419e5"
);
} }
@@ -6,7 +6,7 @@ use crate::ai_serving::planner::report_context::{
insert_provider_stream_event_api_format, LocalExecutionReportContextParts, insert_provider_stream_event_api_format, LocalExecutionReportContextParts,
}; };
use crate::ai_serving::planner::{ 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, AiExecutionDecisionResponseParts,
}; };
use crate::ai_serving::transport::{ use crate::ai_serving::transport::{
@@ -44,6 +44,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
candidate_id, candidate_id,
.. ..
} = attempt; } = attempt;
let upstream_is_stream = upstream_is_stream && eligible.candidate.supports_streaming;
let payload_started_at = std::time::Instant::now(); let payload_started_at = std::time::Instant::now();
let Some(resolved) = resolve_local_openai_chat_candidate_payload_parts( let Some(resolved) = resolve_local_openai_chat_candidate_payload_parts(
state, 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, payload_started_at.elapsed().as_millis() as u64,
); );
let candidate = &eligible.candidate; 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 let prompt_cache_key = resolved
.provider_request_body .provider_request_body
@@ -208,7 +217,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
"stream_candidate_report_context", "stream_candidate_report_context",
report_context_started_at.elapsed().as_millis() as u64, 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 decision_started_at = std::time::Instant::now();
let mut decision = build_ai_execution_decision_response(AiExecutionDecisionResponseParts { 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(), request_id: trace_id.to_string(),
candidate_id: candidate_id.clone(), candidate_id: candidate_id.clone(),
provider_name: transport.provider.name.clone(), provider_name: transport.provider.name.clone(),
provider_type: transport.provider.provider_type.clone(),
provider_id: candidate.provider_id.clone(), provider_id: candidate.provider_id.clone(),
endpoint_id: candidate.endpoint_id.clone(), endpoint_id: candidate.endpoint_id.clone(),
key_id: candidate.key_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: Some(provider_request_body),
provider_request_body_base64: None, provider_request_body_base64: None,
content_type: Some("application/json".to_string()), content_type: Some("application/json".to_string()),
content_encoding: None, content_encoding: request_encoding.content_encoding,
request_gzip, request_gzip: request_encoding.request_gzip,
proxy, proxy,
transport_profile, transport_profile,
timeouts, timeouts,
@@ -246,7 +256,11 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
report_context: Some(report_context), report_context: Some(report_context),
auth_context: input.auth_context.clone(), 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( observe_gateway_stage_ms(
"stream_candidate_decision_build", "stream_candidate_decision_build",
decision_started_at.elapsed().as_millis() as u64, decision_started_at.elapsed().as_millis() as u64,
@@ -25,11 +25,11 @@ use crate::ai_serving::planner::redaction::{
request_identity_response_encoding_when_redacted, resolve_provider_chat_pii_redaction, request_identity_response_encoding_when_redacted, resolve_provider_chat_pii_redaction,
}; };
use crate::ai_serving::planner::standard::{ 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, 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_cross_format_openai_chat_upstream_url, build_local_openai_chat_request_body,
build_local_openai_chat_upstream_url, request_body_build_failure_extra_data, build_local_openai_chat_upstream_url, codex_model_capabilities_for_transport,
request_conversion_failure_extra_data, 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::antigravity::is_antigravity_provider_transport;
use crate::ai_serving::transport::auth::resolve_local_openai_bearer_auth; use crate::ai_serving::transport::auth::resolve_local_openai_bearer_auth;
@@ -56,7 +56,10 @@ use crate::ai_serving::transport::{
use crate::ai_serving::{ use crate::ai_serving::{
ai_local_execution_contract_for_formats, request_conversion_direct_auth, ai_local_execution_contract_for_formats, request_conversion_direct_auth,
request_conversion_kind, CandidateFailureDiagnostic, GatewayProviderTransportSnapshot, 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::ai_serving::{ConversionMode, ExecutionStrategy};
use crate::stage_metrics::observe_gateway_stage_ms; use crate::stage_metrics::observe_gateway_stage_ms;
@@ -88,37 +91,7 @@ pub(crate) struct LocalOpenAiChatCandidatePayloadParts {
} }
#[derive(Default)] #[derive(Default)]
pub(crate) struct LocalOpenAiChatRequestPreparation { 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
}
}
fn is_grok_text_provider_api_format(provider_api_format: &str) -> bool { fn is_grok_text_provider_api_format(provider_api_format: &str) -> bool {
matches!( 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)] #[allow(clippy::too_many_arguments)]
pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts( pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
state: &AppState, state: &AppState,
@@ -134,7 +166,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
trace_id: &str, trace_id: &str,
body_json: &serde_json::Value, body_json: &serde_json::Value,
input: &LocalOpenAiChatDecisionInput, input: &LocalOpenAiChatDecisionInput,
mut preparation: Option<&mut LocalOpenAiChatRequestPreparation>, _preparation: Option<&mut LocalOpenAiChatRequestPreparation>,
eligible: &EligibleLocalExecutionCandidate, eligible: &EligibleLocalExecutionCandidate,
candidate_index: u32, candidate_index: u32,
candidate_id: &str, candidate_id: &str,
@@ -151,18 +183,9 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
let force_body_stream_field = let force_body_stream_field =
endpoint_config_forces_body_stream_field(transport.endpoint.config.as_ref()); endpoint_config_forces_body_stream_field(transport.endpoint.config.as_ref());
let model_directives_started_at = std::time::Instant::now(); let model_directives_started_at = std::time::Instant::now();
let enable_model_directives = if let Some(preparation) = preparation { let model_directive_resolution = input
preparation .model_directive_policy
.model_directives_enabled(state, provider_api_format, &input.requested_model) .resolve_reasoning(provider_api_format, Some(&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
};
observe_gateway_stage_ms( observe_gateway_stage_ms(
"openai_chat_payload_model_directives", "openai_chat_payload_model_directives",
model_directives_started_at.elapsed().as_millis() as u64, 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); 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, body_json,
&prepared_candidate.mapped_model, &prepared_candidate.mapped_model,
upstream_is_stream, upstream_is_stream,
force_body_stream_field, force_body_stream_field,
transport.endpoint.body_rules.as_ref(), transport.endpoint.body_rules.as_ref(),
effective_headers, effective_headers,
enable_model_directives, false,
) else { ) else {
mark_skipped_local_openai_chat_candidate_with_extra_data( mark_skipped_local_openai_chat_candidate_with_extra_data(
state, state,
@@ -247,6 +288,33 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
.await; .await;
return Ok(None); 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 upstream_url = build_grok_upstream_url(transport, GROK_CHAT_PATH);
let Some(mut provider_request_headers) = build_grok_browser_headers(GrokHeaderInput { 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); 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( observe_gateway_stage_ms(
"openai_chat_payload_auth_prepare", "openai_chat_payload_auth_prepare",
auth_prepare_started_at.elapsed().as_millis() as u64, 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, force_body_stream_field,
transport.endpoint.body_rules.as_ref(), transport.endpoint.body_rules.as_ref(),
effective_headers, effective_headers,
enable_model_directives, false,
) else { ) else {
mark_skipped_local_openai_chat_candidate_with_extra_data( mark_skipped_local_openai_chat_candidate_with_extra_data(
state, state,
@@ -415,13 +501,33 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
"openai_chat_payload_body_build", "openai_chat_payload_body_build",
body_build_started_at.elapsed().as_millis() as u64, 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, &mut provider_request_body,
transport.provider.provider_type.as_str(), model_directive_mapping.as_ref(),
transport.endpoint.base_url.as_str(),
"openai:chat", "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 { let Some(upstream_url) = build_local_openai_chat_upstream_url(parts, transport) else {
mark_skipped_local_openai_chat_candidate_with_failure_diagnostic( 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); return Ok(None);
}; };
let mut provider_request_headers = resolved_headers.headers; let mut provider_request_headers = resolved_headers.headers;
apply_codex_openai_responses_special_headers( apply_codex_openai_special_headers(
&mut provider_request_headers, &mut provider_request_headers,
&provider_request_body, &provider_request_body,
effective_headers, 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( let Some(mut provider_request_body) = build_cross_format_openai_chat_request_body(
body_json, 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()), Some(input.auth_context.api_key_id.as_str()),
effective_headers, effective_headers,
enable_model_directives, false,
) else { ) else {
mark_skipped_local_openai_chat_candidate_with_extra_data( mark_skipped_local_openai_chat_candidate_with_extra_data(
state, state,
@@ -689,34 +813,37 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
.await; .await;
return Ok(None); return Ok(None);
}; };
if let Some(mapping) = if !finalize_openai_chat_provider_request_body(
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(
&mut provider_request_body, &mut provider_request_body,
transport.provider.provider_type.as_str(), model_directive_mapping.as_ref(),
transport.endpoint.base_url.as_str(),
provider_api_format.as_str(), 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() { if let Some(kiro_auth) = kiro_auth.as_ref() {
return Ok(build_kiro_openai_chat_cross_format_payload_parts( 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); return Ok(None);
}; };
let mut provider_request_headers = resolved_headers.headers; let mut provider_request_headers = resolved_headers.headers;
apply_codex_openai_responses_special_headers( apply_codex_openai_special_headers(
&mut provider_request_headers, &mut provider_request_headers,
&provider_request_body, &provider_request_body,
effective_headers, effective_headers,
@@ -997,7 +1124,7 @@ async fn build_antigravity_openai_chat_cross_format_payload_parts(
} }
}; };
let mut provider_request_headers = resolved.headers.headers; let mut provider_request_headers = resolved.headers.headers;
apply_codex_openai_responses_special_headers( apply_codex_openai_special_headers(
&mut provider_request_headers, &mut provider_request_headers,
&resolved.body, &resolved.body,
effective_headers, 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; let mut provider_request_headers = resolved.headers.headers;
apply_codex_openai_responses_special_headers( apply_codex_openai_special_headers(
&mut provider_request_headers, &mut provider_request_headers,
&resolved.body, &resolved.body,
effective_headers, effective_headers,
@@ -1255,6 +1382,19 @@ async fn resolve_openai_chat_to_openai_image_payload_parts(
.provider_type .provider_type
.trim() .trim()
.eq_ignore_ascii_case("chatgpt_web"); .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 { let Some((mut provider_request_body, image_request_summary)) = (if is_chatgpt_web {
build_chatgpt_web_image_provider_body_from_openai_chat_body( build_chatgpt_web_image_provider_body_from_openai_chat_body(
body_json, body_json,
@@ -1263,7 +1403,7 @@ async fn resolve_openai_chat_to_openai_image_payload_parts(
} else { } else {
build_openai_image_provider_body_from_openai_chat_body( build_openai_image_provider_body_from_openai_chat_body(
body_json, body_json,
&input.requested_model, &prepared_candidate.mapped_model,
upstream_is_stream, upstream_is_stream,
) )
}) else { }) else {
@@ -1280,24 +1420,70 @@ async fn resolve_openai_chat_to_openai_image_payload_parts(
.await; .await;
return Ok(None); return Ok(None);
}; };
let Some(operation) = openai_image_operation_from_summary(&image_request_summary) else {
return Ok(None);
};
if !is_chatgpt_web { if !is_chatgpt_web {
apply_codex_openai_responses_special_body_edits( let Some(projected) = project_openai_image_api_request_body(
&mut provider_request_body, &provider_request_body,
transport.provider.provider_type.as_str(), &prepared_candidate.mapped_model,
provider_api_format, operation,
transport.endpoint.body_rules.as_ref(), crate::image_capabilities::openai_image_provider_max_generation_count_for_model(
Some(candidate.key_id.as_str()), 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 { let upstream_url = if is_chatgpt_web {
chatgpt_web_image_internal_url(&transport.endpoint.base_url) chatgpt_web_image_internal_url(&transport.endpoint.base_url)
} else { } else {
build_openai_image_upstream_url( let request_path = match operation {
transport, OpenAiImageOperation::Generate => "/v1/images/generations",
Some("/v1/images/generations"), OpenAiImageOperation::Edit => "/v1/images/edits",
parts.uri.query(), };
) build_openai_image_upstream_url(transport, Some(request_path), parts.uri.query())
}; };
let Some(mut provider_request_headers) = let Some(mut provider_request_headers) =
build_openai_image_headers(ProviderOpenAiImageHeadersInput { build_openai_image_headers(ProviderOpenAiImageHeadersInput {
@@ -1305,7 +1491,13 @@ async fn resolve_openai_chat_to_openai_image_payload_parts(
headers: &parts.headers, headers: &parts.headers,
auth_header: &prepared_candidate.auth_header, auth_header: &prepared_candidate.auth_header,
auth_value: &prepared_candidate.auth_value, 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(), header_rules: transport.endpoint.header_rules.as_ref(),
provider_request_body: &provider_request_body, provider_request_body: &provider_request_body,
original_request_body: body_json, original_request_body: body_json,
@@ -1331,7 +1523,7 @@ async fn resolve_openai_chat_to_openai_image_payload_parts(
if is_chatgpt_web { if is_chatgpt_web {
provider_request_headers.insert("x-aether-chatgpt-web-image".to_string(), "1".to_string()); provider_request_headers.insert("x-aether-chatgpt-web-image".to_string(), "1".to_string());
} else { } else {
apply_codex_openai_responses_special_headers( apply_codex_openai_special_headers(
&mut provider_request_headers, &mut provider_request_headers,
&provider_request_body, &provider_request_body,
&parts.headers, &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, "input_fidelity");
copy_openai_chat_image_option(body_json, &mut image_options, "partial_images"); 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(); let mut body = serde_json::Map::new();
if let Some(model) = body_json let requested_model = requested_model.trim();
.get("model") if requested_model.is_empty() {
.and_then(Value::as_str) return None;
.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()));
} }
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( body.insert(
"tools".to_string(), "model".to_string(),
Value::Array(vec![Value::Object(image_tool)]), 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 { if upstream_is_stream {
body.insert("stream".to_string(), Value::Bool(true)); 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))) Some((Value::Object(body), Value::Object(summary)))
} }
fn openai_image_operation_from_summary(summary: &Value) -> Option<OpenAiImageOperation> {
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( fn build_chatgpt_web_image_provider_body_from_openai_chat_body(
body_json: &Value, body_json: &Value,
requested_model: &str, requested_model: &str,
@@ -1603,6 +1781,20 @@ fn openai_image_inputs_as_urls(images: &[Value]) -> Vec<Value> {
.collect() .collect()
} }
fn openai_image_inputs_as_api_urls(images: &[Value]) -> Vec<Value> {
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 { fn chatgpt_web_ratio_for_size(size: &str) -> String {
let Some((width, height)) = size.split_once('x') else { let Some((width, height)) = size.split_once('x') else {
return "1:1".to_string(); return "1:1".to_string();
@@ -2011,6 +2203,7 @@ mod tests {
routing_policy: None, routing_policy: None,
routing_trace_seed: None, routing_trace_seed: None,
routing_context: 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_id: "global-model-1".to_string(),
global_model_name: "gemini-2.5-pro".to_string(), global_model_name: "gemini-2.5-pro".to_string(),
selected_provider_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, mapping_matched_model: None,
}, },
transport: Arc::new(sample_gemini_cli_transport()), transport: Arc::new(sample_gemini_cli_transport()),
@@ -2138,6 +2332,238 @@ mod tests {
eligible 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] #[tokio::test]
async fn openai_chat_to_gemini_cli_wraps_cross_format_body_in_v1internal_envelope() { async fn openai_chat_to_gemini_cli_wraps_cross_format_body_in_v1internal_envelope() {
let state = AppState::new().expect("state should build"); let state = AppState::new().expect("state should build");
@@ -2333,7 +2759,7 @@ mod tests {
} }
#[test] #[test]
fn openai_chat_image_bridge_body_injects_image_generation_tool() { fn openai_chat_image_bridge_builds_images_api_body() {
let body_json = json!({ let body_json = json!({
"model": "gpt-image-2", "model": "gpt-image-2",
"messages": [ "messages": [
@@ -2347,13 +2773,22 @@ mod tests {
build_openai_image_provider_body_from_openai_chat_body(&body_json, "gpt-image-2", true) build_openai_image_provider_body_from_openai_chat_body(&body_json, "gpt-image-2", true)
.expect("chat image body should convert"); .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["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["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["operation"], "generate");
assert_eq!(summary["output_format"], "png"); 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());
} }
} }
@@ -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( build_lazy_requested_model_execution_candidate_attempt_source_with_serving(
planner_state, planner_state,
&input.model_directive_policy,
trace_id, trace_id,
"openai:chat", "openai:chat",
&input.requested_model, &input.requested_model,
@@ -21,6 +21,7 @@ pub(crate) async fn list_local_openai_chat_candidates(
> { > {
let outcome = preselect_local_execution_candidates_with_serving( let outcome = preselect_local_execution_candidates_with_serving(
PlannerAppState::new(state), PlannerAppState::new(state),
&input.model_directive_policy,
"openai:chat", "openai:chat",
&input.requested_model, &input.requested_model,
require_streaming, require_streaming,
@@ -65,7 +65,9 @@ pub(crate) async fn resolve_local_openai_chat_decision_input(
state, state,
auth_context.clone(), auth_context.clone(),
Some(requested_model.as_str()), Some(requested_model.as_str()),
decision.auth_endpoint_signature.as_deref(),
None, None,
&decision.model_directive_policy,
) )
.await .await
{ {
@@ -111,7 +111,7 @@ pub(crate) async fn build_local_openai_chat_stream_attempt_source<'a>(
input, input,
candidates, candidates,
prefetched_attempts: VecDeque::new(), prefetched_attempts: VecDeque::new(),
request_preparation: LocalOpenAiChatRequestPreparation::default(), request_preparation: LocalOpenAiChatRequestPreparation,
}, },
candidate_count, candidate_count,
))) )))
@@ -129,7 +129,7 @@ pub(crate) fn build_openai_chat_stream_plan_from_decision(
headers: std::mem::take(&mut provider_request_headers), headers: std::mem::take(&mut provider_request_headers),
content_type, content_type,
body: RequestBody::from_json(provider_request_body_value), 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), headers: std::mem::take(&mut provider_request_headers),
content_type, content_type,
body: RequestBody::from_json(provider_request_body_value), 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()), request_id: Some("req_123".to_string()),
candidate_id: Some("cand_123".to_string()), candidate_id: Some("cand_123".to_string()),
provider_name: Some("Codex".to_string()), provider_name: Some("Codex".to_string()),
provider_type: Some("codex".to_string()),
provider_id: Some("prov_123".to_string()), provider_id: Some("prov_123".to_string()),
endpoint_id: Some("ep_123".to_string()), endpoint_id: Some("ep_123".to_string()),
key_id: Some("key_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] #[test]
fn build_openai_chat_stream_plan_fallback_preserves_complete_same_format_headers() { fn build_openai_chat_stream_plan_fallback_preserves_complete_same_format_headers() {
let parts = http::Request::builder() let parts = http::Request::builder()
@@ -404,6 +445,7 @@ mod tests {
request_id: Some("req_stream_456".to_string()), request_id: Some("req_stream_456".to_string()),
candidate_id: Some("cand_stream_456".to_string()), candidate_id: Some("cand_stream_456".to_string()),
provider_name: Some("OpenAI".to_string()), provider_name: Some("OpenAI".to_string()),
provider_type: Some("openai".to_string()),
provider_id: Some("prov_stream_456".to_string()), provider_id: Some("prov_stream_456".to_string()),
endpoint_id: Some("ep_stream_456".to_string()), endpoint_id: Some("ep_stream_456".to_string()),
key_id: Some("key_stream_456".to_string()), key_id: Some("key_stream_456".to_string()),
@@ -462,7 +504,7 @@ mod tests {
} }
#[test] #[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<Value>) -> AiExecutionDecision { fn force_non_stream_payload(provider_request_body: Option<Value>) -> AiExecutionDecision {
AiExecutionDecision { AiExecutionDecision {
action: "stream".to_string(), action: "stream".to_string(),
@@ -472,6 +514,7 @@ mod tests {
request_id: Some("req_force_non_stream".to_string()), request_id: Some("req_force_non_stream".to_string()),
candidate_id: Some("cand_force_non_stream".to_string()), candidate_id: Some("cand_force_non_stream".to_string()),
provider_name: Some("OpenAI".to_string()), provider_name: Some("OpenAI".to_string()),
provider_type: Some("openai".to_string()),
provider_id: Some("prov_force_non_stream".to_string()), provider_id: Some("prov_force_non_stream".to_string()),
endpoint_id: Some("ep_force_non_stream".to_string()), endpoint_id: Some("ep_force_non_stream".to_string()),
key_id: Some("key_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 build should succeed")
.expect("plan should be produced"); .expect("plan should be produced");
assert!(built.plan.stream); assert!(!built.plan.stream);
assert_eq!( assert_eq!(
built built
.plan .plan
@@ -548,7 +591,7 @@ mod tests {
.expect("fallback plan build should succeed") .expect("fallback plan build should succeed")
.expect("fallback plan should be produced"); .expect("fallback plan should be produced");
assert!(built.plan.stream); assert!(!built.plan.stream);
assert_eq!( assert_eq!(
built built
.plan .plan
@@ -579,6 +622,7 @@ mod tests {
request_id: Some("req_stream_789".to_string()), request_id: Some("req_stream_789".to_string()),
candidate_id: Some("cand_stream_789".to_string()), candidate_id: Some("cand_stream_789".to_string()),
provider_name: Some("Claude".to_string()), provider_name: Some("Claude".to_string()),
provider_type: Some("anthropic".to_string()),
provider_id: Some("prov_stream_789".to_string()), provider_id: Some("prov_stream_789".to_string()),
endpoint_id: Some("ep_stream_789".to_string()), endpoint_id: Some("ep_stream_789".to_string()),
key_id: Some("key_stream_789".to_string()), key_id: Some("key_stream_789".to_string()),
@@ -257,6 +257,7 @@ mod tests {
request_id: Some("req_123".to_string()), request_id: Some("req_123".to_string()),
candidate_id: Some("cand_123".to_string()), candidate_id: Some("cand_123".to_string()),
provider_name: Some("Codex".to_string()), provider_name: Some("Codex".to_string()),
provider_type: Some("codex".to_string()),
provider_id: Some("prov_123".to_string()), provider_id: Some("prov_123".to_string()),
endpoint_id: Some("ep_123".to_string()), endpoint_id: Some("ep_123".to_string()),
key_id: Some("key_123".to_string()), key_id: Some("key_123".to_string()),
@@ -369,6 +370,7 @@ mod tests {
request_id: Some("req_456".to_string()), request_id: Some("req_456".to_string()),
candidate_id: Some("cand_456".to_string()), candidate_id: Some("cand_456".to_string()),
provider_name: Some("OpenAI".to_string()), provider_name: Some("OpenAI".to_string()),
provider_type: Some("openai".to_string()),
provider_id: Some("prov_456".to_string()), provider_id: Some("prov_456".to_string()),
endpoint_id: Some("ep_456".to_string()), endpoint_id: Some("ep_456".to_string()),
key_id: Some("key_456".to_string()), key_id: Some("key_456".to_string()),
@@ -440,6 +442,7 @@ mod tests {
request_id: Some("req_789".to_string()), request_id: Some("req_789".to_string()),
candidate_id: Some("cand_789".to_string()), candidate_id: Some("cand_789".to_string()),
provider_name: Some("Claude".to_string()), provider_name: Some("Claude".to_string()),
provider_type: Some("anthropic".to_string()),
provider_id: Some("prov_789".to_string()), provider_id: Some("prov_789".to_string()),
endpoint_id: Some("ep_789".to_string()), endpoint_id: Some("ep_789".to_string()),
key_id: Some("key_789".to_string()), key_id: Some("key_789".to_string()),
@@ -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::spec_metadata::local_openai_responses_spec_metadata;
use crate::ai_serving::planner::{ 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, AiExecutionDecisionResponseParts,
}; };
use crate::ai_serving::transport::{ 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: _, image_request_summary: _,
request_redacted: _, request_redacted: _,
} = resolved; } = 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 { let mut decision = build_ai_execution_decision_response(AiExecutionDecisionResponseParts {
decision_is_stream: spec_metadata.require_streaming, 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(), request_id: trace_id.to_string(),
candidate_id: candidate_id.clone(), candidate_id: candidate_id.clone(),
provider_name: transport.provider.name.clone(), provider_name: transport.provider.name.clone(),
provider_type: transport.provider.provider_type.clone(),
provider_id: candidate.provider_id.clone(), provider_id: candidate.provider_id.clone(),
endpoint_id: candidate.endpoint_id.clone(), endpoint_id: candidate.endpoint_id.clone(),
key_id: candidate.key_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: Some(provider_request_body),
provider_request_body_base64: None, provider_request_body_base64: None,
content_type: Some("application/json".to_string()), content_type: Some("application/json".to_string()),
content_encoding: None, content_encoding: request_encoding.content_encoding,
request_gzip, request_gzip: request_encoding.request_gzip,
proxy, proxy,
transport_profile, transport_profile,
timeouts, timeouts,
@@ -241,6 +242,10 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
report_context: Some(report_context), report_context: Some(report_context),
auth_context: input.auth_context.clone(), 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)) Ok(Some(decision))
} }
@@ -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::spec_metadata::local_openai_responses_spec_metadata;
use crate::ai_serving::planner::standard::{ use crate::ai_serving::planner::standard::{
apply_codex_openai_responses_special_body_edits, apply_codex_openai_responses_special_headers, apply_codex_openai_special_headers, apply_deepseek_tool_call_thinking_compat,
apply_deepseek_tool_call_thinking_compat, 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_cross_format_openai_responses_upstream_url,
build_local_openai_responses_upstream_url, request_body_build_failure_extra_data, build_local_openai_responses_request_body_with_codex_model_capabilities,
request_conversion_failure_extra_data, 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::antigravity::is_antigravity_provider_transport;
use crate::ai_serving::transport::auth::{ use crate::ai_serving::transport::auth::{
@@ -58,7 +59,10 @@ use crate::ai_serving::transport::{
use crate::ai_serving::{ use crate::ai_serving::{
ai_local_execution_contract_for_formats, request_conversion_direct_auth, ai_local_execution_contract_for_formats, request_conversion_direct_auth,
request_conversion_kind, CandidateFailureDiagnostic, GatewayProviderTransportSnapshot, 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::ai_serving::{ConversionMode, ExecutionStrategy};
use crate::{AppState, GatewayError}; 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_header = prepared_candidate.auth_header;
let auth_value = prepared_candidate.auth_value; let auth_value = prepared_candidate.auth_value;
let mapped_model = prepared_candidate.mapped_model; let mapped_model = prepared_candidate.mapped_model;
let enable_model_directives = let model_directive_resolution = input
crate::system_features::reasoning_model_directive_enabled_for_api_format_and_model( .model_directive_policy
state, .resolve_reasoning(provider_api_format, Some(&input.requested_model));
provider_api_format, let model_directive_mapping =
Some(&input.requested_model), match model_directive_resolution.mapping_patch_for_mapped_model(&mapped_model) {
) Ok(mapping) => mapping,
.await; 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( let redaction = resolve_provider_chat_pii_redaction(
state, state,
parts, parts,
@@ -311,9 +328,19 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
let force_body_stream_field = let force_body_stream_field =
endpoint_config_forces_body_stream_field(transport.endpoint.config.as_ref()); endpoint_config_forces_body_stream_field(transport.endpoint.config.as_ref());
let effective_headers = input.effective_headers(&parts.headers); 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) = let Some(mut base_provider_request_body) =
(if is_grok && is_grok_text_provider_api_format(provider_api_format) { (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, body_json,
&mapped_model, &mapped_model,
upstream_is_stream, upstream_is_stream,
@@ -321,12 +348,12 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
transport.provider.provider_type.as_str(), transport.provider.provider_type.as_str(),
spec_metadata.api_format, spec_metadata.api_format,
transport.endpoint.body_rules.as_ref(), transport.endpoint.body_rules.as_ref(),
Some(input.auth_context.api_key_id.as_str()),
effective_headers, effective_headers,
enable_model_directives, codex_model_capabilities.as_ref(),
false,
) )
} else if needs_bidirectional_conversion { } 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, body_json,
&mapped_model, &mapped_model,
spec_metadata.api_format, spec_metadata.api_format,
@@ -339,12 +366,12 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
} else { } else {
transport.endpoint.body_rules.as_ref() transport.endpoint.body_rules.as_ref()
}, },
Some(input.auth_context.api_key_id.as_str()),
effective_headers, effective_headers,
enable_model_directives, codex_model_capabilities.as_ref(),
false,
) )
} else { } else {
build_local_openai_responses_request_body( build_local_openai_responses_request_body_with_codex_model_capabilities(
body_json, body_json,
&mapped_model, &mapped_model,
upstream_is_stream, upstream_is_stream,
@@ -356,9 +383,9 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
} else { } else {
transport.endpoint.body_rules.as_ref() transport.endpoint.body_rules.as_ref()
}, },
Some(input.auth_context.api_key_id.as_str()),
effective_headers, effective_headers,
enable_model_directives, codex_model_capabilities.as_ref(),
false,
) )
}) })
else { else {
@@ -383,17 +410,10 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
.await; .await;
return Ok(None); return Ok(None);
}; };
if let Some(mapping) = if let Some(mapping) = model_directive_mapping.as_ref() {
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( crate::ai_serving::apply_model_directive_mapping_patch(
&mut base_provider_request_body, &mut base_provider_request_body,
&mapping, mapping,
); );
// Directive mapping is a deep-merge patch and may overwrite/add `stream`; // Directive mapping is a deep-merge patch and may overwrite/add `stream`;
// re-enforce stream-field policy afterward. // re-enforce stream-field policy afterward.
@@ -411,6 +431,46 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
provider_api_format, provider_api_format,
Some(body_json), 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; let provider_request_body = base_provider_request_body;
if let Some(kiro_auth) = kiro_auth.as_ref() { 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; let mut provider_request_headers = resolved_headers.headers;
if !is_grok { 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, &mut provider_request_headers,
&provider_request_body, &provider_request_body,
effective_headers, effective_headers,
@@ -621,9 +685,13 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
Some(trace_id), Some(trace_id),
transport.key.decrypted_auth_config.as_deref(), 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, &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( 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; 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, &mut provider_request_headers,
&resolved.body, &resolved.body,
effective_headers, effective_headers,
@@ -797,10 +869,6 @@ async fn build_antigravity_openai_responses_payload_parts(
Some(trace_id), Some(trace_id),
resolved.transport.key.decrypted_auth_config.as_deref(), 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()); provider_request_headers.insert("accept".to_string(), "text/event-stream".to_string());
request_identity_response_encoding_when_redacted( request_identity_response_encoding_when_redacted(
&mut provider_request_headers, &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; 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, &mut provider_request_headers,
&resolved.body, &resolved.body,
effective_headers, effective_headers,
@@ -950,10 +1022,6 @@ async fn build_gemini_cli_openai_responses_payload_parts(
Some(trace_id), Some(trace_id),
resolved.transport.key.decrypted_auth_config.as_deref(), 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( request_identity_response_encoding_when_redacted(
&mut provider_request_headers, &mut provider_request_headers,
request_redacted, request_redacted,
@@ -1182,11 +1250,16 @@ async fn resolve_openai_responses_to_openai_image_payload_parts(
.provider_type .provider_type
.trim() .trim()
.eq_ignore_ascii_case("chatgpt_web"); .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( let upstream_is_stream = resolve_upstream_is_stream_for_provider(
transport.endpoint.config.as_ref(), transport.endpoint.config.as_ref(),
transport.provider.provider_type.as_str(), transport.provider.provider_type.as_str(),
provider_api_format, provider_api_format,
spec_metadata.require_streaming, spec_metadata.require_streaming && candidate.supports_streaming,
false, false,
); );
let Some((mut provider_request_body, image_request_summary)) = (if is_chatgpt_web { 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 { } else {
build_openai_image_provider_body_from_openai_responses_body( build_openai_image_provider_body_from_openai_responses_body(
body_json, body_json,
&input.requested_model, &prepared_candidate.mapped_model,
upstream_is_stream, upstream_is_stream,
) )
}) else { }) else {
@@ -1218,25 +1291,31 @@ async fn resolve_openai_responses_to_openai_image_payload_parts(
.await; .await;
return None; return None;
}; };
let operation = openai_image_operation_from_summary(&image_request_summary)?;
if !is_chatgpt_web { if !is_chatgpt_web {
apply_codex_openai_responses_special_body_edits( provider_request_body = project_openai_image_api_request_body(
&mut provider_request_body, &provider_request_body,
transport.provider.provider_type.as_str(), &prepared_candidate.mapped_model,
provider_api_format, operation,
transport.endpoint.body_rules.as_ref(), crate::image_capabilities::openai_image_provider_max_generation_count_for_model(
Some(candidate.key_id.as_str()), 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 { let upstream_url = if is_chatgpt_web {
chatgpt_web_image_internal_url(&transport.endpoint.base_url) chatgpt_web_image_internal_url(&transport.endpoint.base_url)
} else { } else {
build_openai_image_upstream_url( let request_path = match operation {
transport, OpenAiImageOperation::Generate => "/v1/images/generations",
Some("/v1/images/generations"), OpenAiImageOperation::Edit => "/v1/images/edits",
parts.uri.query(), };
) build_openai_image_upstream_url(transport, Some(request_path), parts.uri.query())
}; };
let Some(mut provider_request_headers) = let Some(mut provider_request_headers) =
build_openai_image_headers(ProviderOpenAiImageHeadersInput { build_openai_image_headers(ProviderOpenAiImageHeadersInput {
@@ -1244,7 +1323,13 @@ async fn resolve_openai_responses_to_openai_image_payload_parts(
headers: &parts.headers, headers: &parts.headers,
auth_header: &prepared_candidate.auth_header, auth_header: &prepared_candidate.auth_header,
auth_value: &prepared_candidate.auth_value, 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(), header_rules: transport.endpoint.header_rules.as_ref(),
provider_request_body: &provider_request_body, provider_request_body: &provider_request_body,
original_request_body: body_json, original_request_body: body_json,
@@ -1270,7 +1355,11 @@ async fn resolve_openai_responses_to_openai_image_payload_parts(
if is_chatgpt_web { if is_chatgpt_web {
provider_request_headers.insert("x-aether-chatgpt-web-image".to_string(), "1".to_string()); provider_request_headers.insert("x-aether-chatgpt-web-image".to_string(), "1".to_string());
} else { } 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, &mut provider_request_headers,
&provider_request_body, &provider_request_body,
&parts.headers, &parts.headers,
@@ -1279,10 +1368,6 @@ async fn resolve_openai_responses_to_openai_image_payload_parts(
Some(trace_id), Some(trace_id),
transport.key.decrypted_auth_config.as_deref(), 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) = let (execution_strategy, conversion_mode) =
@@ -1314,59 +1399,68 @@ fn build_openai_image_provider_body_from_openai_responses_body(
upstream_is_stream: bool, upstream_is_stream: bool,
) -> Option<(Value, Value)> { ) -> Option<(Value, Value)> {
let object = body_json.as_object()?; let object = body_json.as_object()?;
let input = object.get("input")?.clone();
let tool = openai_responses_image_generation_tool(object); 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(); let mut body = serde_json::Map::new();
body.insert("input".to_string(), input); let requested_model = requested_model.trim();
if let Some(model) = object if requested_model.is_empty() {
.get("model") return None;
.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()));
} }
body.insert(
"model".to_string(),
Value::String(requested_model.to_string()),
);
body.insert("prompt".to_string(), Value::String(prompt));
for key in [ for key in [
"background",
"quality",
"size",
"output_format",
"output_compression",
"moderation",
"input_fidelity",
"partial_images",
"n",
"user", "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()); 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 { if upstream_is_stream {
body.insert("stream".to_string(), Value::Bool(true)); 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(); let mut summary = serde_json::Map::new();
summary.insert( summary.insert(
"operation".to_string(), "operation".to_string(),
tool.as_ref() Value::String(operation.as_str().to_string()),
.and_then(|tool| tool.get("action"))
.cloned()
.unwrap_or_else(|| json!("generate")),
); );
for key in ["output_format", "partial_images", "size", "quality"] { for key in ["output_format", "partial_images", "size", "quality"] {
let tool_value = tool.as_ref().and_then(|tool| tool.get(key)); 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))) Some((Value::Object(body), Value::Object(summary)))
} }
fn openai_image_operation_from_summary(summary: &Value) -> Option<OpenAiImageOperation> {
match summary.get("operation")?.as_str()? {
"generate" => Some(OpenAiImageOperation::Generate),
"edit" => Some(OpenAiImageOperation::Edit),
_ => None,
}
}
fn openai_responses_image_generation_tool( fn openai_responses_image_generation_tool(
object: &serde_json::Map<String, Value>, object: &serde_json::Map<String, Value>,
) -> Option<serde_json::Map<String, Value>> { ) -> Option<serde_json::Map<String, Value>> {
@@ -1572,6 +1674,20 @@ fn openai_image_inputs_as_urls(images: &[Value]) -> Vec<Value> {
.collect() .collect()
} }
fn openai_image_inputs_as_api_urls(images: &[Value]) -> Vec<Value> {
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 { fn chatgpt_web_ratio_for_size(size: &str) -> String {
let Some((width, height)) = size.split_once('x') else { let Some((width, height)) = size.split_once('x') else {
return "1:1".to_string(); return "1:1".to_string();
@@ -1769,7 +1885,7 @@ mod tests {
use super::*; use super::*;
#[test] #[test]
fn openai_responses_image_bridge_body_preserves_image_generation_tool() { fn openai_responses_image_bridge_builds_images_api_body() {
let body_json = json!({ let body_json = json!({
"model": "gpt-image-2", "model": "gpt-image-2",
"input": "Draw a glass city", "input": "Draw a glass city",
@@ -1792,14 +1908,74 @@ mod tests {
) )
.expect("responses image body should convert"); .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["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_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["operation"], "generate");
assert_eq!(summary["output_format"], "png"); 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] #[test]
@@ -91,7 +91,9 @@ pub(crate) async fn resolve_local_openai_responses_decision_input(
state, state,
auth_context.clone(), auth_context.clone(),
Some(requested_model.as_str()), Some(requested_model.as_str()),
decision.auth_endpoint_signature.as_deref(),
None, None,
&decision.model_directive_policy,
) )
.await .await
{ {
@@ -171,6 +173,7 @@ pub(crate) async fn materialize_local_openai_responses_candidate_attempts(
); );
let preselection = preselect_local_execution_candidates_with_serving( let preselection = preselect_local_execution_candidates_with_serving(
planner_state, planner_state,
&input.model_directive_policy,
spec_metadata.api_format, spec_metadata.api_format,
&input.requested_model, &input.requested_model,
spec_metadata.require_streaming, spec_metadata.require_streaming,
@@ -280,6 +283,7 @@ pub(crate) async fn build_local_openai_responses_candidate_attempt_source<'a>(
Ok( Ok(
build_lazy_requested_model_execution_candidate_attempt_source_with_serving( build_lazy_requested_model_execution_candidate_attempt_source_with_serving(
planner_state, planner_state,
&input.model_directive_policy,
trace_id, trace_id,
spec_metadata.api_format, spec_metadata.api_format,
&input.requested_model, &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( let preselection = preselect_local_execution_candidates_for_api_formats_with_serving(
planner_state, planner_state,
&input.model_directive_policy,
spec_metadata.api_format, spec_metadata.api_format,
&input.requested_model, &input.requested_model,
false, false,
@@ -146,6 +146,7 @@ pub(crate) fn build_standard_stream_plan_from_decision(
&provider_request_headers, &provider_request_headers,
&provider_request_body_value, &provider_request_body_value,
)?; )?;
let stream = payload.upstream_is_stream;
let plan = build_ai_execution_plan_from_decision( let plan = build_ai_execution_plan_from_decision(
&mut payload, &mut payload,
AiExecutionPlanFromDecisionParts { AiExecutionPlanFromDecisionParts {
@@ -155,7 +156,7 @@ pub(crate) fn build_standard_stream_plan_from_decision(
headers: std::mem::take(&mut provider_request_headers), headers: std::mem::take(&mut provider_request_headers),
content_type, content_type,
body: RequestBody::from_json(provider_request_body_value), body: RequestBody::from_json(provider_request_body_value),
stream: true, stream,
}, },
); );
@@ -10,16 +10,15 @@ impl<'a> PlannerAppState<'a> {
api_key_id: &str, api_key_id: &str,
requested_model: Option<&str>, requested_model: Option<&str>,
explicit_required_capabilities: Option<&Value>, explicit_required_capabilities: Option<&Value>,
model_directive_base_model: Option<&str>,
) -> Option<Value> { ) -> Option<Value> {
let enable_model_directives =
crate::system_features::reasoning_model_directive_enabled(self.app()).await;
crate::request_candidate_runtime::resolve_request_candidate_required_capabilities( crate::request_candidate_runtime::resolve_request_candidate_required_capabilities(
self.app(), self.app(),
user_id, user_id,
api_key_id, api_key_id,
requested_model, requested_model,
explicit_required_capabilities, explicit_required_capabilities,
enable_model_directives, model_directive_base_model,
) )
.await .await
} }
@@ -20,14 +20,8 @@ impl<'a> PlannerAppState<'a> {
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>, auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>, client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64, now_unix_secs: u64,
enable_model_directives: bool,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> { ) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, 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( crate::scheduler::candidate::list_selectable_candidates(
self.app().data.as_ref(), self.app().data.as_ref(),
self.app(), self.app(),
@@ -52,6 +46,7 @@ impl<'a> PlannerAppState<'a> {
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>, auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>, client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64, now_unix_secs: u64,
enable_model_directives: bool,
) -> Result< ) -> Result<
( (
Vec<SchedulerMinimalCandidateSelectionCandidate>, Vec<SchedulerMinimalCandidateSelectionCandidate>,
@@ -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_interval = Duration::from_millis(API_KEY_CONCURRENCY_WAIT_POLL_INTERVAL_MS.max(1));
let wait_deadline = Instant::now() + wait_timeout; let wait_deadline = Instant::now() + wait_timeout;
let mut attempt_now_unix_secs = now_unix_secs; 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 { loop {
let result = crate::scheduler::candidate::list_selectable_candidates_with_skip_reasons( let result = crate::scheduler::candidate::list_selectable_candidates_with_skip_reasons(
self.app().data.as_ref(), self.app().data.as_ref(),
+67 -44
View File
@@ -3,10 +3,14 @@ pub(crate) use aether_ai_formats::api::{
aggregate_openai_chat_stream_sync_response, aggregate_openai_responses_stream_sync_response, aggregate_openai_chat_stream_sync_response, aggregate_openai_responses_stream_sync_response,
aggregate_standard_chat_stream_sync_response, aggregate_standard_cli_stream_sync_response, aggregate_standard_chat_stream_sync_response, aggregate_standard_cli_stream_sync_response,
api_format_alias_matches, api_format_storage_aliases, 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_compact_terminal_headers, apply_codex_openai_responses_chat_body_edits,
apply_codex_openai_responses_special_headers, apply_model_directive_mapping_patch, 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_model_directive_overrides_from_model, apply_model_directive_overrides_from_request,
apply_openai_responses_compact_special_body_edits, build_chatgpt_web_image_request_body, 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_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_chat_request_body_with_model_directives,
build_cross_format_openai_responses_request_body, 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, convert_standard_chat_response, convert_standard_cli_response, copy_request_number_field,
copy_request_number_field_as, core_error_background_report_kind, copy_request_number_field_as, core_error_background_report_kind,
core_error_default_client_api_format, core_success_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, default_model_for_openai_image_operation, encode_done_sse, encode_json_sse,
encode_kiro_sse_events, endpoint_config_forces_upstream_stream_policy, encode_kiro_sse_events, endpoint_config_forces_upstream_stream_policy,
enforce_request_body_stream_field, estimate_kiro_tokens, extract_openai_text_content, 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_end_tag, find_kiro_real_thinking_end_tag_at_buffer_end,
find_kiro_real_thinking_start_tag, force_upstream_streaming_for_provider, find_kiro_real_thinking_start_tag, forbid_upstream_streaming_for_provider,
gemini_request_is_image_generation, implicit_sync_finalize_report_kind, force_upstream_streaming_for_provider, gemini_request_is_image_generation,
is_core_error_finalize_kind, is_matching_stream_http_request, is_matching_stream_request, implicit_sync_finalize_report_kind, is_core_error_finalize_kind,
is_openai_image_stream_request, is_openai_responses_family_format, is_openai_responses_format, is_matching_stream_http_request, is_matching_stream_request, is_openai_image_stream_request,
kiro_crc32, map_claude_stop_reason, map_openai_reasoning_effort_to_claude_output, is_openai_responses_compact_format, is_openai_responses_family_format,
map_openai_reasoning_effort_to_gemini_budget, maybe_bridge_standard_sync_json_to_stream, is_openai_responses_format, kiro_crc32, map_claude_stop_reason,
maybe_build_ai_surface_stream_rewriter, 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_chat_cross_format_sync_product_from_normalized_payload,
maybe_build_openai_image_sync_finalize_product, maybe_build_openai_image_sync_finalize_product,
maybe_build_openai_responses_cross_format_sync_product_from_normalized_payload, 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_cross_format_sync_product_from_normalized_payload,
maybe_build_standard_same_format_sync_body_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, maybe_build_standard_sync_finalize_product_from_normalized_payload, model_directive_base_model,
normalize_api_format_alias, normalize_claude_request_to_openai_chat_request, model_directive_builtin_suffix_supported_for_source_model,
normalize_gemini_request_to_openai_chat_request, normalize_openai_image_request, model_directive_suffix_has_builtin_mapping, normalize_api_format_alias,
normalize_openai_image_request_with_options, 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_openai_responses_request_to_openai_chat_request,
normalize_provider_private_report_context, normalize_provider_private_response_value, normalize_provider_private_report_context, normalize_provider_private_response_value,
normalize_standard_request_to_openai_chat_request, openai_image_operation_from_path, 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, parse_codex_auth_identity, parse_direct_request_body, parse_model_directive,
prepare_local_success_response_parts, prepare_local_success_response_parts_owned, parse_model_directive_with_suffixes, parse_openai_stop_sequences,
provider_adaptation_allows_sync_finalize_envelope, provider_adaptation_anchor_api_format, parse_openai_tool_result_content, prepare_local_success_response_parts,
provider_adaptation_descriptor_for_envelope, provider_adaptation_descriptor_for_provider_type, 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_requires_eventstream_accept,
provider_adaptation_should_unwrap_stream_envelope, provider_adaptation_should_unwrap_stream_envelope,
provider_private_response_allows_sync_finalize, request_candidate_api_format_preference, provider_private_response_allows_sync_finalize, request_candidate_api_format_preference,
request_candidate_api_formats, request_conversion_kind, request_candidate_api_formats, request_conversion_kind,
request_conversion_requires_enable_flag, request_path_implies_stream_request, request_conversion_requires_enable_flag, request_path_implies_stream_request,
resolve_claude_stream_spec, resolve_claude_sync_spec, resolve_claude_stream_spec, resolve_claude_sync_spec,
resolve_execution_runtime_stream_plan_kind, resolve_execution_runtime_sync_plan_kind, resolve_codex_responses_model_capabilities, resolve_execution_runtime_stream_plan_kind,
resolve_finalize_stream_rewrite_mode, resolve_gemini_files_stream_spec, resolve_execution_runtime_sync_plan_kind, resolve_finalize_stream_rewrite_mode,
resolve_gemini_files_sync_spec, resolve_gemini_stream_spec, resolve_gemini_sync_spec, resolve_gemini_files_stream_spec, resolve_gemini_files_sync_spec, resolve_gemini_stream_spec,
resolve_local_image_stream_spec, resolve_local_image_sync_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_same_format_stream_spec, resolve_local_same_format_sync_spec,
resolve_local_video_sync_spec, resolve_openai_chat_max_tokens, resolve_local_video_sync_spec, resolve_openai_chat_max_tokens,
resolve_openai_embedding_sync_spec, resolve_openai_responses_stream_spec, resolve_openai_embedding_sync_spec, resolve_openai_responses_stream_spec,
resolve_openai_responses_sync_spec, resolve_requested_gemini_image_model_for_request, resolve_openai_responses_sync_spec, resolve_requested_gemini_image_model_for_request,
resolve_requested_openai_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, resolve_upstream_is_stream_from_endpoint_config, sanitize_request_path,
sanitize_request_path_and_query, sanitize_request_query_string, sanitize_request_path_and_query, sanitize_request_query_string,
stream_body_contains_error_event, supports_stream_execution_decision_kind, stream_body_contains_error_event, supports_stream_execution_decision_kind,
supports_sync_execution_decision_kind, sync_chat_response_conversion_kind, supports_sync_execution_decision_kind, sync_chat_response_conversion_kind,
sync_cli_response_conversion_kind, transform_provider_private_stream_line, value_as_u64, sync_cli_response_conversion_kind, transform_provider_private_stream_line,
AiControlPlanRequest, AiSurfaceFinalizeError, AiSurfaceStreamRewriter, CanonicalStreamFrame, validate_openai_provider_request_contract, value_as_u64, AiControlPlanRequest,
AiSurfaceFinalizeError, AiSurfaceStreamRewriter, CanonicalStreamFrame,
ChatGptWebImageRequestError, ClaudeClientEmitter, ClaudeProviderState, ChatGptWebImageRequestError, ClaudeClientEmitter, ClaudeProviderState,
ExecutionRuntimeAuthContext, FinalizeStreamRewriteMode, FormatContext, GeminiClientEmitter, CodexResponsesModelCapabilities, ExecutionRuntimeAuthContext, FinalizeStreamRewriteMode,
GeminiImageRequestForOpenAi, GeminiProviderState, KiroToClaudeCliStreamState, FormatContext, GeminiClientEmitter, GeminiImageRequestForOpenAi, GeminiProviderState,
LocalCoreSyncErrorKind, LocalGeminiFilesSpec, LocalOpenAiImageSpec, LocalOpenAiResponsesSpec, KiroToClaudeCliStreamState, LocalCoreSyncErrorKind, LocalGeminiFilesSpec, LocalOpenAiImageSpec,
LocalSameFormatProviderFamily, LocalSameFormatProviderSpec, LocalStandardSourceFamily, LocalOpenAiResponsesSpec, LocalSameFormatProviderFamily, LocalSameFormatProviderSpec,
LocalStandardSourceMode, LocalStandardSpec, LocalSyncReportParts, LocalVideoCreateFamily, LocalStandardSourceFamily, LocalStandardSourceMode, LocalStandardSpec, LocalSyncReportParts,
LocalVideoCreateSpec, NormalizedOpenAiImageRequest, OpenAIChatClientEmitter, LocalVideoCreateFamily, LocalVideoCreateSpec, NormalizedOpenAiImageRequest,
OpenAIChatProviderState, OpenAIResponsesClientEmitter, OpenAIResponsesProviderState, OpenAIChatClientEmitter, OpenAIChatProviderState, OpenAIResponsesClientEmitter,
OpenAiImageNormalizeOptions, OpenAiImageOperation, OpenAiImageRequestForGemini, OpenAIResponsesProviderState, OpenAiImageNormalizeOptions, OpenAiImageOperation,
OpenAiImageResponseFormat, OpenAiImageStreamState, OpenAiImageSyncFinalizeProduct, OpenAiImageRequestForGemini, OpenAiImageResponseFormat, OpenAiImageStreamState,
OpenAiImageSyncFinalizeProduct, OpenAiProviderRequestFinalization,
ProviderAdaptationDescriptor, ProviderAdaptationSurface, ProviderPrivateStreamNormalizer, ProviderAdaptationDescriptor, ProviderAdaptationSurface, ProviderPrivateStreamNormalizer,
RequestConversionKind, StandardCrossFormatSyncProduct, StandardSyncFinalizeNormalizedProduct, ReasoningEffort, RequestConversionKind, ServiceTier, StandardCrossFormatSyncProduct,
StreamingStandardFormatMatrix, SyncChatResponseConversionKind, SyncCliResponseConversionKind, StandardSyncFinalizeNormalizedProduct, StreamingStandardFormatMatrix,
SyncToStreamBridgeOutcome, ANTIGRAVITY_V1INTERNAL_ENVELOPE_NAME, CLAUDE_CHAT_STREAM_PLAN_KIND, 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_STREAM_SUCCESS_REPORT_KIND, CLAUDE_CHAT_SYNC_ERROR_REPORT_KIND,
CLAUDE_CHAT_SYNC_FINALIZE_REPORT_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CHAT_SYNC_FINALIZE_REPORT_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND,
CLAUDE_CHAT_SYNC_SUCCESS_REPORT_KIND, CLAUDE_CLI_STREAM_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_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_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, 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, KIRO_ENVELOPE_NAME, KIRO_MAX_THINKING_BUFFER, MODEL_DIRECTIVE_API_FORMATS,
OPENAI_CHAT_STREAM_SUCCESS_REPORT_KIND, OPENAI_CHAT_SYNC_ERROR_REPORT_KIND, OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_STREAM_SUCCESS_REPORT_KIND,
OPENAI_CHAT_SYNC_FINALIZE_REPORT_KIND, OPENAI_CHAT_SYNC_PLAN_KIND, OPENAI_CHAT_SYNC_ERROR_REPORT_KIND, OPENAI_CHAT_SYNC_FINALIZE_REPORT_KIND,
OPENAI_CHAT_SYNC_SUCCESS_REPORT_KIND, OPENAI_EMBEDDING_SYNC_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND, OPENAI_CHAT_SYNC_SUCCESS_REPORT_KIND,
OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_STREAM_SUCCESS_REPORT_KIND, OPENAI_EMBEDDING_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND,
OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_SUCCESS_REPORT_KIND, OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND,
OPENAI_IMAGE_SYNC_SUCCESS_REPORT_KIND, OPENAI_RERANK_SYNC_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND, OPENAI_IMAGE_SYNC_SUCCESS_REPORT_KIND,
OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, OPENAI_RESPONSES_COMPACT_STREAM_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_ERROR_REPORT_KIND,
OPENAI_RESPONSES_COMPACT_SYNC_FINALIZE_REPORT_KIND, OPENAI_RESPONSES_COMPACT_SYNC_PLAN_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_COMPACT_SYNC_SUCCESS_REPORT_KIND, OPENAI_RESPONSES_STREAM_PLAN_KIND,
OPENAI_RESPONSES_STREAM_SUCCESS_REPORT_KIND, OPENAI_RESPONSES_SYNC_ERROR_REPORT_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_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_CONTENT_PLAN_KIND, OPENAI_VIDEO_CREATE_SYNC_FINALIZE_REPORT_KIND,
OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND, OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND,
OPENAI_VIDEO_REMIX_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,
};
+2
View File
@@ -5,6 +5,7 @@ pub(crate) fn normalized_signature(api_format: &str) -> Option<&'static str> {
"openai:rerank" => Some("openai:rerank"), "openai:rerank" => Some("openai:rerank"),
"openai:responses" => Some("openai:responses"), "openai:responses" => Some("openai:responses"),
"openai:responses:compact" => Some("openai:responses:compact"), "openai:responses:compact" => Some("openai:responses:compact"),
"openai:search" => Some("openai:search"),
"openai:image" => Some("openai:image"), "openai:image" => Some("openai:image"),
"openai:video" => Some("openai:video"), "openai:video" => Some("openai:video"),
_ => None, _ => None,
@@ -18,6 +19,7 @@ pub(crate) fn local_path(api_format: &str) -> Option<&'static str> {
"openai:rerank" => Some("/v1/rerank"), "openai:rerank" => Some("/v1/rerank"),
"openai:responses" => Some("/v1/responses"), "openai:responses" => Some("/v1/responses"),
"openai:responses:compact" => Some("/v1/responses/compact"), "openai:responses:compact" => Some("/v1/responses/compact"),
"openai:search" => Some("/v1/alpha/search"),
"openai:image" => Some("/v1/images/generations"), "openai:image" => Some("/v1/images/generations"),
"openai:video" => Some("/v1/videos"), "openai:video" => Some("/v1/videos"),
_ => None, _ => None,
@@ -15,6 +15,7 @@ const AI_POST_ROUTE_PATTERNS: &[&str] = &[
"/v1/messages/count_tokens", "/v1/messages/count_tokens",
"/v1/responses", "/v1/responses",
"/v1/responses/compact", "/v1/responses/compact",
"/v1/alpha/search",
"/v1/images/generations", "/v1/images/generations",
"/v1/images/edits", "/v1/images/edits",
"/v1/interactions", "/v1/interactions",
@@ -120,6 +121,7 @@ mod tests {
"/api/v1/services/embeddings/multimodal-embedding/multimodal-embedding", "/api/v1/services/embeddings/multimodal-embedding/multimodal-embedding",
), ),
("openai:rerank", "openai", "rerank", "/v1/rerank"), ("openai:rerank", "openai", "rerank", "/v1/rerank"),
("openai:search", "openai", "search", "/v1/alpha/search"),
("jina:rerank", "jina", "rerank", "/v1/rerank"), ("jina:rerank", "jina", "rerank", "/v1/rerank"),
] { ] {
assert_eq!( assert_eq!(
+88
View File
@@ -92,6 +92,7 @@ pub(crate) struct CandidatePageCacheKey {
preselection_mode: &'static str, preselection_mode: &'static str,
use_api_format_alias_match: bool, use_api_format_alias_match: bool,
client_session_affinity_hash: String, client_session_affinity_hash: String,
model_directive_policy_hash: String,
} }
#[derive(Debug, Clone, PartialEq, Eq, Hash)] #[derive(Debug, Clone, PartialEq, Eq, Hash)]
@@ -140,6 +141,7 @@ impl CandidatePageCacheKey {
preselection_mode: &'static str, preselection_mode: &'static str,
use_api_format_alias_match: bool, use_api_format_alias_match: bool,
client_session_affinity: Option<&ClientSessionAffinity>, client_session_affinity: Option<&ClientSessionAffinity>,
model_directive_policy_hash: &str,
) -> Self { ) -> Self {
Self { Self {
requested_model: normalize_text_key(requested_model), requested_model: normalize_text_key(requested_model),
@@ -153,6 +155,7 @@ impl CandidatePageCacheKey {
preselection_mode, preselection_mode,
use_api_format_alias_match, use_api_format_alias_match,
client_session_affinity_hash: client_session_affinity_key(client_session_affinity), 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, preselection_mode: &'static str,
use_api_format_alias_match: bool, use_api_format_alias_match: bool,
client_session_affinity: Option<&ClientSessionAffinity>, client_session_affinity: Option<&ClientSessionAffinity>,
model_directive_policy_hash: &str,
resolution_mode: AiCandidateResolutionMode, resolution_mode: AiCandidateResolutionMode,
) -> Self { ) -> Self {
Self { Self {
@@ -186,6 +190,7 @@ impl CandidateResolvedPageCacheKey {
preselection_mode, preselection_mode,
use_api_format_alias_match, use_api_format_alias_match,
client_session_affinity, client_session_affinity,
model_directive_policy_hash,
), ),
resolution_mode: resolution_mode_name(resolution_mode), resolution_mode: resolution_mode_name(resolution_mode),
} }
@@ -569,6 +574,7 @@ mod tests {
"provider_endpoint_key_model", "provider_endpoint_key_model",
true, true,
None, None,
"policy-a",
); );
let different_user = CandidatePageCacheKey::new( let different_user = CandidatePageCacheKey::new(
"gpt-4o", "gpt-4o",
@@ -582,6 +588,7 @@ mod tests {
"provider_endpoint_key_model", "provider_endpoint_key_model",
true, true,
None, None,
"policy-a",
); );
let different_model = CandidatePageCacheKey::new( let different_model = CandidatePageCacheKey::new(
"gpt-4.1", "gpt-4.1",
@@ -595,6 +602,7 @@ mod tests {
"provider_endpoint_key_model", "provider_endpoint_key_model",
true, true,
None, None,
"policy-a",
); );
let different_format = CandidatePageCacheKey::new( let different_format = CandidatePageCacheKey::new(
"gpt-4o", "gpt-4o",
@@ -608,6 +616,7 @@ mod tests {
"provider_endpoint_key_model", "provider_endpoint_key_model",
true, true,
None, None,
"policy-a",
); );
let different_capabilities = CandidatePageCacheKey::new( let different_capabilities = CandidatePageCacheKey::new(
"gpt-4o", "gpt-4o",
@@ -621,11 +630,90 @@ mod tests {
"provider_endpoint_key_model", "provider_endpoint_key_model",
true, true,
None, 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_user);
assert_ne!(base, different_model); assert_ne!(base, different_model);
assert_ne!(base, different_format); assert_ne!(base, different_format);
assert_ne!(base, different_capabilities); 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);
} }
} }
@@ -101,6 +101,29 @@ pub(crate) fn client_session_affinity_from_request(
client_session_scope_from_request(headers, body_json)?.scheduler_affinity() 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<ClientSessionAffinity> {
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<ClientSessionScope> {
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( pub(crate) fn client_session_scope_from_request(
headers: &http::HeaderMap, headers: &http::HeaderMap,
body_json: Option<&Value>, 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())) .or_else(|| extract_scope_from_other_specific_adapters(&request, client_family.as_str()))
} }
fn codex_search_session_scope(request: &ClientSessionRequest<'_>) -> Option<ClientSessionScope> {
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( pub(crate) fn client_session_affinity_from_parts(
parts: &http::request::Parts, parts: &http::request::Parts,
body_json: Option<&Value>, body_json: Option<&Value>,
@@ -307,7 +346,9 @@ impl ClientSessionScopeAdapter for GenericSessionScopeAdapter {
} }
fn extract_scope(&self, request: &ClientSessionRequest<'_>) -> Option<ClientSessionScope> { fn extract_scope(&self, request: &ClientSessionRequest<'_>) -> Option<ClientSessionScope> {
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")) .or_else(|| header_value_str(request.headers, "conversation_id"))
{ {
return Some(ClientSessionScope::new( return Some(ClientSessionScope::new(
@@ -367,7 +408,9 @@ impl ClientSessionScopeAdapter for CodexSessionScopeAdapter {
} }
fn extract_scope(&self, request: &ClientSessionRequest<'_>) -> Option<ClientSessionScope> { fn extract_scope(&self, request: &ClientSessionRequest<'_>) -> Option<ClientSessionScope> {
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")) .or_else(|| header_value_str(request.headers, "conversation_id"))
.map(|root_session| { .map(|root_session| {
ClientSessionScope::new( ClientSessionScope::new(
@@ -785,6 +828,7 @@ fn has_header_with_prefix(headers: &http::HeaderMap, prefix: &str) -> bool {
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::{ use super::{
client_session_affinity_from_api_request,
client_session_affinity_from_report_context_value, client_session_affinity_from_request, client_session_affinity_from_report_context_value, client_session_affinity_from_request,
client_session_affinity_report_context_value, client_session_scope_from_request, client_session_affinity_report_context_value, client_session_scope_from_request,
ClientSessionSignalSource, AETHER_AGENT_ID_HEADER, AETHER_SESSION_ID_HEADER, ClientSessionSignalSource, AETHER_AGENT_ID_HEADER, AETHER_SESSION_ID_HEADER,
@@ -841,9 +885,10 @@ mod tests {
let mut headers = HeaderMap::new(); let mut headers = HeaderMap::new();
headers.insert( headers.insert(
http::header::USER_AGENT, 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 = let affinity =
client_session_affinity_from_request(&headers, None).expect("affinity should build"); client_session_affinity_from_request(&headers, None).expect("affinity should build");
@@ -860,7 +905,7 @@ mod tests {
let mut headers = HeaderMap::new(); let mut headers = HeaderMap::new();
headers.insert( headers.insert(
http::header::USER_AGENT, http::header::USER_AGENT,
HeaderValue::from_static("codex-tui/0.122.0"), HeaderValue::from_static("codex_cli_rs/0.144.1"),
); );
headers.insert( headers.insert(
"x-client-request-id", "x-client-request-id",
@@ -1120,4 +1165,28 @@ mod tests {
assert!(client_session_affinity_from_request(&headers, Some(&body)).is_none()); 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());
}
} }
+1
View File
@@ -125,6 +125,7 @@ pub(crate) const RUST_FRONTDOOR_OWNED_ROUTE_PATTERNS: &[&str] = &[
"/v1/messages/count_tokens", "/v1/messages/count_tokens",
"/v1/responses", "/v1/responses",
"/v1/responses/compact", "/v1/responses/compact",
"/v1/alpha/search",
"/v1/models/{model}:generateContent", "/v1/models/{model}:generateContent",
"/v1/models/{model}:streamGenerateContent", "/v1/models/{model}:streamGenerateContent",
"/v1/models/{model}:predictLongRunning", "/v1/models/{model}:predictLongRunning",
File diff suppressed because it is too large Load Diff
+2 -2
View File
@@ -6,8 +6,8 @@ mod types;
pub(crate) use credentials::extract_requested_model; pub(crate) use credentials::extract_requested_model;
pub(crate) use gate::{ pub(crate) use gate::{
request_model_local_rejection, should_buffer_request_for_local_auth, execution_plan_balance_capacity_rejection, request_model_local_rejection,
trusted_auth_local_rejection, GatewayLocalAuthRejection, should_buffer_request_for_local_auth, trusted_auth_local_rejection, GatewayLocalAuthRejection,
}; };
pub(crate) use resolution::{ pub(crate) use resolution::{
refresh_execution_runtime_auth_context, resolve_execution_runtime_auth_context, refresh_execution_runtime_auth_context, resolve_execution_runtime_auth_context,
@@ -1115,10 +1115,9 @@ fn normalize_api_format_alias(value: &str) -> String {
fn auth_gate_api_format(auth_endpoint_signature: &str) -> String { fn auth_gate_api_format(auth_endpoint_signature: &str) -> String {
let normalized = normalize_api_format_alias(auth_endpoint_signature); let normalized = normalize_api_format_alias(auth_endpoint_signature);
if normalized == "antigravity:v1internal" { match normalized.as_str() {
"gemini:generate_content".to_string() "antigravity:v1internal" => "gemini:generate_content".to_string(),
} else { _ => normalized,
normalized
} }
} }
+2 -1
View File
@@ -8,7 +8,8 @@ mod public;
mod route; mod route;
pub(crate) use auth::{ 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, resolve_execution_runtime_auth_context, should_buffer_request_for_local_auth,
trusted_auth_local_rejection, GatewayAdminPrincipalContext, GatewayControlAuthContext, trusted_auth_local_rejection, GatewayAdminPrincipalContext, GatewayControlAuthContext,
GatewayLocalAuthRejection, GatewayLocalAuthRejection,
@@ -54,6 +54,14 @@ pub(super) fn classify_ai_public_route(
true, 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 } else if method == http::Method::POST
&& matches!( && matches!(
normalized_path, normalized_path,
@@ -25,6 +25,7 @@ pub(crate) struct GatewayControlDecision {
pub(crate) auth_context: Option<GatewayControlAuthContext>, pub(crate) auth_context: Option<GatewayControlAuthContext>,
pub(crate) admin_principal: Option<GatewayAdminPrincipalContext>, pub(crate) admin_principal: Option<GatewayAdminPrincipalContext>,
pub(crate) local_auth_rejection: Option<GatewayLocalAuthRejection>, pub(crate) local_auth_rejection: Option<GatewayLocalAuthRejection>,
pub(crate) model_directive_policy: crate::system_features::ModelDirectivePolicySnapshot,
} }
impl GatewayControlDecision { impl GatewayControlDecision {
@@ -47,6 +48,7 @@ impl GatewayControlDecision {
auth_context: None, auth_context: None,
admin_principal: None, admin_principal: None,
local_auth_rejection: None, local_auth_rejection: None,
model_directive_policy: Default::default(),
} }
} }
@@ -131,6 +133,7 @@ impl ClassifiedRoute {
auth_context: None, auth_context: None,
admin_principal: None, admin_principal: None,
local_auth_rejection: None, local_auth_rejection: None,
model_directive_policy: Default::default(),
} }
} }
} }
@@ -146,6 +149,10 @@ pub(crate) async fn resolve_control_route(
return Ok(None); return Ok(None);
}; };
decision.public_query_string = uri.query().map(ToOwned::to_owned); 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? { match resolve_control_decision_auth(state, headers, uri, trace_id, decision).await? {
ControlDecisionAuthResolution::Resolved(decision) => Ok(Some(decision)), 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(); 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") { if uri.path().starts_with("/v1beta/models") {
return "gemini:generate_content".to_string(); return "gemini:generate_content".to_string();
} }
@@ -56,6 +56,28 @@ fn classifies_openai_rerank_as_rerank_not_chat() {
assert!(decision.is_execution_runtime_candidate()); 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] #[test]
fn classifies_openai_chat_and_responses_separately_from_embedding() { fn classifies_openai_chat_and_responses_separately_from_embedding() {
let headers = headers(&[("authorization", "Bearer sk-test")]); let headers = headers(&[("authorization", "Bearer sk-test")]);
@@ -20,6 +20,39 @@ fn classifies_models_list_as_public_support_route() {
assert!(!decision.is_execution_runtime_candidate()); 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] #[test]
fn classifies_v1beta_models_as_gemini_public_support_route() { fn classifies_v1beta_models_as_gemini_public_support_route() {
let headers = headers(&[]); let headers = headers(&[]);
+220 -70
View File
@@ -12,10 +12,7 @@ use super::{
StoredWalletSnapshot, UpdateManagementTokenRecord, UpsertOAuthProviderConfigRecord, StoredWalletSnapshot, UpdateManagementTokenRecord, UpsertOAuthProviderConfigRecord,
}; };
use crate::LocalMutationOutcome; use crate::LocalMutationOutcome;
use aether_data::repository::auth::{ use aether_data::repository::auth::ResolvedAuthApiKeySnapshotReader;
read_resolved_auth_api_key_snapshot_by_key_hash,
read_resolved_auth_api_key_snapshot_by_user_api_key_ids,
};
#[derive(Debug, Clone, Default)] #[derive(Debug, Clone, Default)]
pub(crate) struct GatewayUserEffectiveListPolicies { pub(crate) struct GatewayUserEffectiveListPolicies {
@@ -1751,15 +1748,14 @@ impl GatewayDataState {
let snapshot = crate::request_diagnostics::observe_db_operation( let snapshot = crate::request_diagnostics::observe_db_operation(
"auth_api_key_snapshot", "auth_api_key_snapshot",
self.database_pool_summary(), self.database_pool_summary(),
read_resolved_auth_api_key_snapshot_by_user_api_key_ids( self.find_stored_auth_api_key_snapshot(AuthApiKeyLookupKey::UserApiKeyIds {
self,
user_id, user_id,
api_key_id, api_key_id,
now_unix_secs, }),
),
) )
.await?; .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( 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( let snapshot = crate::request_diagnostics::observe_db_operation(
"auth_api_key_snapshot_by_hash", "auth_api_key_snapshot_by_hash",
self.database_pool_summary(), 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?; .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( async fn apply_user_group_effective_policies(
&self, &self,
snapshot: Option<GatewayAuthApiKeySnapshot>, snapshot: Option<StoredAuthApiKeySnapshot>,
now_unix_secs: u64,
) -> Result<Option<GatewayAuthApiKeySnapshot>, DataLayerError> { ) -> Result<Option<GatewayAuthApiKeySnapshot>, DataLayerError> {
let Some(mut snapshot) = snapshot else { let Some(mut snapshot) = snapshot else {
return Ok(None); return Ok(None);
}; };
if snapshot.user_role.eq_ignore_ascii_case("admin") && !snapshot.api_key_is_standalone { if snapshot.user_role.eq_ignore_ascii_case("admin") && !snapshot.api_key_is_standalone {
apply_admin_unrestricted_auth_snapshot(&mut snapshot); 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 { 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( let Some(user) = crate::request_diagnostics::observe_db_operation(
"auth_user_policy", "auth_user_policy",
@@ -1797,57 +1801,50 @@ impl GatewayDataState {
) )
.await? .await?
else { 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 { if user.role.eq_ignore_ascii_case("admin") && !snapshot.api_key_is_standalone {
snapshot.user_role = user.role; snapshot.user_role = user.role;
apply_admin_unrestricted_auth_snapshot(&mut snapshot); apply_admin_unrestricted_auth_snapshot(&mut snapshot);
return Ok(Some(snapshot)); return Ok(Some(GatewayAuthApiKeySnapshot::from_stored(
snapshot,
now_unix_secs,
)));
} }
let groups = self let groups = self
.effective_user_groups_for_user(&snapshot.user_id) .effective_user_groups_for_user(&snapshot.user_id)
.await?; .await?;
let mut allowed_providers = let allowed_providers =
resolve_effective_list_policy(None, "unrestricted", &groups, |group| { resolve_effective_list_policy(None, "unrestricted", &groups, |group| {
( (
&group.allowed_providers_mode, &group.allowed_providers_mode,
group.allowed_providers.clone(), group.allowed_providers.clone(),
) )
}); });
let mut allowed_api_formats = let allowed_api_formats =
resolve_effective_list_policy(None, "unrestricted", &groups, |group| { resolve_effective_api_format_policy(None, "unrestricted", &groups, |group| {
( (
&group.allowed_api_formats_mode, &group.allowed_api_formats_mode,
group.allowed_api_formats.clone(), group.allowed_api_formats.clone(),
) )
}); });
let mut allowed_models = let allowed_models =
resolve_effective_list_policy(None, "unrestricted", &groups, |group| { resolve_effective_list_policy(None, "unrestricted", &groups, |group| {
(&group.allowed_models_mode, group.allowed_models.clone()) (&group.allowed_models_mode, group.allowed_models.clone())
}); });
let user_rate_limit = resolve_effective_rate_limit_policy(None, "system", &groups); let user_rate_limit = resolve_effective_rate_limit_policy(None, "system", &groups);
if !snapshot.api_key_is_standalone { snapshot.user_allowed_providers = allowed_providers;
constrain_api_key_list_policy_to_user_policy( snapshot.user_allowed_api_formats = allowed_api_formats;
&mut allowed_providers, snapshot.user_allowed_models = allowed_models;
&mut snapshot.api_key_allowed_providers, snapshot.user_rate_limit = user_rate_limit;
); Ok(Some(GatewayAuthApiKeySnapshot::from_stored(
constrain_api_key_list_policy_to_user_policy( snapshot,
&mut allowed_api_formats, now_unix_secs,
&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))
} }
pub(crate) async fn resolve_user_effective_list_policies( 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.clone(),
&user.allowed_api_formats_mode, &user.allowed_api_formats_mode,
&groups, &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_providers = None;
snapshot.user_allowed_api_formats = None; snapshot.user_allowed_api_formats = None;
snapshot.user_allowed_models = None; snapshot.user_allowed_models = None;
@@ -1995,6 +1992,19 @@ fn resolve_effective_list_policy(
intersect_list_policies(group_policy, user_policy) intersect_list_policies(group_policy, user_policy)
} }
fn resolve_effective_api_format_policy(
user_values: Option<Vec<String>>,
user_mode: &str,
groups: &[aether_data::repository::users::StoredUserGroup],
group_field: impl Fn(
&aether_data::repository::users::StoredUserGroup,
) -> (&str, Option<Vec<String>>),
) -> Option<Vec<String>> {
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( fn union_group_list_policies(
groups: &[aether_data::repository::users::StoredUserGroup], groups: &[aether_data::repository::users::StoredUserGroup],
group_field: impl Fn( group_field: impl Fn(
@@ -2089,6 +2099,19 @@ fn intersect_list_policies(
} }
} }
fn intersect_api_format_list_policies(
left: Option<Vec<String>>,
right: Option<Vec<String>>,
) -> Option<Vec<String>> {
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( fn intersect_rate_limit_policies(
left: Option<RateLimitRestriction>, left: Option<RateLimitRestriction>,
right: Option<RateLimitRestriction>, right: Option<RateLimitRestriction>,
@@ -2133,22 +2156,6 @@ fn rate_limit_policy_value(policy: Option<RateLimitRestriction>) -> Option<i32>
} }
} }
fn constrain_api_key_list_policy_to_user_policy(
user_policy: &mut Option<Vec<String>>,
api_key_policy: &mut Option<Vec<String>>,
) {
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)] #[cfg(test)]
mod tests { mod tests {
use std::sync::Arc; 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] #[test]
fn list_policy_unions_multiple_group_restrictions_legacy_case() { fn list_policy_unions_multiple_group_restrictions_legacy_case() {
let groups = vec![ 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::<String>::new()));
assert_eq!(api_key_policy, Some(Vec::<String>::new()));
}
#[tokio::test] #[tokio::test]
async fn admin_non_standalone_snapshot_bypasses_group_and_key_policies() { async fn admin_non_standalone_snapshot_bypasses_group_and_key_policies() {
let mut snapshot = sample_snapshot_with_role("key-admin", "admin-1", "admin") 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); 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] #[tokio::test]
async fn user_personal_policy_fields_are_ignored_when_groups_are_applied() { 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)); 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)); 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] #[tokio::test]
async fn data_state_lists_auth_api_key_export_records() { async fn data_state_lists_auth_api_key_export_records() {
let repository = Arc::new( let repository = Arc::new(
@@ -532,6 +532,88 @@ impl GatewayDataState {
Ok(updated) 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<u64>,
) -> Result<bool, DataLayerError> {
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<u64>,
last_models_fetch_error: Option<&str>,
updated_at_unix_secs: Option<u64>,
) -> Result<bool, DataLayerError> {
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<u64>,
) -> Result<bool, DataLayerError> {
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( pub(crate) async fn delete_provider_catalog_key(
&self, &self,
key_id: &str, key_id: &str,
@@ -4172,6 +4172,7 @@ mod tests {
global_model_id: "global-model-1".to_string(), global_model_id: "global-model-1".to_string(),
global_model_name: "gpt-5".to_string(), global_model_name: "gpt-5".to_string(),
selected_provider_model_name: "gpt-5".to_string(), selected_provider_model_name: "gpt-5".to_string(),
supports_streaming: true,
mapping_matched_model: None, mapping_matched_model: None,
}, },
provider_api_format: "openai:responses".to_string(), provider_api_format: "openai:responses".to_string(),
@@ -4285,6 +4286,7 @@ mod tests {
global_model_id: "global-model-1".to_string(), global_model_id: "global-model-1".to_string(),
global_model_name: "gpt-5".to_string(), global_model_name: "gpt-5".to_string(),
selected_provider_model_name: "gpt-5".to_string(), selected_provider_model_name: "gpt-5".to_string(),
supports_streaming: true,
mapping_matched_model: None, mapping_matched_model: None,
}, },
provider_api_format: "openai:chat".to_string(), provider_api_format: "openai:chat".to_string(),
+1
View File
@@ -146,6 +146,7 @@ mod tests {
global_model_id: "global-model-1".to_string(), global_model_id: "global-model-1".to_string(),
global_model_name: "gpt-5".to_string(), global_model_name: "gpt-5".to_string(),
selected_provider_model_name: "gpt-5".to_string(), selected_provider_model_name: "gpt-5".to_string(),
supports_streaming: true,
mapping_matched_model: None, mapping_matched_model: None,
}, },
transport: Arc::new(crate::ai_serving::GatewayProviderTransportSnapshot { transport: Arc::new(crate::ai_serving::GatewayProviderTransportSnapshot {
@@ -941,6 +941,7 @@ mod tests {
success_failover_patterns: Vec::new(), success_failover_patterns: Vec::new(),
error_stop_patterns: Vec::new(), error_stop_patterns: Vec::new(),
stop_cyber_policy_errors: false, stop_cyber_policy_errors: false,
retry_client_errors_by_default: true,
} }
); );
} }
@@ -1042,6 +1042,7 @@ fn grok_stream_terminal_summary(
finish_reason: Some("stop".to_string()), finish_reason: Some("stop".to_string()),
response_id: None, response_id: None,
model: plan.model_name.clone(), model: plan.model_name.clone(),
provider_actual_service_tier: None,
observed_finish: true, observed_finish: true,
unknown_event_count: 0, unknown_event_count: 0,
parser_error: None, parser_error: None,
@@ -359,8 +359,7 @@ impl IntoResponse for ExecutionRuntimeAppError {
return build_overloaded_response(&self.0.to_string()); return build_overloaded_response(&self.0.to_string());
} }
ExecutionRuntimeServerError::Transport( ExecutionRuntimeServerError::Transport(
ExecutionRuntimeTransportError::StreamUnsupported ExecutionRuntimeTransportError::RequestBodyRequired
| ExecutionRuntimeTransportError::RequestBodyRequired
| ExecutionRuntimeTransportError::BodyDecode(_) | ExecutionRuntimeTransportError::BodyDecode(_)
| ExecutionRuntimeTransportError::UnsupportedContentEncoding(_) | ExecutionRuntimeTransportError::UnsupportedContentEncoding(_)
| ExecutionRuntimeTransportError::ProxyUnsupported | ExecutionRuntimeTransportError::ProxyUnsupported
@@ -397,7 +396,9 @@ mod tests {
build_execution_runtime_router_with_request_concurrency_limit, build_execution_runtime_router_with_request_concurrency_limit,
build_execution_runtime_router_with_request_gates, DISTRIBUTED_REQUEST_GATE_NAME, 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::{ use aether_runtime_state::{
MemoryRuntimeStateConfig, RuntimeSemaphore, RuntimeSemaphoreConfig, RuntimeState, 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::<StreamFrame>(line)
.expect("execution runtime frame should decode")
.frame_type
})
.collect::<Vec<_>>();
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] #[tokio::test]
async fn execution_runtime_rejects_second_in_flight_stream_request_with_overload() { async fn execution_runtime_rejects_second_in_flight_stream_request_with_overload() {
let upstream_hits = Arc::new(AtomicUsize::new(0)); let upstream_hits = Arc::new(AtomicUsize::new(0));
@@ -851,6 +851,9 @@ fn merge_stream_terminal_summary(
if current_summary.model.is_none() { if current_summary.model.is_none() {
current_summary.model = observed.model; 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.observed_finish |= observed.observed_finish;
current_summary.unknown_event_count = current_summary current_summary.unknown_event_count = current_summary
.unknown_event_count .unknown_event_count
@@ -5484,9 +5487,8 @@ async fn execute_stream_from_frame_stream(
telemetry = Some(frame_telemetry); telemetry = Some(frame_telemetry);
} }
StreamFramePayload::Eof { summary } => { StreamFramePayload::Eof { summary } => {
if summary.is_some() { stream_terminal_summary =
stream_terminal_summary = summary; merge_stream_terminal_summary(stream_terminal_summary.take(), summary);
}
break; break;
} }
StreamFramePayload::Error { error } => { StreamFramePayload::Error { error } => {
@@ -6293,12 +6295,14 @@ mod tests {
Some(ExecutionStreamTerminalSummary { Some(ExecutionStreamTerminalSummary {
standardized_usage: Some(runtime_usage), standardized_usage: Some(runtime_usage),
model: Some("gpt-5.5".to_string()), model: Some("gpt-5.5".to_string()),
provider_actual_service_tier: Some("priority".to_string()),
unknown_event_count: 1, unknown_event_count: 1,
..ExecutionStreamTerminalSummary::default() ..ExecutionStreamTerminalSummary::default()
}), }),
Some(ExecutionStreamTerminalSummary { Some(ExecutionStreamTerminalSummary {
standardized_usage: Some(observed_usage), standardized_usage: Some(observed_usage),
response_id: Some("resp_123".to_string()), response_id: Some("resp_123".to_string()),
provider_actual_service_tier: Some("default".to_string()),
observed_finish: true, observed_finish: true,
unknown_event_count: 2, unknown_event_count: 2,
..ExecutionStreamTerminalSummary::default() ..ExecutionStreamTerminalSummary::default()
@@ -6313,6 +6317,10 @@ mod tests {
assert_eq!(usage.output_tokens, 137); assert_eq!(usage.output_tokens, 137);
assert_eq!(merged.model.as_deref(), Some("gpt-5.5")); assert_eq!(merged.model.as_deref(), Some("gpt-5.5"));
assert_eq!(merged.response_id.as_deref(), Some("resp_123")); 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!(merged.observed_finish);
assert_eq!(merged.unknown_event_count, 3); assert_eq!(merged.unknown_event_count, 3);
} }
@@ -676,6 +676,14 @@ fn should_buffer_non_stream_response(
return false; return false;
} }
if report_context
.get("upstream_is_stream")
.and_then(Value::as_bool)
== Some(false)
{
return true;
}
headers headers
.get("content-length") .get("content-length")
.and_then(|value| value.trim().parse::<u64>().ok()) .and_then(|value| value.trim().parse::<u64>().ok())
@@ -1134,22 +1142,36 @@ mod tests {
} }
#[test] #[test]
fn buffers_non_sse_response_only_when_content_length_is_known() { fn buffers_declared_non_stream_responses_without_relying_on_content_length() {
let report_context = serde_json::json!({ let streaming_context = serde_json::json!({
"provider_api_format": "openai:chat", "provider_api_format": "openai:chat",
"client_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( assert!(!should_buffer_non_stream_response(
&BTreeMap::from([("content-type".into(), "application/json".into())]), &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( assert!(should_buffer_non_stream_response(
&BTreeMap::from([ &BTreeMap::from([
("content-type".into(), "application/json".into()), ("content-type".into(), "application/json".into()),
("content-length".into(), "128".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 .await
.expect("listener should bind"); .expect("listener should bind");
let addr = listener.local_addr().expect("local addr should resolve"); 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 server = tokio::spawn(async move {
let app = Router::new().route( let app = Router::new().route(
"/responses", "/images/generations",
post(|| async { post(move || {
let body = serde_json::json!({ let generated_image = generated_image.clone();
"created": 1776971267_u64, async move {
"data": [{ let body = serde_json::json!({
"b64_json": "aGVsbG8=" "created": 1776971267_u64,
}], "data": [{
"usage": { "b64_json": generated_image
"total_tokens": 100, }],
"input_tokens": 50, "usage": {
"output_tokens": 50, "total_tokens": 100,
"input_tokens_details": { "input_tokens": 50,
"text_tokens": 10, "output_tokens": 50,
"image_tokens": 40 "input_tokens_details": {
"text_tokens": 10,
"image_tokens": 40
}
} }
} });
}); let encoded = serde_json::to_vec(&body).expect("json should encode");
let mut response = axum::http::Response::new(Body::from( let chunks = encoded
serde_json::to_vec(&body).expect("json should encode"), .chunks(4096)
)); .map(Bytes::copy_from_slice)
response.headers_mut().insert( .collect::<Vec<_>>();
header::CONTENT_TYPE, let chunked_body = stream! {
HeaderValue::from_static("application/json"), for chunk in chunks {
); yield Ok::<Bytes, Infallible>(chunk);
response 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) axum::serve(listener, app)
@@ -1611,16 +1648,15 @@ mod tests {
endpoint_id: "endpoint-1".to_string(), endpoint_id: "endpoint-1".to_string(),
key_id: "key-1".to_string(), key_id: "key-1".to_string(),
method: "POST".to_string(), method: "POST".to_string(),
url: format!("http://{addr}/responses"), url: format!("http://{addr}/images/generations"),
headers: BTreeMap::new(), headers: BTreeMap::new(),
content_type: None, content_type: None,
content_encoding: None, content_encoding: None,
body: RequestBody::from_json(serde_json::json!({ body: RequestBody::from_json(serde_json::json!({
"model": "gpt-image-1", "model": "gpt-image-1",
"prompt": "poster", "prompt": "poster"
"stream": true
})), })),
stream: true, stream: false,
client_api_format: "openai:image".to_string(), client_api_format: "openai:image".to_string(),
provider_api_format: "openai:image".to_string(), provider_api_format: "openai:image".to_string(),
model_name: Some("gpt-image-1".into()), 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"); 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("event: image_generation.completed"));
assert!(bridged_text.contains("\"type\":\"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")); assert!(bridged_text.contains("\"total_tokens\":100"));
let eof_frame = frames let eof_frame = frames
@@ -43,6 +43,7 @@ fn missing_exact_provider_request_payload(decision_kind: &str) -> AiExecutionDec
request_id: Some("req_123".to_string()), request_id: Some("req_123".to_string()),
candidate_id: Some("cand_123".to_string()), candidate_id: Some("cand_123".to_string()),
provider_name: Some("provider".to_string()), provider_name: Some("provider".to_string()),
provider_type: None,
provider_id: Some("provider_id".to_string()), provider_id: Some("provider_id".to_string()),
endpoint_id: Some("endpoint_id".to_string()), endpoint_id: Some("endpoint_id".to_string()),
key_id: Some("key_id".to_string()), key_id: Some("key_id".to_string()),
@@ -57,8 +57,8 @@ const TUNNEL_RELAY_PATH_PREFIX: &str = "/api/internal/tunnel/relay";
const DEFAULT_TUNNEL_TIMEOUT_MS: u64 = 60_000; const DEFAULT_TUNNEL_TIMEOUT_MS: u64 = 60_000;
const DEFAULT_STREAM_FIRST_BYTE_TIMEOUT_MS: u64 = 30_000; const DEFAULT_STREAM_FIRST_BYTE_TIMEOUT_MS: u64 = 30_000;
const DEFAULT_NON_STREAM_TOTAL_TIMEOUT_MS: u64 = 300_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 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_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_CLIENT_SHARDS_ENV: &str = "AETHER_GATEWAY_DIRECT_REQWEST_CLIENT_SHARDS";
const DIRECT_REQWEST_H2_TARGET_STREAMS_PER_CLIENT_ENV: &str = 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)] #[derive(Debug, Error)]
pub(crate) enum ExecutionRuntimeTransportError { 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")] #[error("request body must contain json_body or body_bytes_b64")]
RequestBodyRequired, RequestBodyRequired,
#[error("request body base64 is invalid: {0}")] #[error("request body base64 is invalid: {0}")]
@@ -681,10 +679,6 @@ impl DirectSyncExecutionRuntime {
&self, &self,
plan: &ExecutionPlan, plan: &ExecutionPlan,
) -> Result<DirectUpstreamStreamExecution, ExecutionRuntimeTransportError> { ) -> Result<DirectUpstreamStreamExecution, ExecutionRuntimeTransportError> {
if !plan.stream {
return Err(ExecutionRuntimeTransportError::StreamUnsupported);
}
let build_body_started_at = Instant::now(); let build_body_started_at = Instant::now();
let body_bytes = build_request_body(plan)?; let body_bytes = build_request_body(plan)?;
observe_gateway_stage_ms( observe_gateway_stage_ms(
@@ -835,6 +829,7 @@ fn build_stream_summary_report_context(plan: &ExecutionPlan) -> Value {
"provider_api_format": plan.provider_api_format, "provider_api_format": plan.provider_api_format,
"client_api_format": plan.client_api_format, "client_api_format": plan.client_api_format,
"model": plan.model_name, "model": plan.model_name,
"upstream_is_stream": plan.stream,
}) })
} }
@@ -2125,20 +2120,17 @@ pub(crate) fn build_request_body(
Vec::new() Vec::new()
}; };
if should_gzip_request_body(plan) && plan.body.json_body.is_some() { if plan.body.json_body.is_some() {
body_bytes = gzip_bytes(&body_bytes)?; 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) 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<String> { fn normalize_content_encoding(value: Option<&str>) -> Option<String> {
value value
.map(str::trim) .map(str::trim)
@@ -2156,6 +2148,11 @@ fn gzip_bytes(body_bytes: &[u8]) -> Result<Vec<u8>, ExecutionRuntimeTransportErr
.map_err(|err| ExecutionRuntimeTransportError::RelayError(err.to_string())) .map_err(|err| ExecutionRuntimeTransportError::RelayError(err.to_string()))
} }
fn zstd_bytes(body_bytes: &[u8]) -> Result<Vec<u8>, ExecutionRuntimeTransportError> {
zstd::stream::encode_all(std::io::Cursor::new(body_bytes), 3)
.map_err(|err| ExecutionRuntimeTransportError::RelayError(err.to_string()))
}
fn build_relay_client( fn build_relay_client(
timeouts: Option<&aether_contracts::ExecutionTimeouts>, timeouts: Option<&aether_contracts::ExecutionTimeouts>,
) -> Result<reqwest::Client, ExecutionRuntimeTransportError> { ) -> Result<reqwest::Client, ExecutionRuntimeTransportError> {
@@ -2218,28 +2215,49 @@ fn resolve_tunnel_first_byte_timeout(plan: &ExecutionPlan) -> Option<Duration> {
}) })
} }
fn resolve_non_stream_total_timeout(plan: &ExecutionPlan) -> Option<Duration> { pub(crate) fn resolve_non_stream_total_timeout_for_request(
if plan.stream { is_stream: bool,
provider_api_format: &str,
timeouts: Option<&aether_contracts::ExecutionTimeouts>,
) -> Option<Duration> {
if is_stream {
return None; return None;
} }
let timeout_ms = plan let default_timeout_ms =
.timeouts if crate::ai_serving::is_openai_responses_compact_format(provider_api_format) {
.as_ref() DEFAULT_CODEX_COMPACT_TOTAL_TIMEOUT_MS
} else {
DEFAULT_NON_STREAM_TOTAL_TIMEOUT_MS
};
let timeout_ms = timeouts
.and_then(|timeouts| timeouts.total_ms) .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<Duration> {
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<Duration> {
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))) Some(Duration::from_millis(timeout_ms.max(1)))
} }
pub(crate) fn resolve_stream_first_byte_timeout(plan: &ExecutionPlan) -> Option<Duration> { pub(crate) fn resolve_stream_first_byte_timeout(plan: &ExecutionPlan) -> Option<Duration> {
if !plan.stream { resolve_stream_first_byte_timeout_for_request(plan.stream, plan.timeouts.as_ref())
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)))
} }
pub(crate) async fn with_non_stream_total_timeout<T, F>( pub(crate) async fn with_non_stream_total_timeout<T, F>(
@@ -2359,7 +2377,10 @@ fn resolve_tunnel_timeout_metadata(plan: &ExecutionPlan) -> TunnelTimeoutMetadat
fn timeout_ms_to_secs(ms: u64) -> u64 { fn timeout_ms_to_secs(ms: u64) -> u64 {
let secs = ms.div_ceil(1_000); 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<String> { fn resolve_tunnel_node_id(proxy: Option<&ProxySnapshot>) -> Option<String> {
@@ -3552,7 +3573,7 @@ pub(crate) fn build_request_headers(
let mut out = HeaderMap::new(); let mut out = HeaderMap::new();
let normalized_content_encoding = normalize_content_encoding(content_encoding); let normalized_content_encoding = normalize_content_encoding(content_encoding);
if let Some(encoding) = normalized_content_encoding.as_deref() { 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( return Err(ExecutionRuntimeTransportError::UnsupportedContentEncoding(
encoding.to_string(), encoding.to_string(),
)); ));
@@ -4637,6 +4658,25 @@ mod tests {
assert_eq!(timeout, std::time::Duration::from_secs(300)); 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] #[test]
fn tunnel_request_meta_uses_non_stream_default_instead_of_first_byte_default() { fn tunnel_request_meta_uses_non_stream_default_instead_of_first_byte_default() {
let mut plan = tunnel_timeout_plan(false); let mut plan = tunnel_timeout_plan(false);
@@ -6323,13 +6363,21 @@ mod tests {
.and_then(|value| value.to_str().ok()) .and_then(|value| value.to_str().ok())
.unwrap_or_default() .unwrap_or_default()
.to_string(); .to_string();
let mut decoder = flate2::read::GzDecoder::new(body.as_ref()); let decoded = match header_encoding.as_str() {
let mut decoded = String::new(); "gzip" => {
decoder let mut decoder = flate2::read::GzDecoder::new(body.as_ref());
.read_to_string(&mut decoded) let mut decoded = Vec::new();
.expect("gzip body should decode"); 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 = 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, axum::http::StatusCode::OK,
Json(json!({ Json(json!({
@@ -6346,45 +6394,47 @@ mod tests {
}); });
let execution_runtime = DirectSyncExecutionRuntime::new(); let execution_runtime = DirectSyncExecutionRuntime::new();
let result = execution_runtime for encoding in ["gzip", "zstd"] {
.execute_sync(&ExecutionPlan { let result = execution_runtime
request_id: "req-gzip-1".into(), .execute_sync(&ExecutionPlan {
candidate_id: Some("cand-1".into()), request_id: format!("req-{encoding}-1"),
provider_name: Some("openai".into()), candidate_id: Some("cand-1".into()),
provider_id: "prov-1".into(), provider_name: Some("openai".into()),
endpoint_id: "ep-1".into(), provider_id: "prov-1".into(),
key_id: "key-1".into(), endpoint_id: "ep-1".into(),
method: "POST".into(), key_id: "key-1".into(),
url: format!("http://{addr}/chat"), method: "POST".into(),
headers: BTreeMap::from([("content-type".into(), "application/json".into())]), url: format!("http://{addr}/chat"),
content_type: Some("application/json".into()), headers: BTreeMap::from([("content-type".into(), "application/json".into())]),
content_encoding: Some("gzip".into()), content_type: Some("application/json".into()),
body: RequestBody::from_json(json!({"model": "gpt-4.1"})), content_encoding: Some(encoding.into()),
stream: false, body: RequestBody::from_json(json!({"model": "gpt-4.1"})),
client_api_format: "openai:chat".into(), stream: false,
provider_api_format: "openai:chat".into(), client_api_format: "openai:chat".into(),
model_name: Some("gpt-4.1".into()), provider_api_format: "openai:chat".into(),
proxy: None, model_name: Some("gpt-4.1".into()),
transport_profile: None, proxy: None,
timeouts: Some(ExecutionTimeouts { transport_profile: None,
connect_ms: Some(5_000), timeouts: Some(ExecutionTimeouts {
total_ms: Some(LOCAL_HTTP_SUCCESS_TIMEOUT_MS), connect_ms: Some(5_000),
..ExecutionTimeouts::default() total_ms: Some(LOCAL_HTTP_SUCCESS_TIMEOUT_MS),
}), ..ExecutionTimeouts::default()
}) }),
.await })
.expect("gzip sync execution should succeed"); .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(); 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] #[tokio::test]
@@ -166,18 +166,31 @@ where
type Error = GatewayError; type Error = GatewayError;
async fn execute_attempt(&self, attempt: &T) -> Result<Option<Self::Response>, Self::Error> { async fn execute_attempt(&self, attempt: &T) -> Result<Option<Self::Response>, 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 _permit = acquire_upstream_execution_gate(self.state, self.trace_id).await?;
let upstream_execution_gate_held_started_at = std::time::Instant::now(); let upstream_execution_gate_held_started_at = std::time::Instant::now();
let mut response = execute_execution_runtime_sync( let mut response = execute_execution_runtime_sync(
self.state, self.state,
self.parts.uri.path(), self.parts.uri.path(),
attempt.execution_plan().clone(), plan.clone(),
self.trace_id, self.trace_id,
self.decision, self.decision,
self.plan_kind, self.plan_kind,
attempt.report_kind(), attempt.report_kind(),
attempt.report_context(), report_context,
) )
.await?; .await?;
observe_gateway_stage_ms( observe_gateway_stage_ms(
@@ -187,10 +200,7 @@ where
.as_millis() as u64, .as_millis() as u64,
); );
if let Some(response) = response.as_mut() { if let Some(response) = response.as_mut() {
attach_redaction_execution_candidate( attach_redaction_execution_candidate(response, plan.candidate_id.as_deref());
response,
attempt.execution_plan().candidate_id.as_deref(),
);
} }
Ok(response) Ok(response)
} }
@@ -355,7 +365,14 @@ where
}; };
last_attempted = Some((attempt.execution_plan().clone(), attempt.report_context())); last_attempted = Some((attempt.execution_plan().clone(), attempt.report_context()));
let execute_started_at = std::time::Instant::now(); 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( observe_gateway_stage_ms(
"stream_candidate_execute", "stream_candidate_execute",
execute_started_at.elapsed().as_millis() as u64, execute_started_at.elapsed().as_millis() as u64,
@@ -450,6 +467,17 @@ where
candidate_index = candidate_index.as_str(), candidate_index = candidate_index.as_str(),
"candidate loop attempting stream execution candidate" "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); prewarm_direct_reqwest_candidate_client(&plan);
let watchdog_plan = plan.clone(); let watchdog_plan = plan.clone();
let watchdog_report_context = report_context.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<Option<Response<Body>>, 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<T>(state: &AppState, remaining: Vec<T>) pub(crate) async fn mark_unused_local_candidates<T>(state: &AppState, remaining: Vec<T>)
where where
T: AiExecutionAttempt, T: AiExecutionAttempt,
{ {
for plan_and_report in remaining { for plan_and_report in remaining {
let report_context = plan_and_report.report_context(); let report_context = plan_and_report.report_context();
let metadata = mark_unused_local_candidate(
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(
state, state,
plan_and_report.execution_plan(), plan_and_report.execution_plan(),
report_context.as_ref(), 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; .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 { fn should_skip_unused_persistence(report_context: Option<&serde_json::Value>) -> bool {
let metadata = local_execution_candidate_metadata_from_report_context(report_context); let metadata = local_execution_candidate_metadata_from_report_context(report_context);
should_skip_unused_persistence_from_metadata(&metadata) should_skip_unused_persistence_from_metadata(&metadata)
@@ -179,6 +179,7 @@ mod tests {
auth_context: None, auth_context: None,
admin_principal: None, admin_principal: None,
local_auth_rejection: None, local_auth_rejection: None,
model_directive_policy: Default::default(),
} }
} }
@@ -41,6 +41,7 @@ pub(crate) fn frontdoor_self_loop_public_ai_path(path: &str) -> bool {
| "/v1/rerank" | "/v1/rerank"
| "/v1/responses" | "/v1/responses"
| "/v1/responses/compact" | "/v1/responses/compact"
| "/v1/alpha/search"
| "/v1beta/files" | "/v1beta/files"
| "/upload/v1beta/files" | "/upload/v1beta/files"
| "/v1beta/operations" | "/v1beta/operations"
@@ -104,6 +104,7 @@ fn is_known_admin_monitoring_api_format(value: &str) -> bool {
"openai:chat" "openai:chat"
| "openai:responses" | "openai:responses"
| "openai:responses:compact" | "openai:responses:compact"
| "openai:search"
| "openai:image" | "openai:image"
| "openai:video" | "openai:video"
| "openai:embedding" | "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] #[test]
fn parses_v2_scheduler_affinity_key_with_three_segment_api_format() { fn parses_v2_scheduler_affinity_key_with_three_segment_api_format() {
let parsed = parse_admin_monitoring_scheduler_affinity_key( let parsed = parse_admin_monitoring_scheduler_affinity_key(
@@ -6,7 +6,8 @@ use crate::handlers::admin::provider::shared::payloads::{
AdminProviderCreateRequest, AdminProviderUpdatePatch, AdminProviderCreateRequest, AdminProviderUpdatePatch,
}; };
use crate::handlers::admin::provider::write::provider::{ 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::request::{AdminAppState, AdminRequestContext};
use crate::handlers::admin::shared::attach_admin_audit_response; 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() .is_some()
{ {
reconcile_admin_fixed_provider_template_endpoints(state, &created_provider).await?; 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( return Ok(Some(attach_admin_audit_response(
Json(json!({ Json(json!({
@@ -161,8 +161,12 @@ pub(crate) async fn maybe_build_local_admin_provider_writes_response(
.fixed_provider_template(&updated_record.provider_type) .fixed_provider_template(&updated_record.provider_type)
.is_some() .is_some()
{ {
reconcile_admin_fixed_provider_template_endpoints(state, &updated_record).await?; reconcile_admin_fixed_provider_template_endpoints_after_update(
reconcile_admin_fixed_provider_template_keys(state, &updated_record).await?; state,
&existing_provider,
&updated_record,
)
.await?;
} }
return Ok(Some( return Ok(Some(
match state match state
@@ -64,7 +64,8 @@ use aether_data_contracts::repository::provider_catalog::{
}; };
use aether_model_fetch::{ use aether_model_fetch::{
aggregate_models_for_cache, fetch_models_from_transports, json_string_list, 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::{ use axum::{
body::{to_bytes, Body}, body::{to_bytes, Body},
@@ -281,6 +282,7 @@ fn provider_query_attach_model_test_capabilities(
fn provider_query_codex_preset_fallback( fn provider_query_codex_preset_fallback(
provider: &StoredProviderCatalogProvider, provider: &StoredProviderCatalogProvider,
fetch_error: &str,
) -> Option<ProviderQueryKeyFetchResult> { ) -> Option<ProviderQueryKeyFetchResult> {
if !provider.provider_type.trim().eq_ignore_ascii_case("codex") { if !provider.provider_type.trim().eq_ignore_ascii_case("codex") {
return None; return None;
@@ -289,12 +291,59 @@ fn provider_query_codex_preset_fallback(
Some(ProviderQueryKeyFetchResult { Some(ProviderQueryKeyFetchResult {
models: aggregate_models_for_cache(&models), models: aggregate_models_for_cache(&models),
error: None, error: None,
warning: None, warning: Some(format!(
"Codex 动态模型目录不可用,已使用内置模型卡:{fetch_error}"
)),
from_cache: false, from_cache: false,
has_success: true, 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(());
}
<AppState as ModelFetchRuntimeState>::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; mod model_test;
pub(crate) use self::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); let selected_endpoints = selected_models_fetch_endpoints(endpoints, key);
if selected_endpoints.is_empty() { if selected_endpoints.is_empty() {
if let Some(models) = preset_models_for_provider(&provider.provider_type) { if let Some(models) = preset_models_for_provider(&provider.provider_type) {
let models = provider_query_filter_models_for_key( let models = aggregate_models_for_cache(&models);
provider, provider_query_persist_preset_models(state, provider, key, &models).await?;
key, let models = provider_query_filter_models_for_key(provider, key, models);
aggregate_models_for_cache(&models),
);
return Ok(ProviderQueryKeyFetchResult { return Ok(ProviderQueryKeyFetchResult {
models, models,
error: None, error: None,
@@ -492,7 +539,11 @@ async fn provider_query_fetch_models_for_key(
Ok(outcome) => outcome, Ok(outcome) => outcome,
Err(err) => { Err(err) => {
all_errors.push(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(fallback);
} }
return Ok(ProviderQueryKeyFetchResult { return Ok(ProviderQueryKeyFetchResult {
@@ -517,20 +568,14 @@ async fn provider_query_fetch_models_for_key(
.await; .await;
} }
if let Some(upstream_metadata) = outcome.upstream_metadata.as_ref() { if let Some(upstream_metadata) = outcome.upstream_metadata.as_ref() {
let merged_metadata = provider_query_persist_upstream_metadata(state, key, upstream_metadata).await?;
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?;
} }
if unique_models.is_empty() && !all_errors.is_empty() { 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); return Ok(fallback);
} }
} }
@@ -29,10 +29,7 @@ use crate::handlers::shared::{
provider_key_status_snapshot_payload, provider_key_status_snapshot_payload,
}; };
use crate::model_fetch::ModelFetchRuntimeState; use crate::model_fetch::ModelFetchRuntimeState;
use crate::provider_key_auth::{ use crate::provider_key_auth::provider_key_auth_semantics;
provider_key_auth_semantics, provider_key_configured_api_formats,
provider_key_inherits_provider_api_formats,
};
use crate::provider_transport::antigravity::{ use crate::provider_transport::antigravity::{
build_antigravity_safe_v1internal_request, build_antigravity_static_identity_headers, build_antigravity_safe_v1internal_request, build_antigravity_static_identity_headers,
classify_local_antigravity_request_support, AntigravityEnvelopeRequestType, classify_local_antigravity_request_support, AntigravityEnvelopeRequestType,
@@ -543,6 +540,22 @@ fn provider_query_build_test_request_body_for_api_format(
model: &str, model: &str,
route_path: &str, route_path: &str,
client_api_format: &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 { ) -> Value {
let client_api_format = provider_query_normalize_api_format_alias(client_api_format); let client_api_format = provider_query_normalize_api_format_alias(client_api_format);
let override_custom_model = route_path.ends_with("/test-model-failover") 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!( } else if matches!(
client_api_format.as_str(), 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")) ) && !value_has_non_empty_text(object.get("input"))
{ {
if let Some(prompt) = object if let Some(prompt) = object
@@ -580,7 +593,7 @@ fn provider_query_build_test_request_body_for_api_format(
} }
if matches!( if matches!(
client_api_format.as_str(), 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")) ) && value_has_non_empty_text(object.get("input"))
{ {
object.remove("prompt"); object.remove("prompt");
@@ -590,6 +603,9 @@ fn provider_query_build_test_request_body_for_api_format(
{ {
object.remove("messages"); object.remove("messages");
} }
if client_api_format == "openai:search" {
provider_query_ensure_search_test_fields(object, payload, search_session_id);
}
} }
return body; return body;
} }
@@ -620,6 +636,15 @@ fn provider_query_build_test_request_body_for_api_format(
"temperature": 0.7, "temperature": 0.7,
"stream": true, "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!({ "claude:messages" => json!({
"model": model, "model": model,
"messages": [{ "messages": [{
@@ -680,7 +705,7 @@ fn provider_query_insert_default_test_conversation(
.entry("top_n".to_string()) .entry("top_n".to_string())
.or_insert_with(|| Value::from(4_u64)); .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)); object.insert("input".to_string(), Value::String(message));
} }
"claude:messages" => { "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<String, Value>,
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 { 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() { match provider_query_normalize_api_format_alias(provider_api_format).as_str() {
"openai:responses" | "openai:responses:compact" => "openai:responses", "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, client_api_format: &str,
) -> bool { ) -> bool {
match provider_query_normalize_api_format_alias(client_api_format).as_str() { 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("input"))
|| value_has_non_empty_text(body.get("prompt")) || 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_type: &str,
provider_api_format: &str, provider_api_format: &str,
) -> bool { ) -> 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_type,
provider_api_format, provider_api_format,
);
crate::ai_serving::resolve_upstream_is_stream_from_endpoint_config(
endpoint_config,
false, false,
hard_requires_streaming, false,
) )
} }
@@ -876,15 +941,11 @@ fn provider_query_key_supports_endpoint(
provider_type: &str, provider_type: &str,
endpoint_api_format: &str, endpoint_api_format: &str,
) -> bool { ) -> bool {
if provider_key_inherits_provider_api_formats(key, provider_type) { crate::handlers::shared::provider_catalog_key_supports_format(
return true; key,
} provider_type,
let formats = provider_key_configured_api_formats(key); endpoint_api_format,
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)
} }
async fn provider_query_select_preferred_non_kiro_endpoint( async fn provider_query_select_preferred_non_kiro_endpoint(
@@ -1626,6 +1687,15 @@ fn provider_query_standard_execution_response_body(
{ {
return None; 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) Some(body)
} }
@@ -2104,8 +2174,18 @@ async fn provider_query_execute_openai_image_test_candidate(
route_path, route_path,
); );
let incoming_request_headers = provider_query_extract_request_headers(payload); 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() let mut synthetic_request = http::Request::builder()
.uri("/v1/images/generations") .uri(image_request_path)
.body(()) .body(())
.map_err(|err| GatewayError::Internal(err.to_string()))?; .map_err(|err| GatewayError::Internal(err.to_string()))?;
*synthetic_request.headers_mut() = incoming_request_headers; *synthetic_request.headers_mut() = incoming_request_headers;
@@ -2116,11 +2196,18 @@ async fn provider_query_execute_openai_image_test_candidate(
&parts, &parts,
&request_body, &request_body,
None, None,
provider_query_openai_image_normalize_options(provider_type), provider_query_openai_image_normalize_options(
provider_type,
Some(candidate.effective_model.as_str()),
),
) else { ) else {
return Ok(provider_query_skipped_execution_outcome( return Ok(provider_query_skipped_execution_outcome(
request_body.clone(), 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 .provider_type
.trim() .trim()
.eq_ignore_ascii_case("codex"); .eq_ignore_ascii_case("codex");
let mut provider_request_body = if is_chatgpt_web { let upstream_is_stream = crate::ai_serving::resolve_upstream_is_stream_for_provider(
match crate::ai_serving::build_chatgpt_web_image_request_body(&parts, &request_body, None) { transport.endpoint.config.as_ref(),
Ok(body) => body, transport.provider.provider_type.as_str(),
Err(err) => err.to_error_json(), "openai:image",
} request_body
} else if is_codex || is_grok { .get("stream")
crate::ai_serving::build_openai_image_provider_request_body(&normalized_request) .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 { } else {
crate::ai_serving::build_openai_image_api_provider_request_body( crate::ai_serving::build_openai_image_api_provider_request_body(
&normalized_request, &normalized_request,
Some(candidate.effective_model.as_str()), Some(candidate.effective_model.as_str()),
upstream_is_stream,
) )
}; };
if !is_chatgpt_web { let Some(provider_request_body) = provider_request_body else {
crate::ai_serving::apply_codex_openai_responses_special_body_edits( return Ok(provider_query_skipped_execution_outcome(
&mut provider_request_body, request_body,
transport.provider.provider_type.as_str(), "Provider request is outside the Codex Images contract",
"openai:image", ));
transport.endpoint.body_rules.as_ref(), };
Some(candidate.key.id.as_str()),
);
}
let oauth_auth = state.resolve_local_oauth_header_auth(&transport).await?; let oauth_auth = state.resolve_local_oauth_header_auth(&transport).await?;
let Some((auth_header, auth_value)) = let Some((auth_header, auth_value)) =
crate::provider_transport::resolve_openai_image_auth(&transport).or(oauth_auth) 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, headers: &parts.headers,
auth_header: &auth_header, auth_header: &auth_header,
auth_value: &auth_value, auth_value: &auth_value,
accept: if is_codex || is_chatgpt_web { accept: if is_codex {
"text/event-stream" None
} else if upstream_is_stream {
Some("text/event-stream")
} else { } else {
"application/json" Some("application/json")
}, },
header_rules: transport.endpoint.header_rules.as_ref(), header_rules: transport.endpoint.header_rules.as_ref(),
provider_request_body: &provider_request_body, 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()); request_headers.insert("x-aether-chatgpt-web-image".to_string(), "1".to_string());
} else if is_grok { } else if is_grok {
} else { } else {
crate::ai_serving::apply_codex_openai_responses_special_headers( crate::ai_serving::apply_codex_openai_special_headers(
&mut request_headers, &mut request_headers,
&provider_request_body, &provider_request_body,
&parts.headers, &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( let request_url = provider_query_openai_image_test_upstream_url(
&transport, &transport,
Some(parts.uri.path()), Some(image_request_path),
parts.uri.query(), parts.uri.query(),
); );
let upstream_is_stream = provider_request_body
.get("stream")
.and_then(Value::as_bool)
.unwrap_or(true);
let plan = ExecutionPlan { let plan = ExecutionPlan {
request_id: trace_id.to_string(), 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); crate::ai_serving::normalize_api_format_alias(provider_api_format);
let client_api_format = let client_api_format =
provider_query_standard_test_client_api_format(normalized_provider_api_format.as_str()); 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( let original_request_body =
payload, provider_query_build_test_request_body_for_api_format_with_search_session(
&candidate.effective_model, payload,
route_path, &candidate.effective_model,
client_api_format, route_path,
); client_api_format,
Some(trace_id),
);
if crate::provider_transport::is_windsurf_provider_transport(&transport) if crate::provider_transport::is_windsurf_provider_transport(&transport)
&& provider_query_normalize_api_format_alias(candidate.endpoint.api_format.as_str()) && provider_query_normalize_api_format_alias(candidate.endpoint.api_format.as_str())
== "openai:chat" == "openai:chat"
@@ -2992,6 +3100,45 @@ async fn provider_query_execute_standard_test_candidate(
); );
provider_request_body 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" "openai:embedding"
| "gemini:embedding" | "gemini:embedding"
| "jina:embedding" | "jina:embedding"
@@ -3046,6 +3193,37 @@ async fn provider_query_execute_standard_test_candidate(
upstream_is_stream, upstream_is_stream,
require_body_stream_field, 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) if crate::provider_transport::is_gemini_cli_provider_transport(&transport)
&& normalized_provider_api_format == "gemini:generate_content" && normalized_provider_api_format == "gemini:generate_content"
{ {
@@ -3122,6 +3300,7 @@ async fn provider_query_execute_standard_test_candidate(
"openai:chat" "openai:chat"
| "openai:responses" | "openai:responses"
| "openai:responses:compact" | "openai:responses:compact"
| "openai:search"
| "claude:messages" | "claude:messages"
| "gemini:generate_content" | "gemini:generate_content"
| "openai:embedding" | "openai:embedding"
@@ -3137,6 +3316,7 @@ async fn provider_query_execute_standard_test_candidate(
"openai:chat" "openai:chat"
| "openai:responses" | "openai:responses"
| "openai:responses:compact" | "openai:responses:compact"
| "openai:search"
| "openai:embedding" | "openai:embedding"
| "jina:embedding" | "jina:embedding"
| "doubao:embedding" | "doubao:embedding"
@@ -3202,7 +3382,7 @@ async fn provider_query_execute_standard_test_candidate(
&BTreeMap::new(), &BTreeMap::new(),
Some("application/json"), 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( crate::provider_transport::auth::build_complete_passthrough_headers_with_auth(
&parts.headers, &parts.headers,
auth_header.as_deref().unwrap_or_default(), auth_header.as_deref().unwrap_or_default(),
@@ -3268,8 +3448,10 @@ async fn provider_query_execute_standard_test_candidate(
response_body: None, response_body: None,
}); });
} }
if crate::ai_serving::is_openai_responses_format(provider_api_format) { if crate::ai_serving::is_openai_responses_family_format(provider_api_format)
crate::ai_serving::apply_codex_openai_responses_special_headers( || crate::ai_serving::api_format_alias_matches(provider_api_format, "openai:search")
{
crate::ai_serving::apply_codex_openai_special_headers(
&mut request_headers, &mut request_headers,
&provider_request_body, &provider_request_body,
&parts.headers, &parts.headers,
@@ -3278,9 +3460,17 @@ async fn provider_query_execute_standard_test_candidate(
Some(trace_id), Some(trace_id),
transport.key.decrypted_auth_config.as_deref(), 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, &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 { if !uses_vertex_query_auth {
@@ -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); let normalized_api_format = crate::ai_serving::normalize_api_format_alias(provider_api_format);
if normalized_api_format == "openai:responses:compact" { if normalized_api_format == "openai:responses:compact" {
"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) { } else if crate::ai_serving::is_embedding_api_format(&normalized_api_format) {
"openai:embedding" "openai:embedding"
} else if crate::ai_serving::is_rerank_api_format(&normalized_api_format) { } 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"
| "openai:responses:compact" | "openai:responses:compact"
| "openai:search"
| "claude:messages" | "claude:messages"
| "openai:embedding" | "openai:embedding"
| "jina:embedding" | "jina:embedding"
@@ -261,6 +264,7 @@ pub(super) fn provider_query_test_adapter_for_provider_api_format(
"openai:chat" "openai:chat"
| "openai:responses" | "openai:responses"
| "openai:responses:compact" | "openai:responses:compact"
| "openai:search"
| "claude:messages" | "claude:messages"
| "gemini:generate_content" | "gemini:generate_content"
| "gemini:interactions" | "gemini:interactions"
@@ -365,6 +369,7 @@ pub(super) fn provider_query_transport_supports_model_test_execution(
} }
"openai:responses" "openai:responses"
| "openai:responses:compact" | "openai:responses:compact"
| "openai:search"
| "openai:embedding" | "openai:embedding"
| "jina:embedding" | "jina:embedding"
| "doubao:embedding" | "doubao:embedding"
@@ -9,16 +9,19 @@ pub(super) struct ProviderQueryOpenAiImageTestCapability(AdminProviderOpenAiImag
pub(super) fn provider_query_openai_image_test_capability( pub(super) fn provider_query_openai_image_test_capability(
provider_type: &str, provider_type: &str,
provider_model: Option<&str>,
) -> ProviderQueryOpenAiImageTestCapability { ) -> ProviderQueryOpenAiImageTestCapability {
ProviderQueryOpenAiImageTestCapability(admin_provider_openai_image_test_capability( ProviderQueryOpenAiImageTestCapability(admin_provider_openai_image_test_capability(
provider_type, provider_type,
provider_model,
)) ))
} }
pub(super) fn provider_query_openai_image_normalize_options( pub(super) fn provider_query_openai_image_normalize_options(
provider_type: &str, provider_type: &str,
provider_model: Option<&str>,
) -> crate::ai_serving::OpenAiImageNormalizeOptions { ) -> 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<u64> { pub(super) fn provider_query_openai_image_requested_count(request_body: &Value) -> Option<u64> {
@@ -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( pub(super) fn provider_query_openai_image_normalize_failure_message(
provider_type: &str, provider_type: &str,
provider_model: Option<&str>,
request_body: &Value, request_body: &Value,
) -> String { ) -> 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) if provider_query_openai_image_requested_count(request_body)
.is_some_and(|value| !capability.0.supports_generation_count(value)) .is_some_and(|value| !capability.0.supports_generation_count(value))
{ {
@@ -101,8 +101,7 @@ fn provider_query_model_mapping_matches_endpoint(
) -> bool { ) -> bool {
let api_format_matches = mapping.api_formats.as_ref().is_none_or(|api_formats| { let api_format_matches = mapping.api_formats.as_ref().is_none_or(|api_formats| {
api_formats.iter().any(|value| { api_formats.iter().any(|value| {
aether_scheduler_core::normalize_api_format(value) crate::ai_serving::api_format_permission_covers(value, &endpoint.api_format)
== aether_scheduler_core::normalize_api_format(&endpoint.api_format)
}) })
}); });
if !api_format_matches { if !api_format_matches {
@@ -367,3 +366,41 @@ fn provider_query_parse_mapping_string_list_array(
} }
Ok(parsed) 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"),
));
}
}
@@ -337,6 +337,11 @@ fn provider_query_standard_test_resolves_codex_responses_upstream_streaming() {
"codex", "codex",
"openai:responses:compact", "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( assert!(!provider_query_resolve_standard_test_upstream_is_stream(
None, None,
"custom", "custom",
@@ -470,6 +475,118 @@ fn provider_query_compact_test_request_body_defaults_to_responses_input() {
assert!(body.get("messages").is_none()); 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] #[test]
fn provider_query_embedding_test_request_body_defaults_to_embedding_input() { fn provider_query_embedding_test_request_body_defaults_to_embedding_input() {
let payload = json!({"message": "hello from embedding"}); 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"), provider_query_test_adapter_for_provider_api_format("codex", "openai:responses:compact"),
Some(ProviderQueryTestAdapter::Standard) Some(ProviderQueryTestAdapter::Standard)
); );
assert_eq!(
provider_query_test_adapter_for_provider_api_format("codex", "openai:search"),
Some(ProviderQueryTestAdapter::Standard)
);
assert_eq!( assert_eq!(
provider_query_test_adapter_for_provider_api_format("chatgpt_web", "openai:image"), provider_query_test_adapter_for_provider_api_format("chatgpt_web", "openai:image"),
Some(ProviderQueryTestAdapter::OpenAiImage) 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"), provider_query_model_test_endpoint_priority("codex", "openai:responses:compact"),
Some(1) Some(1)
); );
assert_eq!(
provider_query_model_test_endpoint_priority("codex", "openai:search"),
Some(1)
);
assert_eq!( assert_eq!(
provider_query_model_test_endpoint_priority("chatgpt_web", "openai:image"), provider_query_model_test_endpoint_priority("chatgpt_web", "openai:image"),
Some(2) Some(2)
@@ -997,7 +1122,7 @@ fn provider_query_grok_image_test_allows_multi_generation_count() {
&parts, &parts,
&body, &body,
None, 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"); .expect("grok image model tests should allow multi-image generation");
let provider_body = crate::ai_serving::build_openai_image_provider_request_body(&normalized); 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, &parts,
&body, &body,
None, None,
provider_query_openai_image_normalize_options("chatgpt_web"), provider_query_openai_image_normalize_options("chatgpt_web", Some("gpt-image-2")),
) )
.is_none() .is_none()
); );
assert_eq!( 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" "Provider request body could not be normalized for openai:image: selected provider supports n=1..1 for generation"
); );
} }
@@ -1,5 +1,6 @@
use crate::image_capabilities::{ 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}; use serde_json::{json, Value};
@@ -24,16 +25,21 @@ impl AdminProviderOpenAiImageTestCapability {
pub(crate) fn admin_provider_openai_image_test_capability( pub(crate) fn admin_provider_openai_image_test_capability(
provider_type: &str, provider_type: &str,
provider_model: Option<&str>,
) -> AdminProviderOpenAiImageTestCapability { ) -> AdminProviderOpenAiImageTestCapability {
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( pub(crate) fn admin_provider_openai_image_normalize_options(
provider_type: &str, provider_type: &str,
provider_model: Option<&str>,
) -> crate::ai_serving::OpenAiImageNormalizeOptions { ) -> 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( 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; provider_type.eq_ignore_ascii_case("grok") && model_id == GROK_IMAGE_EDIT_MODEL_ID;
let openai_image = if supports_image_generation { let openai_image = if supports_image_generation {
Some(json!({ 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_generation": !is_grok_image_edit,
"supports_edit": is_grok_image_edit, "supports_edit": is_grok_image_edit,
})) }))
@@ -105,6 +111,13 @@ mod tests {
assert!(payload["openai:image"].is_null()); 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] #[test]
fn grok_image_support_uses_catalog_model_ids_not_global_fallback() { fn grok_image_support_uses_catalog_model_ids_not_global_fallback() {
assert!(admin_provider_model_supports_image_generation( assert!(admin_provider_model_supports_image_generation(
@@ -2,7 +2,7 @@ use crate::handlers::admin::provider::shared::payloads::AdminProviderKeyCreateRe
use crate::handlers::admin::provider::write::normalize::{ use crate::handlers::admin::provider::write::normalize::{
normalize_allow_auth_channel_mismatch_formats, normalize_api_format_json_object_keys, 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_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::request::AdminAppState;
use crate::handlers::admin::shared::{ use crate::handlers::admin::shared::{
@@ -165,7 +165,7 @@ pub(crate) async fn build_admin_create_provider_key_record(
}, },
encrypted_api_key, encrypted_api_key,
encrypted_auth_config, encrypted_auth_config,
normalize_api_format_json_object_keys(payload.rate_multipliers, "rate_multipliers")?, normalize_rate_multipliers(payload.rate_multipliers)?,
None, None,
normalize_string_list(payload.allowed_models).map(|value| json!(value)), normalize_string_list(payload.allowed_models).map(|value| json!(value)),
None, None,
@@ -2,7 +2,7 @@ use crate::handlers::admin::provider::shared::payloads::AdminProviderKeyUpdatePa
use crate::handlers::admin::provider::write::normalize::{ use crate::handlers::admin::provider::write::normalize::{
normalize_allow_auth_channel_mismatch_formats, normalize_api_format_json_object_keys, 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_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::request::AdminAppState;
use crate::handlers::admin::shared::{ use crate::handlers::admin::shared::{
@@ -260,8 +260,7 @@ pub(crate) async fn build_admin_update_provider_key_record(
updated.name = trimmed.to_string(); updated.name = trimmed.to_string();
} }
if fields.contains("rate_multipliers") { if fields.contains("rate_multipliers") {
updated.rate_multipliers = updated.rate_multipliers = normalize_rate_multipliers(payload.rate_multipliers)?;
normalize_api_format_json_object_keys(payload.rate_multipliers, "rate_multipliers")?;
} }
if let Some(internal_priority) = payload.internal_priority { if let Some(internal_priority) = payload.internal_priority {
updated.internal_priority = internal_priority; updated.internal_priority = internal_priority;
@@ -42,6 +42,31 @@ pub(crate) fn normalize_api_format_json_object_keys(
Ok(Some(serde_json::Value::Object(normalized))) Ok(Some(serde_json::Value::Object(normalized)))
} }
pub(crate) fn normalize_rate_multipliers(
value: Option<serde_json::Value>,
) -> Result<Option<serde_json::Value>, 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( pub(crate) fn normalize_auth_type_by_format(
value: Option<serde_json::Value>, value: Option<serde_json::Value>,
field_name: &str, field_name: &str,
@@ -212,7 +237,7 @@ mod tests {
normalize_allow_auth_channel_mismatch_formats, normalize_api_format_json_object_keys, 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_api_format_list, normalize_auth_type, normalize_auth_type_by_format,
normalize_chat_pii_redaction_config, normalize_pool_advanced_config, 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; 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] #[test]
fn normalize_pool_advanced_rejects_legacy_booleans() { fn normalize_pool_advanced_rejects_legacy_booleans() {
assert_eq!( assert_eq!(
@@ -3,11 +3,79 @@ mod endpoint;
mod template; mod template;
mod update; mod update;
fn normalize_provider_request_timeout(value: Option<f64>) -> Result<Option<f64>, 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<f64>) -> Result<Option<f64>, 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::create::build_admin_create_provider_record;
pub(crate) use self::endpoint::build_admin_fixed_provider_endpoint_record; pub(crate) use self::endpoint::build_admin_fixed_provider_endpoint_record;
pub(crate) use self::template::{ pub(crate) use self::template::{
apply_admin_fixed_provider_endpoint_template_overrides, apply_admin_fixed_provider_endpoint_template_overrides,
reconcile_admin_fixed_provider_template_endpoints, 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; 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());
}
}
@@ -107,16 +107,9 @@ pub(crate) async fn build_admin_create_provider_record(
None => Some(2), None => Some(2),
}; };
let proxy = normalize_json_object(payload.proxy, "proxy")?; let proxy = normalize_json_object(payload.proxy, "proxy")?;
let stream_first_byte_timeout_secs = match payload.stream_first_byte_timeout { let stream_first_byte_timeout_secs =
Some(value) if (1.0..=300.0).contains(&value) => Some(value), super::normalize_provider_stream_first_byte_timeout(payload.stream_first_byte_timeout)?;
Some(_) => return Err("stream_first_byte_timeout 必须是 1 到 300 之间的数字".to_string()), let request_timeout_secs = super::normalize_provider_request_timeout(payload.request_timeout)?;
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 mut config_map = normalize_json_object(payload.config, "config")? let mut config_map = normalize_json_object(payload.config, "config")?
.and_then(|value| value.as_object().cloned()) .and_then(|value| value.as_object().cloned())
@@ -1,11 +1,11 @@
use super::endpoint::{ use super::endpoint::{
build_admin_fixed_provider_endpoint_defaults, build_admin_fixed_provider_endpoint_record, build_admin_fixed_provider_endpoint_defaults, build_admin_fixed_provider_endpoint_record,
AdminFixedProviderEndpointDefaults,
}; };
use crate::handlers::admin::request::AdminAppState; use crate::handlers::admin::request::AdminAppState;
use crate::provider_key_auth::provider_key_is_oauth_managed;
use crate::GatewayError; use crate::GatewayError;
use aether_data_contracts::repository::provider_catalog::{ use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider, StoredProviderCatalogEndpoint, StoredProviderCatalogProvider,
}; };
use aether_provider_transport::provider_types::{ use aether_provider_transport::provider_types::{
fixed_provider_template, FixedProviderEndpointTemplate, FixedProviderTemplate, fixed_provider_template, FixedProviderEndpointTemplate, FixedProviderTemplate,
@@ -36,6 +36,30 @@ struct FixedProviderEndpointMetadata {
pub(crate) async fn reconcile_admin_fixed_provider_template_endpoints( pub(crate) async fn reconcile_admin_fixed_provider_template_endpoints(
state: &AdminAppState<'_>, state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider, 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> { ) -> Result<(), GatewayError> {
let Some(template) = state.fixed_provider_template(&provider.provider_type) else { let Some(template) = state.fixed_provider_template(&provider.provider_type) else {
return Ok(()); return Ok(());
@@ -53,8 +77,9 @@ pub(crate) async fn reconcile_admin_fixed_provider_template_endpoints(
match existing_endpoint { match existing_endpoint {
Some(existing_endpoint) => { Some(existing_endpoint) => {
matched_endpoint_ids.insert(existing_endpoint.id.clone()); matched_endpoint_ids.insert(existing_endpoint.id.clone());
let updated = reconcile_fixed_provider_endpoint( let updated = reconcile_fixed_provider_endpoint_with_adoption_provider(
provider, provider,
adoption_provider,
existing_endpoint, existing_endpoint,
template, template,
endpoint_template, endpoint_template,
@@ -115,31 +140,6 @@ pub(crate) async fn reconcile_admin_fixed_provider_template_endpoints(
Ok(()) 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( pub(crate) fn apply_admin_fixed_provider_endpoint_template_overrides(
provider: &StoredProviderCatalogProvider, provider: &StoredProviderCatalogProvider,
existing_endpoint: &StoredProviderCatalogEndpoint, existing_endpoint: &StoredProviderCatalogEndpoint,
@@ -156,8 +156,14 @@ pub(crate) fn apply_admin_fixed_provider_endpoint_template_overrides(
let defaults = let defaults =
build_admin_fixed_provider_endpoint_defaults(provider, template, endpoint_template)?; build_admin_fixed_provider_endpoint_defaults(provider, template, endpoint_template)?;
let mut metadata = fixed_provider_endpoint_metadata(existing_endpoint) let mut metadata = fixed_provider_endpoint_metadata(existing_endpoint).unwrap_or_else(|| {
.unwrap_or_else(|| managed_fixed_provider_endpoint_metadata(template, endpoint_template)); adopt_fixed_provider_endpoint_metadata(
existing_endpoint,
&defaults,
template,
endpoint_template,
)
});
let mut overrides = metadata.overrides.clone(); let mut overrides = metadata.overrides.clone();
sync_override_if_changed( sync_override_if_changed(
@@ -250,17 +256,48 @@ pub(crate) fn apply_admin_fixed_provider_endpoint_template_overrides(
Ok(()) Ok(())
} }
#[cfg(test)]
fn reconcile_fixed_provider_endpoint( fn reconcile_fixed_provider_endpoint(
provider: &StoredProviderCatalogProvider, provider: &StoredProviderCatalogProvider,
existing_endpoint: &StoredProviderCatalogEndpoint, existing_endpoint: &StoredProviderCatalogEndpoint,
template: &FixedProviderTemplate, template: &FixedProviderTemplate,
endpoint_template: &FixedProviderEndpointTemplate, endpoint_template: &FixedProviderEndpointTemplate,
) -> Result<StoredProviderCatalogEndpoint, String> {
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<StoredProviderCatalogEndpoint, String> { ) -> Result<StoredProviderCatalogEndpoint, String> {
let defaults = let defaults =
build_admin_fixed_provider_endpoint_defaults(provider, template, endpoint_template)?; build_admin_fixed_provider_endpoint_defaults(provider, template, endpoint_template)?;
let mut updated = existing_endpoint.clone(); let mut updated = existing_endpoint.clone();
let metadata = fixed_provider_endpoint_metadata(existing_endpoint) let metadata = match fixed_provider_endpoint_metadata(existing_endpoint) {
.unwrap_or_else(|| managed_fixed_provider_endpoint_metadata(template, endpoint_template)); 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_format = defaults.api_format.clone();
updated.api_family = Some(defaults.api_family.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( fn upsert_fixed_provider_endpoint_metadata(
endpoint: &mut StoredProviderCatalogEndpoint, endpoint: &mut StoredProviderCatalogEndpoint,
metadata: &FixedProviderEndpointMetadata, metadata: &FixedProviderEndpointMetadata,
@@ -513,22 +621,6 @@ fn current_unix_secs() -> u64 {
.unwrap_or(0) .unwrap_or(0)
} }
fn reconcile_fixed_provider_key(
provider: &StoredProviderCatalogProvider,
existing_key: &StoredProviderCatalogKey,
) -> Option<StoredProviderCatalogKey> {
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<T>(overrides: &mut BTreeSet<String>, key: &str, actual: &T, desired: &T) fn sync_override<T>(overrides: &mut BTreeSet<String>, key: &str, actual: &T, desired: &T)
where where
T: PartialEq, T: PartialEq,
@@ -624,4 +716,68 @@ mod tests {
.expect("endpoint should reconcile"); .expect("endpoint should reconcile");
assert_eq!(reconciled.base_url, "http://127.0.0.1:18181/v1"); 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);
}
} }
@@ -202,21 +202,13 @@ pub(crate) async fn build_admin_update_provider_record(
} }
if fields.contains("stream_first_byte_timeout") { if fields.contains("stream_first_byte_timeout") {
updated.stream_first_byte_timeout_secs = match payload.stream_first_byte_timeout { updated.stream_first_byte_timeout_secs =
Some(value) if (1.0..=300.0).contains(&value) => Some(value), super::normalize_provider_stream_first_byte_timeout(payload.stream_first_byte_timeout)?;
Some(_) => {
return Err("stream_first_byte_timeout 必须是 1 到 300 之间的数字".to_string());
}
None => None,
};
} }
if fields.contains("request_timeout") { if fields.contains("request_timeout") {
updated.request_timeout_secs = match payload.request_timeout { updated.request_timeout_secs =
Some(value) if (1.0..=600.0).contains(&value) => Some(value), super::normalize_provider_request_timeout(payload.request_timeout)?;
Some(_) => return Err("request_timeout 必须是 1 到 600 之间的数字".to_string()),
None => None,
};
} }
if fields.contains("enable_format_conversion") { if fields.contains("enable_format_conversion") {
@@ -1,5 +1,26 @@
use super::*; 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> { impl<'a> AdminAppState<'a> {
pub(crate) async fn build_admin_keys_grouped_by_format_payload( pub(crate) async fn build_admin_keys_grouped_by_format_payload(
&self, &self,
@@ -232,6 +253,7 @@ impl<'a> AdminAppState<'a> {
let (normalized_api_format, api_family, endpoint_kind) = let (normalized_api_format, api_family, endpoint_kind) =
admin_endpoint_signature_parts(&payload.api_format) admin_endpoint_signature_parts(&payload.api_format)
.ok_or_else(|| format!("无效的 api_format: {}", 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 base_url = normalize_admin_base_url(&payload.base_url)?;
let existing_endpoints = self let existing_endpoints = self
@@ -337,6 +359,13 @@ impl<'a> AdminAppState<'a> {
&update_fields, &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" if provider_type == "codex"
&& crate::ai_serving::is_openai_responses_format(&existing_endpoint.api_format) && crate::ai_serving::is_openai_responses_format(&existing_endpoint.api_format)
{ {
@@ -389,3 +418,40 @@ impl<'a> AdminAppState<'a> {
Ok(updated) 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());
}
}
+177 -14
View File
@@ -84,6 +84,8 @@ const OPENAI_RESPONSES_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL: &str =
"当前 OpenAI Responses 请求无法在本地执行:没有匹配到可用的执行路径"; "当前 OpenAI Responses 请求无法在本地执行:没有匹配到可用的执行路径";
const OPENAI_RESPONSES_COMPACT_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL: &str = const OPENAI_RESPONSES_COMPACT_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL: &str =
"当前 OpenAI Responses Compact 请求无法在本地执行:没有匹配到可用的执行路径"; "当前 OpenAI Responses Compact 请求无法在本地执行:没有匹配到可用的执行路径";
const OPENAI_SEARCH_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL: &str =
"当前 OpenAI Search 请求无法在本地执行:没有匹配到可用的执行路径";
const OPENAI_VIDEO_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL: &str = const OPENAI_VIDEO_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL: &str =
"当前 OpenAI Video 请求无法在本地执行:没有匹配到可用的执行路径"; "当前 OpenAI Video 请求无法在本地执行:没有匹配到可用的执行路径";
const CLAUDE_MESSAGES_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL: &str = 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::<serde_json::Value>(body.as_ref()).ok() serde_json::from_slice::<serde_json::Value>(body.as_ref()).ok()
}); });
let client_session_affinity = let client_session_affinity =
crate::client_session_affinity::client_session_affinity_from_parts( crate::client_session_affinity::client_session_affinity_from_api_request(
parts, api_format,
&parts.headers,
body_json.as_ref(), body_json.as_ref(),
); );
let Some(target) = crate::scheduler::affinity::read_cached_scheduler_affinity_target( 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('/'), owner.relay_base_url.trim_end_matches('/'),
request_context.request_path_and_query() 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 { for (name, value) in &parts.headers {
if should_skip_request_header(name.as_str()) || name == http::header::HOST { if should_skip_request_header(name.as_str()) || name == http::header::HOST {
continue; 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()); upstream_request.header(TRUSTED_AUTH_BALANCE_HEADER, balance_remaining.to_string());
} }
let upstream_response = upstream_request let upstream_response = crate::tunnel::send_owner_forward_request(
.body(buffered_body.cloned().unwrap_or_default()) upstream_request.body(buffered_body.cloned().unwrap_or_default()),
.send() stream_first_byte_timeout,
.await )
.map_err(|err| GatewayError::UpstreamUnavailable { .await
trace_id: request_context.trace_id.clone(), .map_err(|message| GatewayError::UpstreamUnavailable {
message: format!("owner gateway affinity forward failed: {err}"), trace_id: request_context.trace_id.clone(),
})?; message: format!("owner gateway affinity forward failed: {message}"),
})?;
let mut response = build_sync_aware_affinity_forward_response( let mut response = build_sync_aware_affinity_forward_response(
request_context, request_context,
@@ -523,6 +547,29 @@ async fn maybe_forward_public_request_to_tunnel_owner(
Ok(Some(response)) 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 { fn upstream_response_is_sse(headers: &reqwest::header::HeaderMap) -> bool {
headers headers
.get(http::header::CONTENT_TYPE) .get(http::header::CONTENT_TYPE)
@@ -2196,6 +2243,7 @@ fn local_execution_runtime_miss_route_label(
"/v1/chat/completions" => "OpenAI Chat Completions", "/v1/chat/completions" => "OpenAI Chat Completions",
"/v1/responses" => "OpenAI Responses", "/v1/responses" => "OpenAI Responses",
"/v1/responses/compact" => "OpenAI Responses Compact", "/v1/responses/compact" => "OpenAI Responses Compact",
"/v1/alpha/search" => "OpenAI Search",
"/v1/messages" => "Claude Messages", "/v1/messages" => "Claude Messages",
path if path.starts_with("/v1/videos") => "OpenAI Video", path if path.starts_with("/v1/videos") => "OpenAI Video",
path if path.starts_with("/upload/v1beta/files") || path.starts_with("/v1beta/files") => { 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" => { "/v1/responses/compact" => {
Some(OPENAI_RESPONSES_COMPACT_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL) 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), "/v1/messages" => Some(CLAUDE_MESSAGES_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL),
path if path.starts_with("/v1/videos") => { path if path.starts_with("/v1/videos") => {
Some(OPENAI_VIDEO_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL) Some(OPENAI_VIDEO_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL)
@@ -2262,15 +2311,129 @@ mod tests {
use super::{ use super::{
api_key_remote_ip_allowed, buffer_and_normalize_request_body, api_key_remote_ip_allowed, buffer_and_normalize_request_body,
diagnostic_is_auth_api_key_concurrency_limited, local_execution_runtime_miss_detail, diagnostic_is_auth_api_key_concurrency_limited, local_execution_runtime_miss_detail,
restore_redacted_stream_execution_response, restore_redacted_sync_execution_response, owner_forward_request_is_stream, restore_redacted_stream_execution_response,
GatewayControlDecision, LocalExecutionRuntimeMissDiagnostic, RequestBodyBufferError, restore_redacted_sync_execution_response, GatewayControlDecision,
RequestBodyBufferPolicy, LocalExecutionRuntimeMissDiagnostic, RequestBodyBufferError, RequestBodyBufferPolicy,
}; };
use axum::body::{to_bytes, Body, Bytes}; use axum::body::{to_bytes, Body, Bytes};
use axum::http::{header, HeaderMap, HeaderValue, Method, Response}; use axum::http::{header, HeaderMap, HeaderValue, Method, Response};
use serde_json::json; use serde_json::json;
use tokio::sync::Semaphore; 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] #[test]
fn api_key_remote_ip_allows_unrestricted_keys() { fn api_key_remote_ip_allows_unrestricted_keys() {
let remote_ip = "203.0.113.10".parse().expect("valid ip"); let remote_ip = "203.0.113.10".parse().expect("valid ip");
@@ -1,9 +1,8 @@
use crate::ai_serving::normalize_openai_image_quality;
use crate::async_task::CancelVideoTaskError; use crate::async_task::CancelVideoTaskError;
use crate::control::GatewayControlDecision; use crate::control::GatewayControlDecision;
use crate::control::GatewayPublicRequestContext; use crate::control::GatewayPublicRequestContext;
use crate::image_capabilities::{ use crate::image_capabilities::openai_image_gateway_max_generation_count;
openai_image_gateway_max_generation_count, openai_image_gateway_max_generation_count_for_model,
};
use crate::{AppState, GatewayError}; use crate::{AppState, GatewayError};
use aether_data_contracts::repository::video_tasks::{ use aether_data_contracts::repository::video_tasks::{
StoredVideoTask, VideoTaskQueryFilter, VideoTaskStatus, 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_STYLE_DETAIL: &str = "当前 Codex 图片反代暂不支持 style 参数";
const OPENAI_IMAGE_RESPONSE_FORMAT_DETAIL: &str = "response_format 仅支持 url 或 b64_json"; 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_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_BACKGROUND_DETAIL: &str = "background 仅支持 auto、opaque 或 transparent";
const OPENAI_IMAGE_MODERATION_DETAIL: &str = "moderation 仅支持 auto 或 low"; const OPENAI_IMAGE_MODERATION_DETAIL: &str = "moderation 仅支持 auto 或 low";
const OPENAI_IMAGE_INPUT_FIDELITY_DETAIL: &str = "input_fidelity 仅支持 low 或 high"; 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 if validation
.quality .quality
.as_deref() .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( return Some(build_ai_public_error_response(
http::StatusCode::BAD_REQUEST, 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<String> { fn validate_openai_image_n(validation: &OpenAiImageValidationInput) -> Option<String> {
let max_generation_count = let max_generation_count = openai_image_gateway_max_generation_count();
openai_image_gateway_max_generation_count_for_model(validation.model.as_deref());
validation validation
.n .n
.is_some_and(|value| value == 0 || value > max_generation_count) .is_some_and(|value| value == 0 || value > max_generation_count)
@@ -1761,7 +1759,7 @@ mod tests {
} }
#[test] #[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_body = Bytes::from_static(br#"{"model":"gpt-image-2","prompt":"draw","n":2}"#);
let openai_validation = parse_openai_image_validation_input( let openai_validation = parse_openai_image_validation_input(
OpenAiImageOperation::Generate, OpenAiImageOperation::Generate,
@@ -1770,10 +1768,7 @@ mod tests {
) )
.expect("valid image payload should parse"); .expect("valid image payload should parse");
assert_eq!( assert!(validate_openai_image_n(&openai_validation).is_none());
validate_openai_image_n(&openai_validation).as_deref(),
Some("当前图片模型仅支持 n=1..1")
);
let grok_body = let grok_body =
Bytes::from_static(br#"{"model":"grok-imagine-image-lite","prompt":"draw","n":4}"#); 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"); .expect("valid grok image payload should parse");
assert!(validate_openai_image_n(&grok_validation).is_none()); 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")
);
} }
} }
@@ -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<Body> { pub(super) fn build_empty_models_list_response(api_format: &str) -> Response<Body> {
match api_format { match api_format {
"openai:responses" => Json(json!({ "models": [] })).into_response(),
"claude:messages" => Json(json!({ "claude:messages" => Json(json!({
"data": [], "data": [],
"has_more": false, "has_more": false,
@@ -101,6 +102,10 @@ pub(super) fn build_empty_models_list_response(api_format: &str) -> Response<Bod
} }
} }
pub(super) fn build_codex_models_list_response(models: Vec<serde_json::Value>) -> Response<Body> {
Json(json!({ "models": models })).into_response()
}
pub(super) fn build_openai_models_list_response( pub(super) fn build_openai_models_list_response(
rows: &[StoredMinimalCandidateSelectionRow], rows: &[StoredMinimalCandidateSelectionRow],
) -> Response<Body> { ) -> Response<Body> {
@@ -1,21 +1,24 @@
use std::collections::{BTreeMap, BTreeSet};
use std::fmt::Debug; use std::fmt::Debug;
use std::future::Future; use std::future::Future;
use std::time::{Duration, SystemTime, UNIX_EPOCH}; use std::time::{Duration, SystemTime, UNIX_EPOCH};
use aether_data_contracts::repository::candidate_selection::StoredMinimalCandidateSelectionRow; use aether_data_contracts::repository::candidate_selection::StoredMinimalCandidateSelectionRow;
use axum::{body::Body, response::Response}; use axum::{body::Body, response::Response};
use serde_json::Value;
use tokio::time::timeout; use tokio::time::timeout;
use tracing::warn; use tracing::warn;
use super::models_responses::{ use super::models_responses::{
build_claude_model_detail_response, build_claude_models_list_response, build_claude_model_detail_response, build_claude_models_list_response,
build_empty_models_list_response, build_gemini_model_detail_response, build_codex_models_list_response, build_empty_models_list_response,
build_gemini_models_list_response, build_models_auth_error_response, build_gemini_model_detail_response, build_gemini_models_list_response,
build_models_not_found_response, build_openai_model_detail_response, build_models_auth_error_response, build_models_not_found_response,
build_openai_models_list_response, build_openai_model_detail_response, build_openai_models_list_response,
}; };
use super::models_shared::{ 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}; 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); const MODELS_ROUTE_READ_TIMEOUT: Duration = Duration::from_secs(5);
#[cfg(test)] #[cfg(test)]
const MODELS_ROUTE_READ_TIMEOUT: Duration = Duration::from_millis(50); 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<T, E, Fut>(operation: &'static str, future: Fut) -> Option<T> async fn await_models_route_read<T, E, Fut>(operation: &'static str, future: Fut) -> Option<T>
where where
@@ -72,7 +76,7 @@ fn build_models_read_fallback_response(
} }
} }
fn sort_and_dedup_model_rows( fn sort_model_rows(
mut rows: Vec<StoredMinimalCandidateSelectionRow>, mut rows: Vec<StoredMinimalCandidateSelectionRow>,
) -> Vec<StoredMinimalCandidateSelectionRow> { ) -> Vec<StoredMinimalCandidateSelectionRow> {
rows.sort_by(|left, right| { 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.key_id.cmp(&right.key_id))
.then(left.model_id.cmp(&right.model_id)) .then(left.model_id.cmp(&right.model_id))
}); });
rows
}
fn sort_and_dedup_model_rows(
rows: Vec<StoredMinimalCandidateSelectionRow>,
) -> Vec<StoredMinimalCandidateSelectionRow> {
let mut deduped = Vec::with_capacity(rows.len()); let mut deduped = Vec::with_capacity(rows.len());
let mut last_model_name: Option<String> = None; let mut last_model_name: Option<String> = 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()) { if last_model_name.as_deref() == Some(row.global_model_name.as_str()) {
continue; continue;
} }
@@ -97,22 +107,157 @@ fn sort_and_dedup_model_rows(
deduped 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<String, Value>) -> 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<Value> {
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<Value> {
let cache_keys = rows
.iter()
.filter(|row| is_codex_provider_row(row))
.map(|row| format!("upstream_models:{}:{}", row.provider_id, row.key_id))
.collect::<BTreeSet<_>>()
.into_iter()
.collect::<Vec<_>>();
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::<Vec<Value>>(raw.as_deref()?).ok()?;
Some((key, models))
})
.collect::<BTreeMap<_, _>>();
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( async fn list_model_rows_for_client_format(
state: &AppState, state: &AppState,
api_format: &str, api_format: &str,
auth_snapshot: Option<&crate::data::auth::GatewayAuthApiKeySnapshot>, auth_snapshot: Option<&crate::data::auth::GatewayAuthApiKeySnapshot>,
) -> Option<Vec<StoredMinimalCandidateSelectionRow>> { ) -> Option<Vec<StoredMinimalCandidateSelectionRow>> {
let mut collected = Vec::new(); 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( let rows = await_models_route_read(
"candidate_selection_by_api_format", "candidate_selection_by_api_format",
state.list_minimal_candidate_selection_rows_for_api_format(query_format), state.list_minimal_candidate_selection_rows_for_api_format(query_format),
) )
.await?; .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); 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( 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() { if rows.is_empty() {
return Some(build_empty_models_list_response(api_format)); 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 { let response = match api_format {
"claude:messages" => { "claude:messages" => {
let before_id = query_param_value( let before_id = query_param_value(

Some files were not shown because too many files have changed in this diff Show More