From 810c3dfe2be328c1a2bf5dfe71c5c1ccb1bfcccc Mon Sep 17 00:00:00 2001 From: zhefox Date: Fri, 14 Aug 2026 17:55:05 +0800 Subject: [PATCH] fix(codex): serve versioned dynamic model catalogs --- Cargo.lock | 2 + Cargo.toml | 1 + apps/aether-gateway/Cargo.toml | 2 + .../src/ai_serving/planner/decision_input.rs | 2 +- .../aether-gateway/src/ai_serving/pure/mod.rs | 9 +- .../src/ai_serving/transport.rs | 5 +- apps/aether-gateway/src/backup/worker.rs | 18 +- apps/aether-gateway/src/cache/auth_runtime.rs | 10 - apps/aether-gateway/src/control/route/mod.rs | 3 +- .../src/control/tests/public_support.rs | 4 +- apps/aether-gateway/src/data/state/catalog.rs | 10 + .../src/data/state/provider_catalog_cache.rs | 57 + .../src/execution_runtime/sync/execution.rs | 35 +- .../src/execution_runtime/transport.rs | 349 +- .../admin/provider/oauth/quota/codex/mod.rs | 2 +- .../aether-gateway/src/handlers/public/mod.rs | 2 + .../src/handlers/public/support.rs | 2 + .../src/handlers/public/support/models.rs | 2 + .../public/support/models/responses.rs | 11 +- .../handlers/public/support/models/route.rs | 285 +- .../handlers/public/support/models/shared.rs | 2 +- .../aether-gateway/src/model_fetch/catalog.rs | 3585 +++++++++++++++++ apps/aether-gateway/src/model_fetch/mod.rs | 7 + .../aether-gateway/src/model_fetch/runtime.rs | 22 +- .../src/model_fetch/runtime/state.rs | 8 + .../src/request_candidate_queue.rs | 4 +- apps/aether-gateway/src/state/catalog.rs | 10 + apps/aether-gateway/src/state/integrations.rs | 107 +- .../src/tests/control/admin/oauth.rs | 29 +- apps/aether-gateway/src/tests/frontdoor/ai.rs | 762 +++- crates/aether-ai/formats/src/api.rs | 6 +- .../src/formats/openai/responses/codex.rs | 135 +- .../src/formats/openai/responses/history.rs | 12 +- crates/aether-ai/formats/src/lib.rs | 7 +- .../src/repository/provider_catalog/types.rs | 12 + crates/aether-model-fetch/src/lib.rs | 16 +- crates/aether-model-fetch/src/logic.rs | 401 +- crates/aether-model-fetch/src/strategy.rs | 419 +- crates/aether-model-fetch/src/transport.rs | 154 +- 39 files changed, 6147 insertions(+), 362 deletions(-) create mode 100644 apps/aether-gateway/src/model_fetch/catalog.rs diff --git a/Cargo.lock b/Cargo.lock index c7b536424..7067f3cd3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -325,6 +325,7 @@ dependencies = [ "axum", "base64 0.22.1", "bcrypt", + "brotli", "bytes", "chrono", "chrono-tz", @@ -346,6 +347,7 @@ dependencies = [ "reqwest", "rsa", "rustls 0.23.37", + "semver", "serde", "serde_json", "sha1", diff --git a/Cargo.toml b/Cargo.toml index 3ba5c5036..ae184e1cb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -106,6 +106,7 @@ async-trait = "0.1" axum = "0.8" base64 = "0.22" bcrypt = "0.16" +brotli = "8" bytes = "1" cbc = "0.1" chrono = { version = "0.4", features = ["serde"] } diff --git a/apps/aether-gateway/Cargo.toml b/apps/aether-gateway/Cargo.toml index a786104cf..2488d6e5f 100644 --- a/apps/aether-gateway/Cargo.toml +++ b/apps/aether-gateway/Cargo.toml @@ -52,6 +52,7 @@ async-trait.workspace = true axum = { version = "0.8", features = ["ws"] } base64.workspace = true bcrypt.workspace = true +brotli.workspace = true bytes.workspace = true chrono.workspace = true chrono-tz.workspace = true @@ -75,6 +76,7 @@ rsa = "0.9.10" rustls.workspace = true serde.workspace = true serde_json.workspace = true +semver.workspace = true sha1 = "0.10" sha2 = { workspace = true, features = ["oid"] } socket2.workspace = true diff --git a/apps/aether-gateway/src/ai_serving/planner/decision_input.rs b/apps/aether-gateway/src/ai_serving/planner/decision_input.rs index 1ac6268eb..007eeeb06 100644 --- a/apps/aether-gateway/src/ai_serving/planner/decision_input.rs +++ b/apps/aether-gateway/src/ai_serving/planner/decision_input.rs @@ -339,7 +339,7 @@ fn apply_codex_oauth_fingerprint_convergence_to_decision( else { return; }; - crate::provider_transport::apply_codex_oauth_fingerprint_convergence( + crate::ai_serving::transport::apply_codex_oauth_fingerprint_convergence( transport, provider_api_format, input.original_client_session_id.as_deref(), diff --git a/apps/aether-gateway/src/ai_serving/pure/mod.rs b/apps/aether-gateway/src/ai_serving/pure/mod.rs index 64523a428..88deedf85 100644 --- a/apps/aether-gateway/src/ai_serving/pure/mod.rs +++ b/apps/aether-gateway/src/ai_serving/pure/mod.rs @@ -82,10 +82,10 @@ pub(crate) use aether_ai_formats::api::{ parse_codex_auth_identity, parse_direct_request_body, parse_model_directive, parse_model_directive_with_suffixes, parse_openai_stop_sequences, parse_openai_tool_result_content, prepare_local_success_response_parts, - prepare_local_success_response_parts_owned, project_codex_openai_image_api_request_body, - project_openai_image_api_request_body, provider_adaptation_allows_sync_finalize_envelope, - provider_adaptation_anchor_api_format, provider_adaptation_descriptor_for_envelope, - provider_adaptation_descriptor_for_provider_type, + prepare_local_success_response_parts_owned, project_codex_catalog_model_card, + project_codex_openai_image_api_request_body, project_openai_image_api_request_body, + provider_adaptation_allows_sync_finalize_envelope, provider_adaptation_anchor_api_format, + provider_adaptation_descriptor_for_envelope, provider_adaptation_descriptor_for_provider_type, provider_adaptation_requires_eventstream_accept, provider_adaptation_should_unwrap_stream_envelope, provider_private_response_allows_sync_finalize, record_converted_response_history, @@ -175,6 +175,7 @@ pub(crate) use aether_ai_formats::{ is_rerank_api_format, openai_responses_request_operation, openai_responses_synthetic_reasoning_item_id, strip_incompatible_openai_responses_reasoning_items, ApiOperation, ClientSurface, + CODEX_CLIENT_VERSION, }; pub(crate) fn plan_kind_matches_api_operation( diff --git a/apps/aether-gateway/src/ai_serving/transport.rs b/apps/aether-gateway/src/ai_serving/transport.rs index a7eb7d5cb..e5267484b 100644 --- a/apps/aether-gateway/src/ai_serving/transport.rs +++ b/apps/aether-gateway/src/ai_serving/transport.rs @@ -59,8 +59,9 @@ pub(crate) mod windsurf { } pub(crate) use aether_provider_transport::{ - append_transport_diagnostics_to_value, apply_local_auth_config_header_overrides, - apply_local_body_rules, apply_local_body_rules_with_request_headers, apply_local_header_rules, + append_transport_diagnostics_to_value, apply_codex_oauth_fingerprint_convergence, + apply_local_auth_config_header_overrides, apply_local_body_rules, + apply_local_body_rules_with_request_headers, apply_local_header_rules, apply_local_header_rules_with_request_headers, apply_standard_provider_request_body_rules, apply_standard_provider_request_body_rules_with_request_headers, apply_transport_request_body_semantics, body_rules_are_locally_supported, diff --git a/apps/aether-gateway/src/backup/worker.rs b/apps/aether-gateway/src/backup/worker.rs index e87865107..9fe12d156 100644 --- a/apps/aether-gateway/src/backup/worker.rs +++ b/apps/aether-gateway/src/backup/worker.rs @@ -85,29 +85,17 @@ async fn read_last_backup_slot(app: &AppState) -> Result, Gateway #[cfg(test)] mod tests { - use crate::backup::schedule::{BackupSchedule, BackupScheduleUnit}; use crate::task_runtime::{task_definition, TASK_KEY_SYSTEM_S3_BACKUP}; #[test] fn backup_worker_skips_already_recorded_slot() { - let schedule = BackupSchedule { - unit: BackupScheduleUnit::Days, - interval: 1, - minute: 0, - hour: 3, - weekday: 1, - month_day: 1, - }; - let now = chrono::DateTime::parse_from_rfc3339("2026-05-24T03:00:30+08:00") - .unwrap() - .with_timezone(&chrono::Utc); - let slot = schedule.due_slot(now).expect("slot should be due"); + let slot = "days:2026-05-23T19:00:00Z"; assert!(super::should_start_scheduled_backup( Some("days:2026-05-22T19:00:00Z"), - &slot + slot )); - assert!(!super::should_start_scheduled_backup(Some(&slot), &slot)); + assert!(!super::should_start_scheduled_backup(Some(slot), slot)); } #[test] diff --git a/apps/aether-gateway/src/cache/auth_runtime.rs b/apps/aether-gateway/src/cache/auth_runtime.rs index 63531106b..0212e7028 100644 --- a/apps/aether-gateway/src/cache/auth_runtime.rs +++ b/apps/aether-gateway/src/cache/auth_runtime.rs @@ -1274,7 +1274,6 @@ mod tests { let first_cache = Arc::clone(&cache); let first_key = key.clone(); let first_calls = Arc::clone(&calls); - let first_started = Instant::now(); let first = tokio::spawn(async move { first_cache .get_or_load_once_stale_while_refreshing::<(), _, _>( @@ -1292,7 +1291,6 @@ mod tests { }); let follower_cache = Arc::clone(&cache); - let follower_started = Instant::now(); let follower_calls = Arc::clone(&calls); let follower = tokio::spawn(async move { follower_cache @@ -1310,15 +1308,7 @@ mod tests { }); assert_eq!(first.await.unwrap().unwrap(), Some(1)); - assert!( - first_started.elapsed() < Duration::from_millis(80), - "stale value should not wait for request-path refresh" - ); assert_eq!(follower.await.unwrap().unwrap(), Some(1)); - assert!( - follower_started.elapsed() < Duration::from_millis(80), - "follower should return stale value without waiting for refresh" - ); assert_eq!(calls.load(Ordering::Acquire), 0); } diff --git a/apps/aether-gateway/src/control/route/mod.rs b/apps/aether-gateway/src/control/route/mod.rs index c18db0f34..529986899 100644 --- a/apps/aether-gateway/src/control/route/mod.rs +++ b/apps/aether-gateway/src/control/route/mod.rs @@ -247,8 +247,7 @@ pub(super) fn detect_public_models_auth_signature(uri: &Uri, headers: &http::Hea 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()) + url::form_urlencoded::parse(query.as_bytes()).any(|(key, _)| key == "client_version") }); if has_codex_client_version { return "openai:responses".to_string(); diff --git a/apps/aether-gateway/src/control/tests/public_support.rs b/apps/aether-gateway/src/control/tests/public_support.rs index 5f1ea8c6e..078f47a2b 100644 --- a/apps/aether-gateway/src/control/tests/public_support.rs +++ b/apps/aether-gateway/src/control/tests/public_support.rs @@ -39,7 +39,7 @@ fn classifies_codex_models_list_with_responses_auth_signature() { } #[test] -fn empty_codex_client_version_keeps_standard_openai_models_signature() { +fn empty_codex_client_version_uses_responses_signature_for_bounded_fallback() { let headers = headers(&[("authorization", "Bearer sk-test")]); let uri: Uri = "/v1/models?client_version=" .parse() @@ -49,7 +49,7 @@ fn empty_codex_client_version_keeps_standard_openai_models_signature() { assert_eq!( decision.auth_endpoint_signature.as_deref(), - Some("openai:chat") + Some("openai:responses") ); } diff --git a/apps/aether-gateway/src/data/state/catalog.rs b/apps/aether-gateway/src/data/state/catalog.rs index a31fa0583..6e1aa6dbf 100644 --- a/apps/aether-gateway/src/data/state/catalog.rs +++ b/apps/aether-gateway/src/data/state/catalog.rs @@ -282,6 +282,16 @@ impl GatewayDataState { } } + pub(crate) async fn list_provider_catalog_keys_by_ids_strong( + &self, + key_ids: &[String], + ) -> Result, DataLayerError> { + match &self.provider_catalog_reader { + Some(repository) => repository.list_keys_by_ids_strong(key_ids).await, + None => Ok(Vec::new()), + } + } + pub(crate) async fn list_provider_catalog_keys_by_provider_ids( &self, provider_ids: &[String], diff --git a/apps/aether-gateway/src/data/state/provider_catalog_cache.rs b/apps/aether-gateway/src/data/state/provider_catalog_cache.rs index 48d1bdb3f..c42a1f977 100644 --- a/apps/aether-gateway/src/data/state/provider_catalog_cache.rs +++ b/apps/aether-gateway/src/data/state/provider_catalog_cache.rs @@ -334,6 +334,13 @@ impl ProviderCatalogReadRepository for CachedProviderCatalogReadRepository { } } + async fn list_keys_by_ids_strong( + &self, + key_ids: &[String], + ) -> Result, DataLayerError> { + self.inner.list_keys_by_ids_strong(key_ids).await + } + async fn list_keys_by_provider_ids( &self, provider_ids: &[String], @@ -538,6 +545,7 @@ fn normalize_ids(ids: &[String]) -> Vec { mod tests { use super::*; use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository; + use aether_data_contracts::repository::provider_catalog::ProviderCatalogWriteRepository; fn cache() -> CachedProviderCatalogReadRepository { CachedProviderCatalogReadRepository::new(Arc::new( @@ -555,6 +563,55 @@ mod tests { .expect("provider should be valid") } + #[tokio::test] + async fn provider_catalog_strong_key_read_bypasses_fresh_cached_generation() { + let old_metadata = serde_json::json!({ + "codex": {"credential_generation": "old"} + }); + let mut key = StoredProviderCatalogKey::new( + "key-1".to_string(), + "provider-1".to_string(), + "key-1".to_string(), + "oauth".to_string(), + None, + true, + ) + .expect("key should be valid"); + key.upstream_metadata = Some(old_metadata.clone()); + let inner = Arc::new(InMemoryProviderCatalogReadRepository::seed( + vec![provider("provider-1")], + Vec::new(), + vec![key], + )); + let cache = CachedProviderCatalogReadRepository::new(inner.clone()); + let key_ids = vec!["key-1".to_string()]; + + let first = cache + .list_keys_by_ids(&key_ids) + .await + .expect("initial key read should succeed"); + assert_eq!(first[0].upstream_metadata.as_ref(), Some(&old_metadata)); + + let new_metadata = serde_json::json!({ + "codex": {"credential_generation": "new"} + }); + assert!(inner + .upsert_key_upstream_metadata_namespace("key-1", "codex", &new_metadata["codex"], None,) + .await + .expect("inner metadata update should succeed")); + + let cached = cache + .list_keys_by_ids(&key_ids) + .await + .expect("cached key read should succeed"); + assert_eq!(cached[0].upstream_metadata.as_ref(), Some(&old_metadata)); + let strong = cache + .list_keys_by_ids_strong(&key_ids) + .await + .expect("strong key read should succeed"); + assert_eq!(strong[0].upstream_metadata.as_ref(), Some(&new_metadata)); + } + #[tokio::test] async fn provider_catalog_follower_observes_completion_before_first_poll() { let cache = cache(); diff --git a/apps/aether-gateway/src/execution_runtime/sync/execution.rs b/apps/aether-gateway/src/execution_runtime/sync/execution.rs index fe30181f7..0f3d36e3a 100644 --- a/apps/aether-gateway/src/execution_runtime/sync/execution.rs +++ b/apps/aether-gateway/src/execution_runtime/sync/execution.rs @@ -55,8 +55,9 @@ use crate::execution_runtime::submission::{ resolve_local_sync_error_status_code, submit_local_core_error_or_sync_finalize, }; use crate::execution_runtime::transport::{ - append_upstream_response_body_chunk, build_execution_response_body, build_request_body, - collect_response_headers, decode_response_body_bytes, execution_response_body_mode, + append_upstream_response_body_chunk_with_limit, build_execution_response_body, + build_request_body, collect_response_headers, decode_response_body_bytes_with_limit, + execution_plan_response_body_limit_bytes, execution_response_body_mode, format_hyper_error_chain, format_upstream_request_error, format_wreq_upstream_request_error, response_body_is_json, send_request, DirectHttpResponse, DirectSyncExecutionRuntime, ExecutionRuntimeTransportError, @@ -1491,6 +1492,7 @@ async fn execute_openai_image_sync_upstream_sse_candidate( progress_snapshot: Option>>, ) -> Result { let request_body = build_request_body(plan).map_err(SyncExecutionFailure::from_transport)?; + let response_body_limit_bytes = execution_plan_response_body_limit_bytes(plan); let started_at = Instant::now(); let mut progress = OpenAiImageSyncProgressRecorder::new(state, plan, report_context, progress_snapshot); @@ -1529,8 +1531,12 @@ async fn execute_openai_image_sync_upstream_sse_candidate( ), ) })?; - append_upstream_response_body_chunk(&mut body_bytes, &chunk) - .map_err(SyncExecutionFailure::from_transport)?; + append_upstream_response_body_chunk_with_limit( + &mut body_bytes, + &chunk, + response_body_limit_bytes, + ) + .map_err(SyncExecutionFailure::from_transport)?; let elapsed_ms = started_at.elapsed().as_millis() as u64; progress .observe_chunk(&chunk, status_code, elapsed_ms) @@ -1547,8 +1553,12 @@ async fn execute_openai_image_sync_upstream_sse_candidate( )), ) })?; - append_upstream_response_body_chunk(&mut body_bytes, &chunk) - .map_err(SyncExecutionFailure::from_transport)?; + append_upstream_response_body_chunk_with_limit( + &mut body_bytes, + &chunk, + response_body_limit_bytes, + ) + .map_err(SyncExecutionFailure::from_transport)?; let elapsed_ms = started_at.elapsed().as_millis() as u64; progress .observe_chunk(&chunk, status_code, elapsed_ms) @@ -1565,8 +1575,12 @@ async fn execute_openai_image_sync_upstream_sse_candidate( ), ) })?; - append_upstream_response_body_chunk(&mut body_bytes, &chunk) - .map_err(SyncExecutionFailure::from_transport)?; + append_upstream_response_body_chunk_with_limit( + &mut body_bytes, + &chunk, + response_body_limit_bytes, + ) + .map_err(SyncExecutionFailure::from_transport)?; let elapsed_ms = started_at.elapsed().as_millis() as u64; progress .observe_chunk(&chunk, status_code, elapsed_ms) @@ -1575,8 +1589,9 @@ async fn execute_openai_image_sync_upstream_sse_candidate( } } - let decoded_body_bytes = decode_response_body_bytes(&headers, &body_bytes) - .map_err(SyncExecutionFailure::from_transport)?; + let decoded_body_bytes = + decode_response_body_bytes_with_limit(&headers, &body_bytes, response_body_limit_bytes) + .map_err(SyncExecutionFailure::from_transport)?; let elapsed_ms = started_at.elapsed().as_millis() as u64; let upstream_bytes = body_bytes.len() as u64; progress.finish(status_code, elapsed_ms).await; diff --git a/apps/aether-gateway/src/execution_runtime/transport.rs b/apps/aether-gateway/src/execution_runtime/transport.rs index e17b394e3..8df7018d7 100644 --- a/apps/aether-gateway/src/execution_runtime/transport.rs +++ b/apps/aether-gateway/src/execution_runtime/transport.rs @@ -21,6 +21,7 @@ use aether_http::{apply_http_client_config, HttpClientConfig}; use aether_runtime::{MetricKind, MetricSample}; use axum::body::Bytes; use base64::Engine as _; +use brotli::Decompressor as BrotliDecoder; use flate2::read::{DeflateDecoder, GzDecoder}; use flate2::write::GzEncoder; use flate2::Compression; @@ -62,6 +63,10 @@ const DEFAULT_STREAM_FIRST_BYTE_TIMEOUT_MS: u64 = 30_000; const DEFAULT_NON_STREAM_TOTAL_TIMEOUT_MS: u64 = 300_000; const DEFAULT_CODEX_COMPACT_TOTAL_TIMEOUT_MS: u64 = 1_200_000; const MIN_TUNNEL_TIMEOUT_SECS: u64 = 1; +const EXECUTION_RESPONSE_BODY_LIMIT_HEADER: &str = "x-aether-execution-response-body-limit-bytes"; +const DEFAULT_SCOPED_RESPONSE_BODY_LIMIT_BYTES: usize = 8 * 1024 * 1024; +const MIN_SCOPED_RESPONSE_BODY_LIMIT_BYTES: usize = 64 * 1024; +const MAX_SCOPED_RESPONSE_BODY_LIMIT_BYTES: usize = 64 * 1024 * 1024; const DIRECT_REQWEST_H2_CLIENT_SHARDS_ENV: &str = "AETHER_GATEWAY_DIRECT_REQWEST_H2_CLIENT_SHARDS"; const DIRECT_REQWEST_CLIENT_SHARDS_ENV: &str = "AETHER_GATEWAY_DIRECT_REQWEST_CLIENT_SHARDS"; const DIRECT_REQWEST_H2_TARGET_STREAMS_PER_CLIENT_ENV: &str = @@ -607,6 +612,56 @@ impl std::fmt::Display for UpstreamResponseBodyPhase { } } +pub(crate) fn with_upstream_response_body_limit( + plan: &ExecutionPlan, + limit_bytes: usize, +) -> ExecutionPlan { + let mut bounded_plan = plan.clone(); + bounded_plan + .headers + .retain(|name, _| !name.eq_ignore_ascii_case(EXECUTION_RESPONSE_BODY_LIMIT_HEADER)); + bounded_plan.headers.insert( + EXECUTION_RESPONSE_BODY_LIMIT_HEADER.to_string(), + normalize_scoped_response_body_limit(limit_bytes) + .unwrap_or(DEFAULT_SCOPED_RESPONSE_BODY_LIMIT_BYTES) + .to_string(), + ); + bounded_plan +} + +pub(crate) fn execution_plan_response_body_limit_bytes(plan: &ExecutionPlan) -> usize { + effective_response_body_limit_bytes( + execution_transport_header_value(&plan.headers, EXECUTION_RESPONSE_BODY_LIMIT_HEADER), + crate::headers::max_internal_buffered_body_bytes(), + ) +} + +fn effective_response_body_limit_bytes( + raw_scoped_limit: Option<&str>, + global_limit: usize, +) -> usize { + let Some(raw_scoped_limit) = raw_scoped_limit else { + return global_limit; + }; + parse_scoped_response_body_limit(raw_scoped_limit) + .unwrap_or(DEFAULT_SCOPED_RESPONSE_BODY_LIMIT_BYTES) + .min(global_limit) +} + +fn parse_scoped_response_body_limit(value: &str) -> Option { + let raw_limit = value.trim().parse::().ok()?; + usize::try_from(raw_limit) + .ok() + .and_then(normalize_scoped_response_body_limit) +} + +fn normalize_scoped_response_body_limit(limit_bytes: usize) -> Option { + (limit_bytes > 0).then_some(limit_bytes.clamp( + MIN_SCOPED_RESPONSE_BODY_LIMIT_BYTES, + MAX_SCOPED_RESPONSE_BODY_LIMIT_BYTES, + )) +} + pub(crate) fn append_upstream_response_body_chunk( body: &mut Vec, chunk: &[u8], @@ -618,7 +673,7 @@ pub(crate) fn append_upstream_response_body_chunk( ) } -fn append_upstream_response_body_chunk_with_limit( +pub(crate) fn append_upstream_response_body_chunk_with_limit( body: &mut Vec, chunk: &[u8], limit_bytes: usize, @@ -724,6 +779,7 @@ impl DirectSyncExecutionRuntime { F: FnOnce(DirectSyncResponseStarted), { let body_bytes = build_request_body(plan)?; + let response_body_limit_bytes = execution_plan_response_body_limit_bytes(plan); let started_at = Instant::now(); let request_started_at_unix_ms = crate::clock::current_unix_ms(); @@ -744,9 +800,14 @@ impl DirectSyncExecutionRuntime { ttfb_ms, response_observation: response_observation.clone(), }); - let (body_bytes, stream_ttfb_ms) = - response.bytes_with_stream_timeout(plan, started_at).await?; - let decoded_body_bytes = decode_response_body_bytes(&headers, &body_bytes)?; + let (body_bytes, stream_ttfb_ms) = response + .bytes_with_stream_timeout(plan, started_at, response_body_limit_bytes) + .await?; + let decoded_body_bytes = decode_response_body_bytes_with_limit( + &headers, + &body_bytes, + response_body_limit_bytes, + )?; let elapsed_ms = started_at.elapsed().as_millis() as u64; let upstream_bytes = body_bytes.len() as u64; @@ -1058,6 +1119,7 @@ async fn execute_sync_plan_via_local_tunnel_inner( } let body_bytes = build_request_body(plan)?; + let response_body_limit_bytes = execution_plan_response_body_limit_bytes(plan); let transport_controls = resolve_execution_transport_controls(&plan.headers); let headers = build_request_headers( &plan.headers, @@ -1113,8 +1175,10 @@ async fn execute_sync_plan_via_local_tunnel_inner( ); let proxy_timing = execution_header_for_log(&headers, "x-proxy-timing").unwrap_or("-"); let (body_bytes, stream_ttfb_ms) = - collect_local_tunnel_response_body(response, plan, started_at).await?; - let decoded_body_bytes = decode_response_body_bytes(&headers, &body_bytes)?; + collect_local_tunnel_response_body(response, plan, started_at, response_body_limit_bytes) + .await?; + let decoded_body_bytes = + decode_response_body_bytes_with_limit(&headers, &body_bytes, response_body_limit_bytes)?; let elapsed_ms = started_at.elapsed().as_millis() as u64; let upstream_bytes = body_bytes.len() as u64; if status_code >= 400 { @@ -1179,6 +1243,7 @@ async fn collect_local_tunnel_response_body( mut response: tunnel::DirectRelayResponse, plan: &ExecutionPlan, started_at: Instant, + response_body_limit_bytes: usize, ) -> Result<(Vec, Option), ExecutionRuntimeTransportError> { let mut body_bytes = Vec::new(); let mut first_byte_ms = None; @@ -1201,7 +1266,11 @@ async fn collect_local_tunnel_response_body( if plan.stream && first_byte_ms.is_none() && !chunk.is_empty() { first_byte_ms = Some(started_at.elapsed().as_millis() as u64); } - append_upstream_response_body_chunk(&mut body_bytes, &chunk)?; + append_upstream_response_body_chunk_with_limit( + &mut body_bytes, + &chunk, + response_body_limit_bytes, + )?; } Ok((body_bytes, first_byte_ms)) @@ -1362,20 +1431,28 @@ impl DirectHttpResponse { } pub(crate) async fn bytes(self) -> Result { + self.bytes_with_limit(crate::headers::max_internal_buffered_body_bytes()) + .await + } + + async fn bytes_with_limit( + self, + response_body_limit_bytes: usize, + ) -> Result { let started_at = Instant::now(); match self { DirectHttpResponse::Reqwest(response) => { - collect_reqwest_stream_body(response, started_at, None) + collect_reqwest_stream_body(response, started_at, None, response_body_limit_bytes) .await .map(|(body, _)| body) } DirectHttpResponse::HyperH2c(response) => { - collect_hyper_stream_body(response, started_at, None) + collect_hyper_stream_body(response, started_at, None, response_body_limit_bytes) .await .map(|(body, _)| body) } DirectHttpResponse::BrowserWreq(response) => { - collect_wreq_stream_body(response, started_at, None) + collect_wreq_stream_body(response, started_at, None, response_body_limit_bytes) .await .map(|(body, _)| body) } @@ -1386,21 +1463,43 @@ impl DirectHttpResponse { self, plan: &ExecutionPlan, started_at: Instant, + response_body_limit_bytes: usize, ) -> Result<(Bytes, Option), ExecutionRuntimeTransportError> { if !plan.stream { - return self.bytes().await.map(|bytes| (bytes, None)); + return self + .bytes_with_limit(response_body_limit_bytes) + .await + .map(|bytes| (bytes, None)); } let first_byte_timeout = resolve_stream_first_byte_timeout(plan); match self { DirectHttpResponse::Reqwest(response) => { - collect_reqwest_stream_body(response, started_at, first_byte_timeout).await + collect_reqwest_stream_body( + response, + started_at, + first_byte_timeout, + response_body_limit_bytes, + ) + .await } DirectHttpResponse::HyperH2c(response) => { - collect_hyper_stream_body(response, started_at, first_byte_timeout).await + collect_hyper_stream_body( + response, + started_at, + first_byte_timeout, + response_body_limit_bytes, + ) + .await } DirectHttpResponse::BrowserWreq(response) => { - collect_wreq_stream_body(response, started_at, first_byte_timeout).await + collect_wreq_stream_body( + response, + started_at, + first_byte_timeout, + response_body_limit_bytes, + ) + .await } } } @@ -1446,6 +1545,7 @@ async fn collect_reqwest_stream_body( response: reqwest::Response, started_at: Instant, first_byte_timeout: Option, + response_body_limit_bytes: usize, ) -> Result<(Bytes, Option), ExecutionRuntimeTransportError> { let mut stream = response.bytes_stream(); let mut body_bytes = Vec::new(); @@ -1466,7 +1566,11 @@ async fn collect_reqwest_stream_body( if first_byte_ms.is_none() && !chunk.is_empty() { first_byte_ms = Some(started_at.elapsed().as_millis() as u64); } - append_upstream_response_body_chunk(&mut body_bytes, &chunk)?; + append_upstream_response_body_chunk_with_limit( + &mut body_bytes, + &chunk, + response_body_limit_bytes, + )?; } Ok((Bytes::from(body_bytes), first_byte_ms)) @@ -1476,6 +1580,7 @@ async fn collect_hyper_stream_body( response: hyper::Response, started_at: Instant, first_byte_timeout: Option, + response_body_limit_bytes: usize, ) -> Result<(Bytes, Option), ExecutionRuntimeTransportError> { let mut stream = response.into_body().into_data_stream(); let mut body_bytes = Vec::new(); @@ -1496,7 +1601,11 @@ async fn collect_hyper_stream_body( if first_byte_ms.is_none() && !chunk.is_empty() { first_byte_ms = Some(started_at.elapsed().as_millis() as u64); } - append_upstream_response_body_chunk(&mut body_bytes, &chunk)?; + append_upstream_response_body_chunk_with_limit( + &mut body_bytes, + &chunk, + response_body_limit_bytes, + )?; } Ok((Bytes::from(body_bytes), first_byte_ms)) @@ -1506,6 +1615,7 @@ async fn collect_wreq_stream_body( response: wreq::Response, started_at: Instant, first_byte_timeout: Option, + response_body_limit_bytes: usize, ) -> Result<(Bytes, Option), ExecutionRuntimeTransportError> { let mut stream = response.bytes_stream(); let mut body_bytes = Vec::new(); @@ -1526,7 +1636,11 @@ async fn collect_wreq_stream_body( if first_byte_ms.is_none() && !chunk.is_empty() { first_byte_ms = Some(started_at.elapsed().as_millis() as u64); } - append_upstream_response_body_chunk(&mut body_bytes, &chunk)?; + append_upstream_response_body_chunk_with_limit( + &mut body_bytes, + &chunk, + response_body_limit_bytes, + )?; } Ok((Bytes::from(body_bytes), first_byte_ms)) @@ -2378,10 +2492,31 @@ async fn send_via_tunnel_relay( error_kind = %kind, "gateway execution runtime tunnel relay returned relay error" ); - let message = response - .text() - .await - .unwrap_or_else(|_| format!("hub relay error: {kind}")); + let response_headers = collect_response_headers(response.headers()); + let response_body_limit_bytes = execution_plan_response_body_limit_bytes(plan); + let (wire_body, _) = + collect_reqwest_stream_body(response, Instant::now(), None, response_body_limit_bytes) + .await + .map_err(|error| { + ExecutionRuntimeTransportError::RelayError(format!( + "hub relay error: {kind}: bounded error body read failed: {error}" + )) + })?; + let decoded_body = decode_response_body_bytes_with_limit( + &response_headers, + &wire_body, + response_body_limit_bytes, + ) + .map_err(|error| { + ExecutionRuntimeTransportError::RelayError(format!( + "hub relay error: {kind}: bounded error body decode failed: {error}" + )) + })?; + let message = if decoded_body.is_empty() { + format!("hub relay error: {kind}") + } else { + String::from_utf8_lossy(decoded_body.as_ref()).into_owned() + }; return Err(ExecutionRuntimeTransportError::RelayError(message)); } @@ -3903,6 +4038,7 @@ pub(crate) fn build_request_headers( || normalized_key == EXECUTION_REQUEST_HTTP1_ONLY_HEADER || normalized_key == EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER || normalized_key == EXECUTION_RESPONSE_BODY_MODE_HEADER + || normalized_key == EXECUTION_RESPONSE_BODY_LIMIT_HEADER { continue; } @@ -4051,7 +4187,7 @@ pub(crate) fn decode_response_body_bytes<'a>( ) } -fn decode_response_body_bytes_with_limit<'a>( +pub(crate) fn decode_response_body_bytes_with_limit<'a>( headers: &BTreeMap, body_bytes: &'a [u8], limit_bytes: usize, @@ -4073,6 +4209,11 @@ fn decode_response_body_bytes_with_limit<'a>( read_upstream_response_decoder_with_limit("deflate", &mut decoder, limit_bytes) .map(Cow::Owned) } + Some("br") => { + let mut decoder = BrotliDecoder::new(body_bytes, 4_096); + read_upstream_response_decoder_with_limit("br", &mut decoder, limit_bytes) + .map(Cow::Owned) + } _ => Ok(Cow::Borrowed(body_bytes)), } } @@ -4192,12 +4333,16 @@ mod tests { use super::{ append_upstream_response_body_chunk_with_limit, build_browser_wreq_client, build_client, build_direct_tunnel_request_meta, build_execution_response_body, build_request_headers, - decode_response_body_bytes_with_limit, execute_sync_plan, execution_response_body_mode, + decode_response_body_bytes_with_limit, effective_response_body_limit_bytes, + execute_sync_plan, execution_plan_response_body_limit_bytes, execution_response_body_mode, record_manual_proxy_request_failure, record_manual_proxy_request_outcome, record_manual_proxy_request_success, record_manual_proxy_stream_error, resolve_execution_transport_controls, resolve_non_stream_total_timeout, - resolve_stream_first_byte_timeout, response_body_is_json, DirectSyncExecutionRuntime, + resolve_stream_first_byte_timeout, response_body_is_json, + with_upstream_response_body_limit, DirectSyncExecutionRuntime, ExecutionRuntimeTransportError, ExecutionTransportControls, UpstreamResponseBodyPhase, + DEFAULT_SCOPED_RESPONSE_BODY_LIMIT_BYTES, EXECUTION_RESPONSE_BODY_LIMIT_HEADER, + MAX_SCOPED_RESPONSE_BODY_LIMIT_BYTES, MIN_SCOPED_RESPONSE_BODY_LIMIT_BYTES, }; use crate::constants::{ EXECUTION_RUNTIME_LOOP_GUARD_HEADER, EXECUTION_RUNTIME_LOOP_GUARD_VIA_TOKEN, @@ -4252,6 +4397,162 @@ mod tests { assert!(!materialized.contains_key("x-aether-future-control")); } + #[test] + fn scoped_response_body_limit_injection_preserves_transport_profile_and_extra() { + let mut plan = tunnel_timeout_plan(false); + let original_profile = ResolvedTransportProfile { + profile_id: "existing-profile".into(), + backend: TRANSPORT_BACKEND_BROWSER_WREQ.into(), + http_mode: TRANSPORT_HTTP_MODE_HTTP1_ONLY.into(), + pool_scope: "provider".into(), + header_fingerprint: Some(json!({"user_agent": "existing"})), + extra: Some(json!({"existing": {"nested": true}})), + }; + plan.transport_profile = Some(original_profile.clone()); + + let bounded_plan = + with_upstream_response_body_limit(&plan, DEFAULT_SCOPED_RESPONSE_BODY_LIMIT_BYTES); + + assert_eq!(plan.transport_profile, Some(original_profile.clone())); + assert_eq!(bounded_plan.transport_profile, Some(original_profile)); + assert_eq!( + bounded_plan + .headers + .get(EXECUTION_RESPONSE_BODY_LIMIT_HEADER) + .and_then(|value| value.parse::().ok()), + Some(DEFAULT_SCOPED_RESPONSE_BODY_LIMIT_BYTES) + ); + assert_eq!( + execution_plan_response_body_limit_bytes(&bounded_plan), + DEFAULT_SCOPED_RESPONSE_BODY_LIMIT_BYTES + ); + + let unprofiled_plan = tunnel_timeout_plan(false); + let bounded_unprofiled_plan = with_upstream_response_body_limit( + &unprofiled_plan, + DEFAULT_SCOPED_RESPONSE_BODY_LIMIT_BYTES, + ); + assert!(unprofiled_plan.transport_profile.is_none()); + assert!(bounded_unprofiled_plan.transport_profile.is_none()); + assert_eq!( + execution_plan_response_body_limit_bytes(&bounded_unprofiled_plan), + DEFAULT_SCOPED_RESPONSE_BODY_LIMIT_BYTES + ); + + let mut shadowed_plan = tunnel_timeout_plan(false); + shadowed_plan.headers.insert( + EXECUTION_RESPONSE_BODY_LIMIT_HEADER.to_ascii_uppercase(), + "65536".to_string(), + ); + let bounded_shadowed_plan = with_upstream_response_body_limit( + &shadowed_plan, + DEFAULT_SCOPED_RESPONSE_BODY_LIMIT_BYTES, + ); + assert_eq!( + bounded_shadowed_plan + .headers + .keys() + .filter(|name| name.eq_ignore_ascii_case(EXECUTION_RESPONSE_BODY_LIMIT_HEADER)) + .count(), + 1 + ); + } + + #[test] + fn scoped_response_body_limit_parsing_rejects_invalid_values_and_clamps_bounds() { + let scoped_plan = |raw_limit: &str| { + let mut plan = tunnel_timeout_plan(false); + plan.headers.insert( + EXECUTION_RESPONSE_BODY_LIMIT_HEADER.to_string(), + raw_limit.to_string(), + ); + plan + }; + + for invalid in ["0", "-1", "1.5", "", "invalid"] { + assert_eq!( + execution_plan_response_body_limit_bytes(&scoped_plan(invalid)), + DEFAULT_SCOPED_RESPONSE_BODY_LIMIT_BYTES + ); + } + assert_eq!( + execution_plan_response_body_limit_bytes(&scoped_plan("1")), + MIN_SCOPED_RESPONSE_BODY_LIMIT_BYTES + ); + assert_eq!( + execution_plan_response_body_limit_bytes(&scoped_plan( + &(MAX_SCOPED_RESPONSE_BODY_LIMIT_BYTES as u64 + 1).to_string() + )), + MAX_SCOPED_RESPONSE_BODY_LIMIT_BYTES + ); + assert_eq!( + execution_plan_response_body_limit_bytes(&scoped_plan("1048576")), + 1_048_576 + ); + assert_eq!( + effective_response_body_limit_bytes( + Some(&(DEFAULT_SCOPED_RESPONSE_BODY_LIMIT_BYTES * 2).to_string()), + 1024 * 1024, + ), + 1024 * 1024, + "a scoped limit must never raise the operator's global cap" + ); + } + + #[test] + fn scoped_response_body_wire_limit_rejects_overflow() { + let bounded_plan = with_upstream_response_body_limit( + &tunnel_timeout_plan(false), + MIN_SCOPED_RESPONSE_BODY_LIMIT_BYTES, + ); + let limit_bytes = execution_plan_response_body_limit_bytes(&bounded_plan); + let mut body = vec![b'x'; limit_bytes]; + + let error = + append_upstream_response_body_chunk_with_limit(&mut body, b"overflow", limit_bytes) + .expect_err("wire body above the plan-scoped limit should fail"); + + assert!(matches!( + error, + ExecutionRuntimeTransportError::UpstreamResponseTooLarge { + phase: UpstreamResponseBodyPhase::Wire, + limit_bytes: MIN_SCOPED_RESPONSE_BODY_LIMIT_BYTES, + } + )); + } + + #[test] + fn scoped_response_body_limit_rejects_gzip_bomb_after_wire_check() { + let bounded_plan = with_upstream_response_body_limit( + &tunnel_timeout_plan(false), + MIN_SCOPED_RESPONSE_BODY_LIMIT_BYTES, + ); + let limit_bytes = execution_plan_response_body_limit_bytes(&bounded_plan); + let payload = vec![b'x'; limit_bytes + 1]; + let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default()); + encoder + .write_all(&payload) + .expect("gzip payload should encode"); + let encoded = encoder.finish().expect("gzip payload should finish"); + assert!(encoded.len() < limit_bytes); + + let mut wire_body = Vec::new(); + append_upstream_response_body_chunk_with_limit(&mut wire_body, &encoded, limit_bytes) + .expect("compressed wire body should fit within the plan-scoped limit"); + let headers = BTreeMap::from([("content-encoding".to_string(), "gzip".to_string())]); + + let error = decode_response_body_bytes_with_limit(&headers, &wire_body, limit_bytes) + .expect_err("decoded body above the plan-scoped limit should fail"); + + assert!(matches!( + error, + ExecutionRuntimeTransportError::UpstreamResponseTooLarge { + phase: UpstreamResponseBodyPhase::Decoded, + limit_bytes: MIN_SCOPED_RESPONSE_BODY_LIMIT_BYTES, + } + )); + } + #[test] fn upstream_response_wire_limit_allows_exact_body_and_rejects_next_byte() { let mut body = Vec::new(); diff --git a/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/codex/mod.rs b/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/codex/mod.rs index ea0f33418..b5c490fea 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/codex/mod.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/codex/mod.rs @@ -1241,7 +1241,7 @@ async fn refresh_codex_provider_quota_locally_with_reset_fence( continue; } let persisted_key = state - .read_provider_catalog_keys_by_ids(&[key.id.clone()]) + .read_provider_catalog_keys_by_ids(std::slice::from_ref(&key.id)) .await? .into_iter() .next(); diff --git a/apps/aether-gateway/src/handlers/public/mod.rs b/apps/aether-gateway/src/handlers/public/mod.rs index 9824be1d4..d666a3dd9 100644 --- a/apps/aether-gateway/src/handlers/public/mod.rs +++ b/apps/aether-gateway/src/handlers/public/mod.rs @@ -22,6 +22,8 @@ pub(crate) use self::system_modules_helpers::{ serialize_public_capability, supported_capability_names, PUBLIC_CAPABILITY_DEFINITIONS, }; +#[cfg(test)] +pub(crate) use self::support::filter_eligible_model_rows; pub(crate) use self::support::{ build_api_key_install_session_response, build_proxy_node_install_session_response, build_unhandled_public_support_response, matches_model_mapping_for_models, diff --git a/apps/aether-gateway/src/handlers/public/support.rs b/apps/aether-gateway/src/handlers/public/support.rs index 353f16ea3..4a14e2819 100644 --- a/apps/aether-gateway/src/handlers/public/support.rs +++ b/apps/aether-gateway/src/handlers/public/support.rs @@ -52,6 +52,8 @@ mod support_user_me; mod support_wallet; pub(crate) use self::support_announcements::maybe_build_local_admin_announcements_response; +#[cfg(test)] +pub(crate) use self::support_models::filter_eligible_model_rows; pub(crate) use self::support_models::matches_model_mapping_for_models; use self::support_announcements::{ diff --git a/apps/aether-gateway/src/handlers/public/support/models.rs b/apps/aether-gateway/src/handlers/public/support/models.rs index 0f6e07038..81196c8c2 100644 --- a/apps/aether-gateway/src/handlers/public/support/models.rs +++ b/apps/aether-gateway/src/handlers/public/support/models.rs @@ -10,6 +10,8 @@ mod models_route; mod models_shared; pub(crate) use self::models_responses::build_models_auth_error_response; +#[cfg(test)] +pub(crate) use self::models_shared::filter_eligible_model_rows; pub(crate) use self::models_shared::{matches_model_mapping_for_models, models_api_format}; pub(super) async fn maybe_build_local_models_response( diff --git a/apps/aether-gateway/src/handlers/public/support/models/responses.rs b/apps/aether-gateway/src/handlers/public/support/models/responses.rs index cb470ba19..e758173ea 100644 --- a/apps/aether-gateway/src/handlers/public/support/models/responses.rs +++ b/apps/aether-gateway/src/handlers/public/support/models/responses.rs @@ -102,8 +102,15 @@ pub(super) fn build_empty_models_list_response(api_format: &str) -> Response) -> Response { - Json(json!({ "models": models })).into_response() +pub(super) fn build_codex_models_list_response( + models: Vec, + etag: Option<&str>, +) -> Response { + let mut response = Json(json!({ "models": models })).into_response(); + if let Some(etag) = etag.and_then(|value| http::HeaderValue::from_str(value).ok()) { + response.headers_mut().insert(http::header::ETAG, etag); + } + response } pub(super) fn build_openai_models_list_response( diff --git a/apps/aether-gateway/src/handlers/public/support/models/route.rs b/apps/aether-gateway/src/handlers/public/support/models/route.rs index 6588fc50b..39b795646 100644 --- a/apps/aether-gateway/src/handlers/public/support/models/route.rs +++ b/apps/aether-gateway/src/handlers/public/support/models/route.rs @@ -1,4 +1,4 @@ -use std::collections::{BTreeMap, BTreeSet}; +use std::collections::BTreeSet; use std::fmt::Debug; use std::future::Future; use std::time::{Duration, SystemTime, UNIX_EPOCH}; @@ -25,8 +25,18 @@ use super::{query_param_value, AppState, GatewayPublicRequestContext}; #[cfg(not(test))] const MODELS_ROUTE_READ_TIMEOUT: Duration = Duration::from_secs(5); #[cfg(test)] -const MODELS_ROUTE_READ_TIMEOUT: Duration = Duration::from_millis(50); +const MODELS_ROUTE_READ_TIMEOUT: Duration = Duration::from_secs(1); const CODEX_MODELS_QUERY_API_FORMATS: &[&str] = &["openai:responses"]; +const CODEX_MODELS_MAX_RESPONSE_MODELS: usize = 512; +const CODEX_MODELS_MAX_RESPONSE_JSON_BYTES: usize = 8 * 1024 * 1024; + +fn codex_projected_catalog_fits_response_limits(cards: &[Value]) -> bool { + if cards.len() > CODEX_MODELS_MAX_RESPONSE_MODELS { + return false; + } + serde_json::to_vec(&serde_json::json!({ "models": cards })) + .is_ok_and(|body| body.len() <= CODEX_MODELS_MAX_RESPONSE_JSON_BYTES) +} async fn await_models_route_read(operation: &'static str, future: Fut) -> Option where @@ -115,124 +125,130 @@ fn is_codex_provider_row(row: &StoredMinimalCandidateSelectionRow) -> bool { row.provider_type.trim().eq_ignore_ascii_case("codex") } -fn codex_model_card_is_complete(card: &serde_json::Map) -> bool { - card.get("slug").and_then(Value::as_str).is_some() - && card.get("display_name").and_then(Value::as_str).is_some() - && card - .get("supported_reasoning_levels") - .and_then(Value::as_array) - .is_some() - && card.get("shell_type").and_then(Value::as_str).is_some() - && card.get("visibility").and_then(Value::as_str).is_some() - && card - .get("supported_in_api") - .and_then(Value::as_bool) - .is_some() - && card.get("priority").and_then(Value::as_i64).is_some() - && card - .get("base_instructions") - .and_then(Value::as_str) - .is_some() - && card - .get("supports_reasoning_summary_parameter") - .is_none_or(Value::is_boolean) - && card - .get("support_verbosity") - .and_then(Value::as_bool) - .is_some() - && card - .get("truncation_policy") - .and_then(Value::as_object) - .is_some() - && card - .get("supports_parallel_tool_calls") - .and_then(Value::as_bool) - .is_some() - && card - .get("experimental_supported_tools") - .and_then(Value::as_array) - .is_some() -} - -fn project_codex_model_card( - cached_models: &[Value], - source_model: &str, - global_model: &str, -) -> Option { - let mut card = cached_models - .iter() - .find(|model| { - model.get("id").and_then(Value::as_str) == Some(source_model) - || model.get("slug").and_then(Value::as_str) == Some(source_model) - })? - .as_object()? - .clone(); - if !codex_model_card_is_complete(&card) { - return None; - } - - card.remove("id"); - card.remove("api_formats"); - card.insert("slug".to_string(), Value::String(global_model.to_string())); - Some(Value::Object(card)) -} - async fn load_codex_model_cards( state: &AppState, rows: &[StoredMinimalCandidateSelectionRow], -) -> Vec { - let cache_keys = rows + targets: &[crate::model_fetch::CodexCatalogTarget], + client_version: crate::model_fetch::NormalizedCodexClientVersion, +) -> (Vec, Option) { + let catalogs = crate::model_fetch::load_codex_catalogs(state, targets, &client_version).await; + for target in catalogs.stale_targets() { + let state = state.clone(); + let target = target.clone(); + let client_version = client_version.clone(); + tokio::spawn(async move { + crate::model_fetch::refresh_codex_catalog_target(&state, &target, &client_version) + .await; + }); + } + if !catalogs.is_complete() { + warn!( + event_name = "codex_catalog_aggregate_incomplete", + client_version = %client_version.as_str(), + target_count = targets.len(), + "Codex catalog aggregation was incomplete; returning an empty remote catalog so the client can use its bundled fallback" + ); + return (Vec::new(), None); + } + let mut seen_global_models = BTreeSet::new(); + let possible_inference_catalogs = rows .iter() .filter(|row| is_codex_provider_row(row)) - .map(|row| format!("upstream_models:{}:{}", row.provider_id, row.key_id)) - .collect::>() - .into_iter() - .collect::>(); - let cached_values = await_models_route_read( - "codex_models_cache", - state.runtime_state.kv_get_many(&cache_keys), - ) - .await - .unwrap_or_default(); - let cached_models_by_key = cache_keys - .into_iter() - .zip(cached_values) - .filter_map(|(key, raw)| { - let models = serde_json::from_str::>(raw.as_deref()?).ok()?; - Some((key, models)) - }) - .collect::>(); - - let mut seen_global_models = BTreeSet::new(); + .map(|row| (row.provider_id.clone(), row.key_id.clone())) + .collect::>(); + let expected_global_models = rows + .iter() + .filter(|row| is_codex_provider_row(row)) + .map(|row| row.global_model_name.clone()) + .collect::>(); 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 { + let Some(snapshot) = catalogs.snapshot(&row.provider_id, &row.key_id) else { continue; }; let source_model = aether_scheduler_core::select_provider_model_name(row, "openai:responses"); - let Some(card) = project_codex_model_card( - cached_models, + let Some(card) = crate::ai_serving::project_codex_catalog_model_card( + &snapshot.models, source_model.as_str(), row.global_model_name.as_str(), ) else { + warn!( + event_name = "codex_catalog_authorized_model_missing", + provider_id = %row.provider_id, + key_id = %row.key_id, + client_version = %client_version.as_str(), + source_model = %source_model, + global_model = %row.global_model_name, + "authorized Codex model was not present in this upstream catalog mapping" + ); continue; }; seen_global_models.insert(row.global_model_name.clone()); cards.push(card); + if cards.len() > CODEX_MODELS_MAX_RESPONSE_MODELS { + warn!( + event_name = "codex_catalog_aggregate_model_limit", + client_version = %client_version.as_str(), + model_count = cards.len(), + limit = CODEX_MODELS_MAX_RESPONSE_MODELS, + "Codex projected catalog exceeded the aggregate model limit; returning an empty remote catalog" + ); + return (Vec::new(), None); + } } - cards + let missing_model_count = expected_global_models + .difference(&seen_global_models) + .count(); + if missing_model_count > 0 { + warn!( + event_name = "codex_catalog_authorized_models_incomplete", + client_version = %client_version.as_str(), + expected_model_count = expected_global_models.len(), + projected_model_count = cards.len(), + missing_model_count, + "Codex upstream catalogs omitted authorized mappings; returning an empty remote catalog so the client can use its bundled fallback" + ); + return (Vec::new(), None); + } + if !codex_projected_catalog_fits_response_limits(&cards) { + warn!( + event_name = "codex_catalog_aggregate_body_limit", + client_version = %client_version.as_str(), + model_count = cards.len(), + limit_bytes = CODEX_MODELS_MAX_RESPONSE_JSON_BYTES, + "Codex projected catalog exceeded the aggregate response body limit; returning an empty remote catalog" + ); + return (Vec::new(), None); + } + if cards.is_empty() { + return (cards, None); + } + let etag = if possible_inference_catalogs.len() == 1 { + possible_inference_catalogs + .iter() + .next() + .and_then(|(provider_id, key_id)| catalogs.snapshot(provider_id, key_id)) + .and_then(|snapshot| snapshot.etag.clone()) + } else { + None + }; + (cards, etag) +} + +struct ModelRowsForClientFormat { + rows: Vec, + codex_catalog_targets: Vec, } async fn list_model_rows_for_client_format( state: &AppState, api_format: &str, auth_snapshot: Option<&crate::data::auth::GatewayAuthApiKeySnapshot>, -) -> Option> { +) -> Option { let mut collected = Vec::new(); let query_api_formats = if is_codex_models_api_format(api_format) { CODEX_MODELS_QUERY_API_FORMATS @@ -254,9 +270,16 @@ async fn list_model_rows_for_client_format( } if is_codex_models_api_format(api_format) { collected.retain(is_codex_provider_row); - Some(sort_model_rows(collected)) + let codex_catalog_targets = crate::model_fetch::codex_catalog_targets(&collected); + Some(ModelRowsForClientFormat { + rows: sort_model_rows(collected), + codex_catalog_targets, + }) } else { - Some(sort_and_dedup_model_rows(collected)) + Some(ModelRowsForClientFormat { + rows: sort_and_dedup_model_rows(collected), + codex_catalog_targets: Vec::new(), + }) } } @@ -296,6 +319,9 @@ pub(super) async fn maybe_build_local_models_route_response( } let auth_context = decision.auth_context.as_ref()?; + if !auth_context.access_allowed || auth_context.local_rejection.is_some() { + return Some(build_models_auth_error_response(api_format)); + } let now_unix_secs = SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap_or_default() @@ -318,11 +344,23 @@ pub(super) async fn maybe_build_local_models_route_response( )) } }; - let auth_snapshot = auth_snapshot.as_ref(); + let Some(auth_snapshot) = auth_snapshot.as_ref() else { + warn!( + event_name = "models_route_auth_snapshot_missing", + user_id = %auth_context.user_id, + api_key_id = %auth_context.api_key_id, + "gateway models route rejected a request whose authenticated API key snapshot disappeared" + ); + return Some(build_models_auth_error_response(api_format)); + }; + if !auth_snapshot.currently_usable { + return Some(build_models_auth_error_response(api_format)); + } + let auth_snapshot = Some(auth_snapshot); match decision.route_kind.as_deref() { Some("list") => { - let rows = + let listed = match list_model_rows_for_client_format(state, api_format, auth_snapshot).await { Some(rows) => rows, None => { @@ -332,12 +370,34 @@ pub(super) async fn maybe_build_local_models_route_response( )) } }; + let rows = listed.rows; if rows.is_empty() { return Some(build_empty_models_list_response(api_format)); } if is_codex_models_api_format(api_format) { - let models = load_codex_model_cards(state, &rows).await; - return Some(build_codex_models_list_response(models)); + let raw_client_version = query_param_value( + request_context.request_query_string.as_deref(), + "client_version", + ); + let client_version = crate::model_fetch::normalize_codex_client_version( + raw_client_version.as_deref(), + ); + if client_version.used_fallback() { + warn!( + event_name = "codex_catalog_invalid_client_version", + raw_length = raw_client_version.as_ref().map_or(0, String::len), + fallback_version = %client_version.as_str(), + "invalid Codex client_version used the bounded fallback version" + ); + } + let (models, etag) = load_codex_model_cards( + state, + &rows, + &listed.codex_catalog_targets, + client_version, + ) + .await; + return Some(build_codex_models_list_response(models, etag.as_deref())); } let response = match api_format { "claude:messages" => { @@ -410,3 +470,32 @@ pub(super) async fn maybe_build_local_models_route_response( _ => Some(build_models_auth_error_response(api_format)), } } + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::{ + codex_projected_catalog_fits_response_limits, CODEX_MODELS_MAX_RESPONSE_JSON_BYTES, + CODEX_MODELS_MAX_RESPONSE_MODELS, + }; + + #[test] + fn projected_codex_catalog_enforces_aggregate_count_and_body_limits() { + assert!(codex_projected_catalog_fits_response_limits(&[json!({ + "slug": "gpt-future-dynamic", + "model_messages": {"instructions_template": "opaque"} + })])); + + let too_many = (0..=CODEX_MODELS_MAX_RESPONSE_MODELS) + .map(|index| json!({"slug": format!("gpt-future-{index}")})) + .collect::>(); + assert!(!codex_projected_catalog_fits_response_limits(&too_many)); + + let oversized = vec![json!({ + "slug": "gpt-future-oversized", + "future_capability": "x".repeat(CODEX_MODELS_MAX_RESPONSE_JSON_BYTES) + })]; + assert!(!codex_projected_catalog_fits_response_limits(&oversized)); + } +} diff --git a/apps/aether-gateway/src/handlers/public/support/models/shared.rs b/apps/aether-gateway/src/handlers/public/support/models/shared.rs index 1369c5a11..3767e8c43 100644 --- a/apps/aether-gateway/src/handlers/public/support/models/shared.rs +++ b/apps/aether-gateway/src/handlers/public/support/models/shared.rs @@ -194,7 +194,7 @@ fn row_exposes_global_model_for_models( false } -pub(super) fn filter_eligible_model_rows( +pub(crate) fn filter_eligible_model_rows( rows: Vec, auth_snapshot: Option<&crate::data::auth::GatewayAuthApiKeySnapshot>, api_format: &str, diff --git a/apps/aether-gateway/src/model_fetch/catalog.rs b/apps/aether-gateway/src/model_fetch/catalog.rs new file mode 100644 index 000000000..c4ab9347b --- /dev/null +++ b/apps/aether-gateway/src/model_fetch/catalog.rs @@ -0,0 +1,3585 @@ +use std::collections::{BTreeMap, BTreeSet}; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, LazyLock, Mutex as StdMutex, Weak}; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; + +use aether_data_contracts::repository::candidate_selection::StoredMinimalCandidateSelectionRow; +use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey; +use aether_model_fetch::{ + fetch_models_from_transports_for_client_version, ModelFetchTransportRuntime, +}; +use aether_provider_transport::GatewayProviderTransportSnapshot; +use aether_runtime_state::RuntimeState; +use async_trait::async_trait; +use futures_util::{stream, StreamExt}; +use semver::Version; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use sha2::{Digest, Sha256}; +use tokio::sync::{Mutex, Semaphore}; +use tracing::{debug, info, warn}; + +const CODEX_CATALOG_SCHEMA_VERSION: u32 = 2; +const CODEX_CATALOG_CREDENTIAL_SCOPE_DOMAIN: &str = "aether-codex-catalog-credential-v2"; +const CODEX_CLIENT_VERSION_MAX_LEN: usize = 64; +const CODEX_CATALOG_MAX_MODELS: usize = 512; +const CODEX_CATALOG_MAX_BODY_BYTES: usize = 8 * 1024 * 1024; +const CODEX_CATALOG_MAX_AGGREGATE_BODY_BYTES: usize = 8 * 1024 * 1024; +const CODEX_CATALOG_MAX_ETAG_LEN: usize = 256; +const CODEX_CATALOG_MAX_VERSIONS_PER_KEY: usize = 8; +const CODEX_CATALOG_MAX_FAILED_VERSIONS_PER_KEY: usize = 8; +const CODEX_CATALOG_MAX_TARGETS_PER_LOAD: usize = 256; +const CODEX_CATALOG_MAX_COLD_FETCHES_PER_LOAD: usize = 8; +const CODEX_CATALOG_MAX_STALE_REFRESHES_PER_LOAD: usize = 8; +const CODEX_CATALOG_MAX_CONCURRENT_FETCHES: usize = 8; +const CODEX_CATALOG_MAX_CONCURRENT_LOADS: usize = 16; +const CODEX_CATALOG_RETENTION_LOCK_TTL: Duration = Duration::from_secs(5); +const CODEX_CATALOG_RETENTION_LOCK_WAIT: Duration = Duration::from_secs(2); + +#[cfg(not(test))] +const CODEX_CATALOG_FRESH_TTL: Duration = Duration::from_secs(5 * 60); +#[cfg(test)] +const CODEX_CATALOG_FRESH_TTL: Duration = Duration::from_secs(1); +#[cfg(not(test))] +const CODEX_CATALOG_RETRY_TTL: Duration = Duration::from_secs(30); +#[cfg(test)] +const CODEX_CATALOG_RETRY_TTL: Duration = Duration::from_millis(100); +#[cfg(not(test))] +const CODEX_CATALOG_FETCH_TIMEOUT: Duration = Duration::from_secs(8); +#[cfg(test)] +const CODEX_CATALOG_FETCH_TIMEOUT: Duration = Duration::from_secs(3); +#[cfg(not(test))] +const CODEX_CATALOG_FETCH_PACE_TTL: Duration = Duration::from_secs(1); +#[cfg(test)] +const CODEX_CATALOG_FETCH_PACE_TTL: Duration = Duration::from_millis(20); +#[cfg(not(test))] +const CODEX_CATALOG_FLIGHT_WAIT_TIMEOUT: Duration = Duration::from_secs(9); +#[cfg(test)] +const CODEX_CATALOG_FLIGHT_WAIT_TIMEOUT: Duration = Duration::from_secs(4); +#[cfg(not(test))] +const CODEX_CATALOG_FLIGHT_LOCK_TTL: Duration = Duration::from_secs(20); +#[cfg(test)] +const CODEX_CATALOG_FLIGHT_LOCK_TTL: Duration = Duration::from_secs(8); +#[cfg(not(test))] +const CODEX_CATALOG_FLIGHT_LEADER_TIMEOUT: Duration = Duration::from_secs(18); +#[cfg(test)] +const CODEX_CATALOG_FLIGHT_LEADER_TIMEOUT: Duration = Duration::from_secs(7); + +static CODEX_CATALOG_FETCH_GATE: Semaphore = + Semaphore::const_new(CODEX_CATALOG_MAX_CONCURRENT_FETCHES); +static CODEX_CATALOG_LOCAL_FALLBACK_FLIGHTS: LazyLock>>>> = + LazyLock::new(|| StdMutex::new(BTreeMap::new())); +static CODEX_CATALOG_LOCAL_RETENTION: Mutex<()> = Mutex::const_new(()); + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct NormalizedCodexClientVersion { + value: String, + used_fallback: bool, +} + +impl NormalizedCodexClientVersion { + pub(crate) fn as_str(&self) -> &str { + self.value.as_str() + } + + pub(crate) fn used_fallback(&self) -> bool { + self.used_fallback + } +} + +pub(crate) fn normalize_codex_client_version(raw: Option<&str>) -> NormalizedCodexClientVersion { + let normalized = raw + .map(str::trim) + .filter(|value| !value.is_empty() && value.len() <= CODEX_CLIENT_VERSION_MAX_LEN) + .and_then(|value| Version::parse(value).ok()) + .map(|version| format!("{}.{}.{}", version.major, version.minor, version.patch)); + match normalized { + Some(value) => NormalizedCodexClientVersion { + value, + used_fallback: false, + }, + None => NormalizedCodexClientVersion { + value: crate::ai_serving::CODEX_CLIENT_VERSION.to_string(), + used_fallback: true, + }, + } +} + +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] +struct CodexCatalogIdentity { + provider_id: String, + key_id: String, +} + +impl CodexCatalogIdentity { + fn new(provider_id: &str, key_id: &str) -> Self { + Self { + provider_id: provider_id.to_string(), + key_id: key_id.to_string(), + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct CodexCatalogTarget { + identity: CodexCatalogIdentity, + endpoint_ids: Vec, + credential_scope: Option, +} + +impl CodexCatalogTarget { + fn credential_scope(&self) -> Option<&str> { + self.credential_scope.as_deref() + } + + fn with_credential_scope(&self, credential_scope: String) -> Self { + Self { + identity: self.identity.clone(), + endpoint_ids: self.endpoint_ids.clone(), + credential_scope: Some(credential_scope), + } + } +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub(crate) struct CodexCatalogSnapshot { + schema_version: u32, + pub(crate) provider_id: String, + pub(crate) key_id: String, + credential_scope: String, + pub(crate) client_version: String, + pub(crate) models: Vec, + pub(crate) etag: Option, + pub(crate) fetched_at_unix_secs: u64, + pub(crate) last_checked_at_unix_secs: u64, + pub(crate) content_sha256: String, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +struct CodexCatalogStatus { + schema_version: u32, + provider_id: String, + key_id: String, + credential_scope: String, + client_version: String, + last_checked_at_unix_secs: u64, + last_success_at_unix_secs: Option, + last_error: Option, + last_http_status: Option, + model_count: Option, + fetch_duration_ms: u64, + used_lkg: bool, +} + +#[derive(Debug, Default)] +pub(crate) struct CodexCatalogLoad { + snapshots: BTreeMap, + stale_targets: Vec, + complete: bool, +} + +impl CodexCatalogLoad { + pub(crate) fn snapshot( + &self, + provider_id: &str, + key_id: &str, + ) -> Option<&CodexCatalogSnapshot> { + self.snapshots + .get(&CodexCatalogIdentity::new(provider_id, key_id)) + } + + pub(crate) fn stale_targets(&self) -> &[CodexCatalogTarget] { + &self.stale_targets + } + + pub(crate) fn is_complete(&self) -> bool { + self.complete + } +} + +#[async_trait] +pub(crate) trait CodexCatalogRuntime: ModelFetchTransportRuntime + Send + Sync { + fn codex_catalog_runtime_state(&self) -> &RuntimeState; + + fn codex_catalog_fetch_pace_ttl(&self) -> Duration { + CODEX_CATALOG_FETCH_PACE_TTL + } + + async fn read_codex_catalog_transport_snapshot( + &self, + provider_id: &str, + endpoint_id: &str, + key_id: &str, + ) -> Result, String>; + + async fn read_codex_catalog_credential_scope_strong( + &self, + provider_id: &str, + key_id: &str, + ) -> Result, String>; +} + +fn credential_scope_digest(kind: &str, value: &str) -> Option { + let value = value.trim(); + (!value.is_empty()).then(|| { + sha256_hex(format!("{CODEX_CATALOG_CREDENTIAL_SCOPE_DOMAIN}\0{kind}\0{value}").as_bytes()) + }) +} + +fn codex_catalog_credential_scope_from_parts( + upstream_metadata: Option<&Value>, + decrypted_auth_config: Option<&str>, + decrypted_api_key: Option<&str>, +) -> Option { + let codex_metadata = upstream_metadata.and_then(|metadata| metadata.get("codex")); + if let Some(generation) = + aether_admin::provider::quota::codex_credential_generation(codex_metadata) + { + if let Some(scope) = credential_scope_digest("generation", generation) { + return Some(scope); + } + } + + let auth_config = decrypted_auth_config + .map(str::trim) + .filter(|value| !value.is_empty()); + if let Some(account_id) = crate::ai_serving::parse_codex_auth_identity(auth_config).account_id { + return credential_scope_digest("account_id", &account_id); + } + let parsed_auth_config = auth_config.and_then(|raw| serde_json::from_str::(raw).ok()); + if let Some(refresh_token) = parsed_auth_config + .as_ref() + .and_then(|value| value.get("refresh_token")) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + { + return credential_scope_digest("refresh_token", refresh_token); + } + if let Some(auth_config) = auth_config { + return credential_scope_digest("auth_config", auth_config); + } + decrypted_api_key + .map(str::trim) + .filter(|value| !value.is_empty()) + .and_then(|value| credential_scope_digest("api_key", value)) +} + +pub(crate) fn codex_catalog_credential_scope_from_stored_key( + key: &StoredProviderCatalogKey, + decrypted_auth_config: Option<&str>, + decrypted_api_key: Option<&str>, +) -> Option { + codex_catalog_credential_scope_from_parts( + key.upstream_metadata.as_ref(), + decrypted_auth_config, + decrypted_api_key, + ) +} + +fn codex_catalog_credential_scope_from_transport( + transport: &GatewayProviderTransportSnapshot, +) -> Option { + codex_catalog_credential_scope_from_parts( + transport.key.upstream_metadata.as_ref(), + transport.key.decrypted_auth_config.as_deref(), + Some(&transport.key.decrypted_api_key), + ) +} + +pub(crate) fn codex_catalog_targets( + rows: &[StoredMinimalCandidateSelectionRow], +) -> Vec { + let mut endpoints_by_identity = BTreeMap::>::new(); + for row in rows + .iter() + .filter(|row| row.provider_type.trim().eq_ignore_ascii_case("codex")) + { + endpoints_by_identity + .entry(CodexCatalogIdentity::new(&row.provider_id, &row.key_id)) + .or_default() + .insert(row.endpoint_id.clone()); + } + endpoints_by_identity + .into_iter() + .map(|(identity, endpoint_ids)| CodexCatalogTarget { + identity, + endpoint_ids: endpoint_ids.into_iter().collect(), + credential_scope: None, + }) + .collect() +} + +async fn bind_codex_catalog_target( + runtime: &R, + target: &CodexCatalogTarget, +) -> Option +where + R: CodexCatalogRuntime + ?Sized, +{ + match runtime + .read_codex_catalog_credential_scope_strong( + &target.identity.provider_id, + &target.identity.key_id, + ) + .await + { + Ok(Some(scope)) if !scope.is_empty() => Some(target.with_credential_scope(scope)), + Ok(_) => { + warn!( + event_name = "codex_catalog_credential_scope_missing", + provider_id = %target.identity.provider_id, + key_id = %target.identity.key_id, + "Codex catalog key was missing, inactive, or lacked a safe credential scope" + ); + None + } + Err(error) => { + warn!( + event_name = "codex_catalog_credential_scope_error", + provider_id = %target.identity.provider_id, + key_id = %target.identity.key_id, + error = %safe_catalog_error(&error), + "Codex catalog credential scope strong read failed" + ); + None + } + } +} + +pub(crate) async fn load_codex_catalogs( + runtime: &R, + targets: &[CodexCatalogTarget], + client_version: &NormalizedCodexClientVersion, +) -> CodexCatalogLoad +where + R: CodexCatalogRuntime + ?Sized, +{ + let version = client_version.as_str().to_string(); + if targets.len() > CODEX_CATALOG_MAX_TARGETS_PER_LOAD { + warn!( + event_name = "codex_catalog_target_limit", + client_version = %version, + target_count = targets.len(), + limit = CODEX_CATALOG_MAX_TARGETS_PER_LOAD, + "Codex catalog request exceeded the target safety limit" + ); + return CodexCatalogLoad::default(); + } + let remember_version = !client_version.used_fallback(); + let cold_fetch_budget = Arc::new(AtomicUsize::new(CODEX_CATALOG_MAX_COLD_FETCHES_PER_LOAD)); + let results = stream::iter(targets.iter().cloned()) + .map(|target| { + let version = version.clone(); + let cold_fetch_budget = Arc::clone(&cold_fetch_budget); + async move { + let Some(target) = bind_codex_catalog_target(runtime, &target).await else { + return (target, None); + }; + let loaded = load_codex_catalog( + runtime, + &target, + &version, + remember_version, + cold_fetch_budget.as_ref(), + ) + .await; + (target, loaded) + } + }) + .buffer_unordered(CODEX_CATALOG_MAX_CONCURRENT_LOADS); + futures_util::pin_mut!(results); + + let mut load = CodexCatalogLoad { + complete: true, + ..CodexCatalogLoad::default() + }; + let mut aggregate_body_bytes = 0usize; + while let Some((target, loaded)) = results.next().await { + let Some((snapshot, is_stale, should_refresh)) = loaded else { + load.complete = false; + continue; + }; + let snapshot_body_bytes = match serde_json::to_vec(&snapshot.models) { + Ok(serialized) => serialized.len(), + Err(error) => { + warn!( + event_name = "codex_catalog_aggregate_serialize_error", + provider_id = %target.identity.provider_id, + key_id = %target.identity.key_id, + client_version = %version, + error = %error, + "Codex source catalog could not be measured for aggregate safety" + ); + load.complete = false; + load.snapshots.clear(); + load.stale_targets.clear(); + break; + } + }; + let Some(next_aggregate_body_bytes) = aggregate_body_bytes.checked_add(snapshot_body_bytes) + else { + warn!( + event_name = "codex_catalog_source_aggregate_body_limit", + provider_id = %target.identity.provider_id, + key_id = %target.identity.key_id, + client_version = %version, + limit_bytes = CODEX_CATALOG_MAX_AGGREGATE_BODY_BYTES, + "Codex source catalog aggregate size overflowed; aggregation was aborted" + ); + load.complete = false; + load.snapshots.clear(); + load.stale_targets.clear(); + break; + }; + if next_aggregate_body_bytes > CODEX_CATALOG_MAX_AGGREGATE_BODY_BYTES { + warn!( + event_name = "codex_catalog_source_aggregate_body_limit", + provider_id = %target.identity.provider_id, + key_id = %target.identity.key_id, + client_version = %version, + aggregate_body_bytes = next_aggregate_body_bytes, + limit_bytes = CODEX_CATALOG_MAX_AGGREGATE_BODY_BYTES, + "Codex source catalogs exceeded the aggregate body limit; aggregation was aborted" + ); + load.complete = false; + load.snapshots.clear(); + load.stale_targets.clear(); + break; + } + aggregate_body_bytes = next_aggregate_body_bytes; + load.snapshots.insert(target.identity.clone(), snapshot); + if is_stale + && should_refresh + && load.stale_targets.len() < CODEX_CATALOG_MAX_STALE_REFRESHES_PER_LOAD + { + load.stale_targets.push(target); + } + } + load +} + +pub(crate) async fn refresh_codex_catalog_target( + runtime: &R, + target: &CodexCatalogTarget, + client_version: &NormalizedCodexClientVersion, +) where + R: CodexCatalogRuntime + ?Sized, +{ + let version = client_version.as_str(); + let Some(target) = bind_codex_catalog_target(runtime, target).await else { + return; + }; + if runtime + .codex_catalog_runtime_state() + .kv_get(&catalog_retry_key(&target, version)) + .await + .ok() + .flatten() + .is_some() + { + return; + } + let _ = fetch_catalog_singleflight(runtime, &target, version, false).await; +} + +async fn load_codex_catalog( + runtime: &R, + target: &CodexCatalogTarget, + client_version: &str, + remember_version: bool, + cold_fetch_budget: &AtomicUsize, +) -> Option<(CodexCatalogSnapshot, bool, bool)> +where + R: CodexCatalogRuntime + ?Sized, +{ + if remember_version { + remember_seen_version( + runtime.codex_catalog_runtime_state(), + target, + client_version, + ) + .await; + } + if let Some(snapshot) = read_lkg_snapshot(runtime, target, client_version).await { + touch_catalog_version_best_effort( + runtime.codex_catalog_runtime_state(), + target, + client_version, + ); + let fresh_hash = runtime + .codex_catalog_runtime_state() + .kv_get(&catalog_fresh_key(target, client_version)) + .await + .ok() + .flatten(); + let is_fresh = fresh_hash.as_deref() == Some(snapshot.content_sha256.as_str()); + if is_fresh { + debug!( + event_name = "codex_catalog_cache", + provider_id = %target.identity.provider_id, + key_id = %target.identity.key_id, + client_version, + cache_state = "fresh", + model_count = snapshot.models.len(), + "Codex catalog fresh cache hit" + ); + return Some((snapshot, false, false)); + } + let retry_suppressed = runtime + .codex_catalog_runtime_state() + .kv_get(&catalog_retry_key(target, client_version)) + .await + .ok() + .flatten() + .is_some(); + info!( + event_name = "codex_catalog_cache", + provider_id = %target.identity.provider_id, + key_id = %target.identity.key_id, + client_version, + cache_state = "stale", + model_count = snapshot.models.len(), + refresh_suppressed = retry_suppressed, + "Codex catalog stale LKG hit" + ); + return Some((snapshot, true, !retry_suppressed)); + } + + info!( + event_name = "codex_catalog_cache", + provider_id = %target.identity.provider_id, + key_id = %target.identity.key_id, + client_version, + cache_state = "cold_miss", + "Codex catalog cold cache miss" + ); + if runtime + .codex_catalog_runtime_state() + .kv_get(&catalog_retry_key(target, client_version)) + .await + .ok() + .flatten() + .is_some() + { + debug!( + event_name = "codex_catalog_cold_retry_suppressed", + provider_id = %target.identity.provider_id, + key_id = %target.identity.key_id, + client_version, + "Codex catalog cold retry was suppressed by the failure cooldown" + ); + return None; + } + if cold_fetch_budget + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |remaining| { + remaining.checked_sub(1) + }) + .is_err() + { + debug!( + event_name = "codex_catalog_cold_fetch_budget_exhausted", + provider_id = %target.identity.provider_id, + key_id = %target.identity.key_id, + client_version, + limit = CODEX_CATALOG_MAX_COLD_FETCHES_PER_LOAD, + "Codex catalog cold fetch was deferred to a later request" + ); + return None; + } + if let Some(snapshot) = fetch_catalog_singleflight(runtime, target, client_version, true).await + { + return Some((snapshot, false, false)); + } + None +} + +async fn fetch_catalog_singleflight( + runtime: &R, + target: &CodexCatalogTarget, + client_version: &str, + wait_for_existing: bool, +) -> Option +where + R: CodexCatalogRuntime + ?Sized, +{ + let state = runtime.codex_catalog_runtime_state(); + let lock_key = catalog_flight_key(target, client_version); + let owner = format!("codex-catalog:{}", std::process::id()); + match state + .lock_try_acquire(&lock_key, &owner, CODEX_CATALOG_FLIGHT_LOCK_TTL) + .await + { + Ok(Some(lease)) => { + let snapshot = match tokio::time::timeout( + CODEX_CATALOG_FLIGHT_LEADER_TIMEOUT, + load_or_fetch_catalog_leader(runtime, target, client_version, wait_for_existing), + ) + .await + { + Ok(snapshot) => snapshot, + Err(_) => { + warn!( + event_name = "codex_catalog_singleflight_leader_timeout", + provider_id = %target.identity.provider_id, + key_id = %target.identity.key_id, + client_version, + "Codex catalog singleflight leader exceeded its lease-safe deadline" + ); + None + } + }; + match state.lock_release(&lease).await { + Ok(true) => {} + Ok(false) => warn!( + event_name = "codex_catalog_singleflight_lease_lost", + provider_id = %target.identity.provider_id, + key_id = %target.identity.key_id, + client_version, + "Codex catalog singleflight lease was no longer owned at release" + ), + Err(error) => warn!( + event_name = "codex_catalog_singleflight_release_error", + provider_id = %target.identity.provider_id, + key_id = %target.identity.key_id, + client_version, + error = %error, + "Codex catalog singleflight lease release failed" + ), + } + snapshot + } + Ok(None) if wait_for_existing => { + wait_for_catalog_flight(runtime, target, client_version).await + } + Ok(None) => None, + Err(error) => { + warn!( + event_name = "codex_catalog_singleflight_error", + provider_id = %target.identity.provider_id, + key_id = %target.identity.key_id, + client_version, + error = %error, + "Codex catalog singleflight backend unavailable" + ); + let local_flight = local_catalog_fallback_flight(&lock_key); + let local_guard = + match tokio::time::timeout(CODEX_CATALOG_FLIGHT_WAIT_TIMEOUT, local_flight.lock()) + .await + { + Ok(guard) => guard, + Err(_) => { + prune_local_catalog_fallback_flight(&lock_key, &local_flight); + return read_lkg_snapshot(runtime, target, client_version).await; + } + }; + let snapshot = + load_or_fetch_catalog_leader(runtime, target, client_version, wait_for_existing) + .await; + drop(local_guard); + prune_local_catalog_fallback_flight(&lock_key, &local_flight); + snapshot + } + } +} + +fn local_catalog_fallback_flight(lock_key: &str) -> Arc> { + let mut flights = CODEX_CATALOG_LOCAL_FALLBACK_FLIGHTS + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + flights.retain(|_, flight| flight.strong_count() > 0); + if let Some(flight) = flights.get(lock_key).and_then(Weak::upgrade) { + return flight; + } + let flight = Arc::new(Mutex::new(())); + flights.insert(lock_key.to_string(), Arc::downgrade(&flight)); + flight +} + +fn prune_local_catalog_fallback_flight(lock_key: &str, flight: &Arc>) { + let mut flights = CODEX_CATALOG_LOCAL_FALLBACK_FLIGHTS + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let weak = Arc::downgrade(flight); + let remove = flights + .get(lock_key) + .is_some_and(|stored| stored.ptr_eq(&weak) && Arc::strong_count(flight) == 1); + if remove { + flights.remove(lock_key); + } +} + +async fn load_or_fetch_catalog_leader( + runtime: &R, + target: &CodexCatalogTarget, + client_version: &str, + wait_for_existing: bool, +) -> Option +where + R: CodexCatalogRuntime + ?Sized, +{ + let state = runtime.codex_catalog_runtime_state(); + let existing = read_lkg_snapshot(runtime, target, client_version).await; + let fresh_matches = if let Some(snapshot) = existing.as_ref() { + state + .kv_get(&catalog_fresh_key(target, client_version)) + .await + .ok() + .flatten() + .as_deref() + == Some(snapshot.content_sha256.as_str()) + } else { + false + }; + let retry_suppressed = state + .kv_get(&catalog_retry_key(target, client_version)) + .await + .ok() + .flatten() + .is_some(); + if fresh_matches || retry_suppressed || (wait_for_existing && existing.is_some()) { + return existing; + } + if !acquire_catalog_fetch_pace(runtime, target, client_version).await { + return existing; + } + fetch_catalog_with_deadline(runtime, target, client_version).await +} + +async fn acquire_catalog_fetch_pace( + runtime: &R, + target: &CodexCatalogTarget, + client_version: &str, +) -> bool +where + R: CodexCatalogRuntime + ?Sized, +{ + let state = runtime.codex_catalog_runtime_state(); + let lock_key = catalog_fetch_pace_key(target); + let owner = format!("codex-catalog-fetch-pace:{}", std::process::id()); + let pace_ttl = runtime.codex_catalog_fetch_pace_ttl(); + match state.lock_try_acquire(&lock_key, &owner, pace_ttl).await { + Ok(Some(_lease)) => true, + Ok(None) => { + debug!( + event_name = "codex_catalog_fetch_paced", + provider_id = %target.identity.provider_id, + key_id = %target.identity.key_id, + client_version, + pace_ms = duration_ms(pace_ttl), + "Codex catalog upstream fetch was suppressed by the per-key pace limit" + ); + false + } + Err(error) => { + warn!( + event_name = "codex_catalog_fetch_pace_error", + provider_id = %target.identity.provider_id, + key_id = %target.identity.key_id, + client_version, + error = %error, + "Codex catalog upstream fetch was suppressed because the pace limiter was unavailable" + ); + false + } + } +} + +async fn fetch_catalog_with_deadline( + runtime: &R, + target: &CodexCatalogTarget, + client_version: &str, +) -> Option +where + R: CodexCatalogRuntime + ?Sized, +{ + let started = Instant::now(); + match tokio::time::timeout( + CODEX_CATALOG_FETCH_TIMEOUT, + fetch_and_store_catalog(runtime, target, client_version), + ) + .await + { + Ok(snapshot) => snapshot, + Err(_) => { + record_catalog_failure( + runtime, + target, + client_version, + current_unix_secs(), + started.elapsed(), + None, + "upstream catalog fetch timed out", + ) + .await; + read_lkg_snapshot(runtime, target, client_version).await + } + } +} + +async fn wait_for_catalog_flight( + runtime: &R, + target: &CodexCatalogTarget, + client_version: &str, +) -> Option +where + R: CodexCatalogRuntime + ?Sized, +{ + let deadline = tokio::time::Instant::now() + CODEX_CATALOG_FLIGHT_WAIT_TIMEOUT; + loop { + if let Some(snapshot) = read_lkg_snapshot(runtime, target, client_version).await { + return Some(snapshot); + } + if runtime + .codex_catalog_runtime_state() + .kv_get(&catalog_retry_key(target, client_version)) + .await + .ok() + .flatten() + .is_some() + { + return None; + } + if tokio::time::Instant::now() >= deadline { + warn!( + event_name = "codex_catalog_singleflight_wait_timeout", + provider_id = %target.identity.provider_id, + key_id = %target.identity.key_id, + client_version, + "Codex catalog cold miss timed out waiting for in-flight fetch" + ); + return None; + } + tokio::time::sleep(Duration::from_millis(25)).await; + } +} + +async fn fetch_and_store_catalog( + runtime: &R, + target: &CodexCatalogTarget, + client_version: &str, +) -> Option +where + R: CodexCatalogRuntime + ?Sized, +{ + let started = Instant::now(); + let now = current_unix_secs(); + let expected_scope = target.credential_scope()?; + let mut transports = Vec::new(); + let mut transport_load_failed = false; + for endpoint_id in &target.endpoint_ids { + match runtime + .read_codex_catalog_transport_snapshot( + &target.identity.provider_id, + endpoint_id, + &target.identity.key_id, + ) + .await + { + Ok(Some(transport)) => { + let transport_scope = codex_catalog_credential_scope_from_transport(&transport); + if transport_scope.as_deref() != Some(expected_scope) { + transport_load_failed = true; + warn!( + event_name = "codex_catalog_transport_credential_scope_mismatch", + provider_id = %target.identity.provider_id, + endpoint_id, + key_id = %target.identity.key_id, + client_version, + "Codex catalog cached transport did not match the strong credential scope" + ); + } else { + transports.push(transport); + } + } + Ok(None) => { + transport_load_failed = true; + warn!( + event_name = "codex_catalog_transport_missing", + provider_id = %target.identity.provider_id, + endpoint_id, + key_id = %target.identity.key_id, + client_version, + "Codex catalog transport snapshot unavailable" + ); + } + Err(error) => { + transport_load_failed = true; + warn!( + event_name = "codex_catalog_transport_error", + provider_id = %target.identity.provider_id, + endpoint_id, + key_id = %target.identity.key_id, + client_version, + error = %safe_catalog_error(&error), + "Codex catalog transport snapshot failed" + ); + } + } + } + if transport_load_failed || transports.is_empty() { + record_catalog_failure( + runtime, + target, + client_version, + now, + started.elapsed(), + None, + "one or more catalog transports were unavailable", + ) + .await; + return None; + } + + let permit = match CODEX_CATALOG_FETCH_GATE.acquire().await { + Ok(permit) => permit, + Err(_) => { + record_catalog_failure( + runtime, + target, + client_version, + now, + started.elapsed(), + None, + "catalog fetch concurrency limit unavailable", + ) + .await; + return None; + } + }; + + let fetched = + fetch_models_from_transports_for_client_version(runtime, &transports, Some(client_version)) + .await; + drop(permit); + + let outcome = match fetched { + Err(error) => { + record_catalog_failure( + runtime, + target, + client_version, + now, + started.elapsed(), + None, + &error, + ) + .await; + return None; + } + Ok(outcome) => outcome, + }; + + if !outcome.has_success { + let error = if outcome.errors.is_empty() { + "upstream catalog fetch failed".to_string() + } else { + outcome.errors.join("; ") + }; + record_catalog_failure( + runtime, + target, + client_version, + now, + started.elapsed(), + outcome.upstream_status, + &error, + ) + .await; + return None; + } + if !outcome.errors.is_empty() { + let error = outcome.errors.join("; "); + record_catalog_failure( + runtime, + target, + client_version, + now, + started.elapsed(), + outcome.upstream_status, + &error, + ) + .await; + return None; + } + + let models = outcome.cached_models; + let serialized_models = match validate_catalog_models(&models) { + Ok(serialized) => serialized, + Err(error) => { + record_catalog_failure( + runtime, + target, + client_version, + now, + started.elapsed(), + outcome.upstream_status, + &error, + ) + .await; + return None; + } + }; + let snapshot = CodexCatalogSnapshot { + schema_version: CODEX_CATALOG_SCHEMA_VERSION, + provider_id: target.identity.provider_id.clone(), + key_id: target.identity.key_id.clone(), + credential_scope: expected_scope.to_string(), + client_version: client_version.to_string(), + models, + etag: outcome.etag.as_deref().and_then(normalize_etag), + fetched_at_unix_secs: now, + last_checked_at_unix_secs: now, + content_sha256: sha256_hex(&serialized_models), + }; + if !codex_catalog_credential_scope_is_current(runtime, target).await { + warn!( + event_name = "codex_catalog_credential_scope_changed", + provider_id = %target.identity.provider_id, + key_id = %target.identity.key_id, + client_version, + "Codex catalog fetch result was discarded because the credential scope changed" + ); + return None; + } + persist_catalog_success( + runtime, + target, + client_version, + &snapshot, + outcome.upstream_status, + started.elapsed(), + ) + .await; + if !codex_catalog_credential_scope_is_current(runtime, target).await { + discard_catalog_success_after_scope_change( + runtime.codex_catalog_runtime_state(), + target, + client_version, + ) + .await; + warn!( + event_name = "codex_catalog_credential_scope_changed_after_persist", + provider_id = %target.identity.provider_id, + key_id = %target.identity.key_id, + client_version, + "Codex catalog fetch result was removed because the credential scope changed during persistence" + ); + return None; + } + Some(snapshot) +} + +async fn discard_catalog_success_after_scope_change( + state: &RuntimeState, + target: &CodexCatalogTarget, + client_version: &str, +) { + let keys = vec![ + catalog_lkg_key(target, client_version), + catalog_fresh_key(target, client_version), + catalog_status_key(target, client_version), + catalog_retry_key(target, client_version), + ]; + if let Err(error) = state.kv_delete_many(&keys).await { + warn!( + event_name = "codex_catalog_scope_changed_cleanup_error", + provider_id = %target.identity.provider_id, + key_id = %target.identity.key_id, + client_version, + error = %error, + "Codex catalog stale credential-scope KV cleanup failed" + ); + } + let member = catalog_version_member(target, client_version); + let _ = state + .score_remove(&catalog_versions_key(&target.identity), &member) + .await; + let _ = state + .score_remove(&catalog_failed_versions_key(&target.identity), &member) + .await; +} + +async fn codex_catalog_credential_scope_is_current( + runtime: &R, + target: &CodexCatalogTarget, +) -> bool +where + R: CodexCatalogRuntime + ?Sized, +{ + let Some(expected_scope) = target.credential_scope() else { + return false; + }; + matches!( + runtime + .read_codex_catalog_credential_scope_strong( + &target.identity.provider_id, + &target.identity.key_id, + ) + .await, + Ok(Some(current_scope)) if current_scope == expected_scope + ) +} + +fn validate_catalog_models(models: &[Value]) -> Result, String> { + if models.is_empty() { + return Err("upstream catalog returned no models".to_string()); + } + if models.len() > CODEX_CATALOG_MAX_MODELS { + return Err(format!( + "upstream catalog model count exceeds {CODEX_CATALOG_MAX_MODELS}" + )); + } + if models + .iter() + .any(|model| !catalog_model_has_valid_identity(model)) + { + return Err("upstream catalog contains a model without a valid id or slug".to_string()); + } + let serialized = serde_json::to_vec(models) + .map_err(|_| "upstream catalog models could not be serialized".to_string())?; + if serialized.len() > CODEX_CATALOG_MAX_BODY_BYTES { + return Err(format!( + "upstream catalog body exceeds {CODEX_CATALOG_MAX_BODY_BYTES} bytes" + )); + } + Ok(serialized) +} + +fn catalog_model_has_valid_identity(model: &Value) -> bool { + let Some(object) = model.as_object() else { + return false; + }; + ["slug", "id"].iter().any(|field| { + object + .get(*field) + .and_then(Value::as_str) + .filter(|value| *value == value.trim()) + .is_some_and(valid_model_identity) + }) +} + +fn valid_model_identity(value: &str) -> bool { + !value.is_empty() + && value.len() <= 256 + && value == value.trim() + && value + .chars() + .all(|character| !character.is_whitespace() && !character.is_control()) +} + +async fn persist_catalog_success( + runtime: &R, + target: &CodexCatalogTarget, + client_version: &str, + snapshot: &CodexCatalogSnapshot, + status_code: Option, + elapsed: Duration, +) where + R: CodexCatalogRuntime + ?Sized, +{ + let state = runtime.codex_catalog_runtime_state(); + let serialized = match serde_json::to_string(snapshot) { + Ok(serialized) => serialized, + Err(error) => { + warn!( + event_name = "codex_catalog_snapshot_serialize_error", + provider_id = %target.identity.provider_id, + key_id = %target.identity.key_id, + client_version, + error = %error, + "Codex catalog snapshot serialization failed" + ); + return; + } + }; + if let Err(error) = state + .kv_set(&catalog_lkg_key(target, client_version), serialized, None) + .await + { + warn!( + event_name = "codex_catalog_lkg_write_error", + provider_id = %target.identity.provider_id, + key_id = %target.identity.key_id, + client_version, + error = %error, + "Codex catalog LKG write failed" + ); + return; + } + if let Err(error) = state + .kv_set( + &catalog_fresh_key(target, client_version), + snapshot.content_sha256.clone(), + Some(CODEX_CATALOG_FRESH_TTL), + ) + .await + { + warn!( + event_name = "codex_catalog_fresh_marker_write_error", + provider_id = %target.identity.provider_id, + key_id = %target.identity.key_id, + client_version, + error = %error, + "Codex catalog fresh marker write failed" + ); + } + let _ = state + .kv_delete(&catalog_retry_key(target, client_version)) + .await; + let status = CodexCatalogStatus { + schema_version: CODEX_CATALOG_SCHEMA_VERSION, + provider_id: target.identity.provider_id.clone(), + key_id: target.identity.key_id.clone(), + credential_scope: target.credential_scope().unwrap_or_default().to_string(), + client_version: client_version.to_string(), + last_checked_at_unix_secs: snapshot.last_checked_at_unix_secs, + last_success_at_unix_secs: Some(snapshot.fetched_at_unix_secs), + last_error: None, + last_http_status: status_code, + model_count: Some(snapshot.models.len()), + fetch_duration_ms: duration_ms(elapsed), + used_lkg: false, + }; + write_catalog_status(state, target, client_version, &status).await; + retain_catalog_version(state, target, client_version, true).await; + info!( + event_name = "codex_catalog_fetch_success", + provider_id = %target.identity.provider_id, + key_id = %target.identity.key_id, + client_version, + status_code, + model_count = snapshot.models.len(), + duration_ms = duration_ms(elapsed), + "Codex catalog upstream fetch succeeded" + ); +} + +#[allow(clippy::too_many_arguments)] +async fn record_catalog_failure( + runtime: &R, + target: &CodexCatalogTarget, + client_version: &str, + checked_at: u64, + elapsed: Duration, + status_code: Option, + error: &str, +) where + R: CodexCatalogRuntime + ?Sized, +{ + let state = runtime.codex_catalog_runtime_state(); + let safe_error = safe_catalog_error(error); + let existing_snapshot = read_lkg_snapshot(runtime, target, client_version).await; + let status = CodexCatalogStatus { + schema_version: CODEX_CATALOG_SCHEMA_VERSION, + provider_id: target.identity.provider_id.clone(), + key_id: target.identity.key_id.clone(), + credential_scope: target.credential_scope().unwrap_or_default().to_string(), + client_version: client_version.to_string(), + last_checked_at_unix_secs: checked_at, + last_success_at_unix_secs: existing_snapshot + .as_ref() + .map(|snapshot| snapshot.fetched_at_unix_secs), + last_error: Some(safe_error.clone()), + last_http_status: status_code, + model_count: existing_snapshot + .as_ref() + .map(|snapshot| snapshot.models.len()), + fetch_duration_ms: duration_ms(elapsed), + used_lkg: existing_snapshot.is_some(), + }; + write_catalog_status(state, target, client_version, &status).await; + if let Err(marker_error) = state + .kv_set( + &catalog_retry_key(target, client_version), + checked_at.to_string(), + Some(CODEX_CATALOG_RETRY_TTL), + ) + .await + { + debug!( + event_name = "codex_catalog_retry_marker_write_error", + provider_id = %target.identity.provider_id, + key_id = %target.identity.key_id, + client_version, + error = %marker_error, + "Codex catalog retry marker write failed" + ); + } + retain_catalog_version(state, target, client_version, existing_snapshot.is_some()).await; + warn!( + event_name = "codex_catalog_fetch_failure", + provider_id = %target.identity.provider_id, + key_id = %target.identity.key_id, + client_version, + status_code, + duration_ms = duration_ms(elapsed), + used_lkg = existing_snapshot.is_some(), + error = %safe_error, + "Codex catalog upstream fetch failed" + ); +} + +async fn read_lkg_snapshot( + runtime: &R, + target: &CodexCatalogTarget, + client_version: &str, +) -> Option +where + R: CodexCatalogRuntime + ?Sized, +{ + let raw = match runtime + .codex_catalog_runtime_state() + .kv_get(&catalog_lkg_key(target, client_version)) + .await + { + Ok(raw) => raw?, + Err(error) => { + warn!( + event_name = "codex_catalog_lkg_read_error", + provider_id = %target.identity.provider_id, + key_id = %target.identity.key_id, + client_version, + error = %error, + "Codex catalog LKG read failed" + ); + return None; + } + }; + let snapshot = match serde_json::from_str::(&raw) { + Ok(snapshot) => snapshot, + Err(error) => { + warn!( + event_name = "codex_catalog_lkg_decode_error", + provider_id = %target.identity.provider_id, + key_id = %target.identity.key_id, + client_version, + error = %error, + "Codex catalog LKG snapshot is invalid" + ); + return None; + } + }; + let validated_models = validate_catalog_models(&snapshot.models); + let content_hash_matches = validated_models + .as_ref() + .is_ok_and(|serialized| sha256_hex(serialized) == snapshot.content_sha256); + if snapshot.schema_version != CODEX_CATALOG_SCHEMA_VERSION + || snapshot.provider_id != target.identity.provider_id + || snapshot.key_id != target.identity.key_id + || snapshot.credential_scope != target.credential_scope().unwrap_or_default() + || snapshot.client_version != client_version + || !content_hash_matches + { + warn!( + event_name = "codex_catalog_lkg_identity_mismatch", + provider_id = %target.identity.provider_id, + key_id = %target.identity.key_id, + client_version, + "Codex catalog LKG snapshot failed validation" + ); + return None; + } + if !codex_catalog_credential_scope_is_current(runtime, target).await { + warn!( + event_name = "codex_catalog_lkg_credential_scope_changed", + provider_id = %target.identity.provider_id, + key_id = %target.identity.key_id, + client_version, + "Codex catalog LKG was rejected because the credential scope changed during the read" + ); + return None; + } + Some(snapshot) +} + +async fn write_catalog_status( + state: &RuntimeState, + target: &CodexCatalogTarget, + client_version: &str, + status: &CodexCatalogStatus, +) { + let Ok(serialized) = serde_json::to_string(status) else { + return; + }; + if let Err(error) = state + .kv_set( + &catalog_status_key(target, client_version), + serialized, + None, + ) + .await + { + debug!( + event_name = "codex_catalog_status_write_error", + provider_id = %target.identity.provider_id, + key_id = %target.identity.key_id, + client_version, + error = %error, + "Codex catalog status write failed" + ); + } +} + +async fn remember_seen_version( + state: &RuntimeState, + target: &CodexCatalogTarget, + client_version: &str, +) { + if let Err(error) = state + .kv_set( + &catalog_recent_version_key(&target.identity), + catalog_version_member(target, client_version), + None, + ) + .await + { + debug!( + event_name = "codex_catalog_recent_version_write_error", + provider_id = %target.identity.provider_id, + key_id = %target.identity.key_id, + client_version, + error = %error, + "Codex catalog recent client version write failed" + ); + } +} + +fn touch_catalog_version_best_effort( + state: &RuntimeState, + target: &CodexCatalogTarget, + client_version: &str, +) { + let state = state.clone(); + let target = target.clone(); + let client_version = client_version.to_string(); + drop(tokio::spawn(async move { + if let Err(error) = state + .score_set( + &catalog_versions_key(&target.identity), + &catalog_version_member(&target, &client_version), + current_unix_ms() as f64, + ) + .await + { + debug!( + event_name = "codex_catalog_version_touch_error", + provider_id = %target.identity.provider_id, + key_id = %target.identity.key_id, + client_version, + error = %error, + "Codex catalog version best-effort touch failed" + ); + } + })); +} + +pub(crate) async fn read_recent_codex_catalog_client_version( + state: &RuntimeState, + provider_id: &str, + key_id: &str, + credential_scope: &str, +) -> Option { + let identity = CodexCatalogIdentity::new(provider_id, key_id); + let raw = state + .kv_get(&catalog_recent_version_key(&identity)) + .await + .ok() + .flatten()?; + let (stored_scope, stored_version) = parse_catalog_version_member(&raw)?; + if stored_scope != credential_scope { + return None; + } + let normalized = normalize_codex_client_version(Some(stored_version)); + (!normalized.used_fallback()).then_some(normalized.value) +} + +async fn retain_catalog_version( + state: &RuntimeState, + target: &CodexCatalogTarget, + client_version: &str, + has_lkg_hint: bool, +) { + let identity = &target.identity; + let lock_key = catalog_retention_lock_key(identity); + let owner = format!("codex-catalog-retention:{}", std::process::id()); + let started = Instant::now(); + let lease = loop { + match state + .lock_try_acquire(&lock_key, &owner, CODEX_CATALOG_RETENTION_LOCK_TTL) + .await + { + Ok(Some(lease)) => break Some(lease), + Ok(None) if started.elapsed() < CODEX_CATALOG_RETENTION_LOCK_WAIT => { + tokio::time::sleep(Duration::from_millis(10)).await; + } + Ok(None) => break None, + Err(error) => { + debug!( + event_name = "codex_catalog_retention_lock_error", + provider_id = %identity.provider_id, + key_id = %identity.key_id, + client_version, + error = %error, + "Codex catalog retention lock backend was unavailable" + ); + let _local_guard = CODEX_CATALOG_LOCAL_RETENTION.lock().await; + retain_catalog_version_locked(state, target, client_version, has_lkg_hint).await; + return; + } + } + }; + + let Some(lease) = lease else { + let index_key = if has_lkg_hint { + catalog_versions_key(identity) + } else { + catalog_failed_versions_key(identity) + }; + let _ = state + .score_set( + &index_key, + &catalog_version_member(target, client_version), + current_unix_ms() as f64, + ) + .await; + warn!( + event_name = "codex_catalog_retention_lock_timeout", + provider_id = %identity.provider_id, + key_id = %identity.key_id, + client_version, + "Codex catalog retention update was indexed without destructive trimming" + ); + return; + }; + + retain_catalog_version_locked(state, target, client_version, has_lkg_hint).await; + match state.lock_release(&lease).await { + Ok(true) => {} + Ok(false) => debug!( + event_name = "codex_catalog_retention_lease_lost", + provider_id = %identity.provider_id, + key_id = %identity.key_id, + client_version, + "Codex catalog retention lease was no longer owned at release" + ), + Err(error) => debug!( + event_name = "codex_catalog_retention_release_error", + provider_id = %identity.provider_id, + key_id = %identity.key_id, + client_version, + error = %error, + "Codex catalog retention lease release failed" + ), + } +} + +async fn retain_catalog_version_locked( + state: &RuntimeState, + target: &CodexCatalogTarget, + client_version: &str, + has_lkg_hint: bool, +) { + let identity = &target.identity; + let member = catalog_version_member(target, client_version); + let has_lkg = has_lkg_hint + || state + .kv_get(&catalog_lkg_key(target, client_version)) + .await + .ok() + .flatten() + .is_some(); + let index_key = if has_lkg { + catalog_versions_key(identity) + } else { + catalog_failed_versions_key(identity) + }; + if let Err(error) = state + .score_set(&index_key, &member, current_unix_ms() as f64) + .await + { + debug!( + event_name = "codex_catalog_version_touch_error", + provider_id = %identity.provider_id, + key_id = %identity.key_id, + client_version, + has_lkg, + error = %error, + "Codex catalog version index update failed" + ); + return; + } + + if has_lkg { + let _ = state + .score_remove(&catalog_failed_versions_key(identity), &member) + .await; + trim_catalog_versions_locked(state, identity).await; + } else { + trim_failed_catalog_versions_locked(state, identity).await; + } +} + +async fn trim_catalog_versions_locked(state: &RuntimeState, identity: &CodexCatalogIdentity) { + let index_key = catalog_versions_key(identity); + let Ok(versions) = state.score_range_by_min(&index_key, 0.0).await else { + return; + }; + let trim_count = versions + .len() + .saturating_sub(CODEX_CATALOG_MAX_VERSIONS_PER_KEY); + for member in versions.into_iter().take(trim_count) { + let Some((credential_scope, version)) = parse_catalog_version_member(&member) else { + let _ = state.score_remove(&index_key, &member).await; + continue; + }; + let keys = vec![ + catalog_versioned_key("lkg", identity, credential_scope, version), + catalog_versioned_key("fresh", identity, credential_scope, version), + catalog_versioned_key("status", identity, credential_scope, version), + catalog_versioned_key("retry", identity, credential_scope, version), + ]; + let _ = state.kv_delete_many(&keys).await; + let _ = state.score_remove(&index_key, &member).await; + } +} + +async fn trim_failed_catalog_versions_locked( + state: &RuntimeState, + identity: &CodexCatalogIdentity, +) { + let index_key = catalog_failed_versions_key(identity); + let Ok(versions) = state.score_range_by_min(&index_key, 0.0).await else { + return; + }; + let trim_count = versions + .len() + .saturating_sub(CODEX_CATALOG_MAX_FAILED_VERSIONS_PER_KEY); + for member in versions.into_iter().take(trim_count) { + // The retry marker has its own short TTL. Keep it until natural expiry so rotating more + // than the retained failure-status count cannot bypass the per-version cooldown. + if let Some((credential_scope, version)) = parse_catalog_version_member(&member) { + let _ = state + .kv_delete(&catalog_versioned_key( + "status", + identity, + credential_scope, + version, + )) + .await; + } + let _ = state.score_remove(&index_key, &member).await; + } +} + +fn normalize_etag(value: &str) -> Option { + let value = value.trim(); + let opaque = value.strip_prefix("W/").unwrap_or(value); + let valid_syntax = opaque.len() >= 2 + && opaque.starts_with('"') + && opaque.ends_with('"') + && opaque[1..opaque.len() - 1] + .bytes() + .all(|byte| byte == 0x21 || (0x23..=0x7e).contains(&byte) || byte >= 0x80); + (!value.is_empty() && value.len() <= CODEX_CATALOG_MAX_ETAG_LEN && valid_syntax) + .then(|| value.to_string()) +} + +fn safe_catalog_error(error: &str) -> String { + let lower = error.to_ascii_lowercase(); + if lower.contains("401") || lower.contains("unauthorized") { + "upstream catalog authentication failed (status 401)".to_string() + } else if lower.contains("403") || lower.contains("forbidden") { + "upstream catalog authorization failed (status 403)".to_string() + } else if lower.contains("429") { + "upstream catalog rate limited (status 429)".to_string() + } else if lower.contains("timeout") || lower.contains("timed out") { + "upstream catalog fetch timed out".to_string() + } else if lower.contains("no models") || lower.contains("missing") || lower.contains("invalid") + { + "upstream catalog response was invalid".to_string() + } else { + "upstream catalog fetch failed".to_string() + } +} + +fn catalog_lkg_key(target: &CodexCatalogTarget, client_version: &str) -> String { + catalog_versioned_key( + "lkg", + &target.identity, + target.credential_scope().unwrap_or_default(), + client_version, + ) +} + +fn catalog_fresh_key(target: &CodexCatalogTarget, client_version: &str) -> String { + catalog_versioned_key( + "fresh", + &target.identity, + target.credential_scope().unwrap_or_default(), + client_version, + ) +} + +fn catalog_status_key(target: &CodexCatalogTarget, client_version: &str) -> String { + catalog_versioned_key( + "status", + &target.identity, + target.credential_scope().unwrap_or_default(), + client_version, + ) +} + +fn catalog_retry_key(target: &CodexCatalogTarget, client_version: &str) -> String { + catalog_versioned_key( + "retry", + &target.identity, + target.credential_scope().unwrap_or_default(), + client_version, + ) +} + +fn catalog_flight_key(target: &CodexCatalogTarget, client_version: &str) -> String { + catalog_versioned_key( + "flight", + &target.identity, + target.credential_scope().unwrap_or_default(), + client_version, + ) +} + +fn catalog_fetch_pace_key(target: &CodexCatalogTarget) -> String { + format!( + "codex_catalog:fetch_pace:v2:{}", + catalog_scoped_identity_digest(target) + ) +} + +fn catalog_versioned_key( + namespace: &str, + identity: &CodexCatalogIdentity, + credential_scope: &str, + client_version: &str, +) -> String { + let digest = sha256_hex( + format!( + "{}\0{}\0{}\0{}", + identity.provider_id, identity.key_id, credential_scope, client_version + ) + .as_bytes(), + ); + format!("codex_catalog:{namespace}:v2:{digest}") +} + +fn catalog_versions_key(identity: &CodexCatalogIdentity) -> String { + format!( + "codex_catalog:versions:v2:{}", + catalog_identity_digest(identity) + ) +} + +fn catalog_failed_versions_key(identity: &CodexCatalogIdentity) -> String { + format!( + "codex_catalog:failed_versions:v2:{}", + catalog_identity_digest(identity) + ) +} + +fn catalog_retention_lock_key(identity: &CodexCatalogIdentity) -> String { + format!( + "codex_catalog:retention_lock:v2:{}", + catalog_identity_digest(identity) + ) +} + +fn catalog_recent_version_key(identity: &CodexCatalogIdentity) -> String { + format!( + "codex_catalog:recent:v2:{}", + catalog_identity_digest(identity) + ) +} + +fn catalog_scoped_identity_digest(target: &CodexCatalogTarget) -> String { + sha256_hex( + format!( + "{}\0{}\0{}", + target.identity.provider_id, + target.identity.key_id, + target.credential_scope().unwrap_or_default() + ) + .as_bytes(), + ) +} + +fn catalog_version_member(target: &CodexCatalogTarget, client_version: &str) -> String { + format!( + "{}:{client_version}", + target.credential_scope().unwrap_or_default() + ) +} + +fn parse_catalog_version_member(member: &str) -> Option<(&str, &str)> { + let (credential_scope, client_version) = member.split_once(':')?; + (credential_scope.len() == 64 + && credential_scope + .bytes() + .all(|byte| byte.is_ascii_hexdigit()) + && !client_version.is_empty()) + .then_some((credential_scope, client_version)) +} + +fn catalog_identity_digest(identity: &CodexCatalogIdentity) -> String { + sha256_hex(format!("{}\0{}", identity.provider_id, identity.key_id).as_bytes()) +} + +fn sha256_hex(value: &[u8]) -> String { + let digest = Sha256::digest(value); + digest.iter().map(|byte| format!("{byte:02x}")).collect() +} + +fn current_unix_secs() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs() +} + +fn current_unix_ms() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_millis() + .try_into() + .unwrap_or(u64::MAX) +} + +fn duration_ms(duration: Duration) -> u64 { + duration.as_millis().try_into().unwrap_or(u64::MAX) +} + +#[cfg(test)] +mod tests { + use std::collections::VecDeque; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::{Arc, Mutex}; + + use aether_contracts::{ExecutionPlan, ExecutionResult, ProxySnapshot, ResponseBody}; + use aether_provider_transport::snapshot::{ + GatewayProviderTransportEndpoint, GatewayProviderTransportKey, + GatewayProviderTransportProvider, + }; + use aether_provider_transport::LocalResolvedOAuthRequestAuth; + use aether_runtime_state::MemoryRuntimeStateConfig; + use async_trait::async_trait; + use serde_json::json; + + use super::*; + + const TEST_PROVIDER_ID: &str = "provider-codex"; + const TEST_ENDPOINT_ID: &str = "endpoint-responses"; + const TEST_SECONDARY_ENDPOINT_ID: &str = "endpoint-responses-secondary"; + const TEST_KEY_ID: &str = "key-oauth"; + const TEST_OAUTH_TOKEN: &str = "oauth-super-secret-token"; + const TEST_CREDENTIAL_GENERATION_A: &str = "0198-aether-catalog-generation-a"; + const TEST_CREDENTIAL_GENERATION_B: &str = "0198-aether-catalog-generation-b"; + + struct QueuedExecution { + delay: Duration, + result: Result, + } + + impl QueuedExecution { + fn result(result: ExecutionResult) -> Self { + Self { + delay: Duration::ZERO, + result: Ok(result), + } + } + + fn delayed(delay: Duration, result: ExecutionResult) -> Self { + Self { + delay, + result: Ok(result), + } + } + + fn error(error: impl Into) -> Self { + Self { + delay: Duration::ZERO, + result: Err(error.into()), + } + } + } + + #[derive(Clone)] + struct TestRuntime { + state: RuntimeState, + transport: Arc>, + strong_credential_scope: Arc, String>>>, + strong_read_count: Arc, + rebind_after_strong_read: Arc>>, + execution_results: Arc>>, + executed_plans: Arc>>, + execution_count: Arc, + unavailable_endpoint_ids: Arc>>, + fetch_pace_ttl: Duration, + } + + impl TestRuntime { + fn new(execution_results: Vec) -> Self { + Self { + state: RuntimeState::memory(MemoryRuntimeStateConfig::default()), + transport: Arc::new(Mutex::new(sample_codex_transport())), + strong_credential_scope: Arc::new(Mutex::new(Ok(Some(test_credential_scope( + TEST_CREDENTIAL_GENERATION_A, + ))))), + strong_read_count: Arc::new(AtomicUsize::new(0)), + rebind_after_strong_read: Arc::new(Mutex::new(None)), + execution_results: Arc::new(Mutex::new(VecDeque::from(execution_results))), + executed_plans: Arc::new(Mutex::new(Vec::new())), + execution_count: Arc::new(AtomicUsize::new(0)), + unavailable_endpoint_ids: Arc::new(Mutex::new(BTreeSet::new())), + fetch_pace_ttl: CODEX_CATALOG_FETCH_PACE_TTL, + } + } + + fn with_fetch_pace_ttl(mut self, fetch_pace_ttl: Duration) -> Self { + self.fetch_pace_ttl = fetch_pace_ttl; + self + } + + fn enqueue(&self, execution: QueuedExecution) { + self.execution_results + .lock() + .expect("execution results mutex") + .push_back(execution); + } + + fn execution_count(&self) -> usize { + self.execution_count.load(Ordering::SeqCst) + } + + fn executed_plans(&self) -> Vec { + self.executed_plans + .lock() + .expect("executed plans mutex") + .clone() + } + + fn mark_endpoint_unavailable(&self, endpoint_id: &str) { + self.unavailable_endpoint_ids + .lock() + .expect("unavailable endpoint ids mutex") + .insert(endpoint_id.to_string()); + } + + fn set_credential_generation(&self, generation: &str) { + self.transport + .lock() + .expect("transport mutex") + .key + .upstream_metadata = Some(json!({ + "codex": { + "account_id": format!("account-{generation}"), + "credential_generation": generation, + } + })); + *self + .strong_credential_scope + .lock() + .expect("strong credential scope mutex") = + Ok(Some(test_credential_scope(generation))); + } + + fn set_strong_credential_scope(&self, scope: Result, String>) { + *self + .strong_credential_scope + .lock() + .expect("strong credential scope mutex") = scope; + } + + fn refresh_transport_secrets(&self) { + let mut transport = self.transport.lock().expect("transport mutex"); + transport.key.decrypted_api_key = "refreshed-oauth-placeholder".to_string(); + transport.key.decrypted_auth_config = Some( + json!({"refresh_token": "refreshed-token", "access_token": "new-access"}) + .to_string(), + ); + } + + fn rebind_after_strong_read(&self, read_number: usize, generation: &str) { + *self + .rebind_after_strong_read + .lock() + .expect("strong read rebind hook mutex") = + Some((read_number, generation.to_string())); + } + } + + #[async_trait] + impl ModelFetchTransportRuntime for TestRuntime { + async fn resolve_local_oauth_request_auth( + &self, + _transport: &GatewayProviderTransportSnapshot, + ) -> Result, String> { + Ok(Some(LocalResolvedOAuthRequestAuth::Header { + name: "authorization".to_string(), + value: format!("Bearer {TEST_OAUTH_TOKEN}"), + })) + } + + async fn resolve_model_fetch_proxy( + &self, + _transport: &GatewayProviderTransportSnapshot, + ) -> Option { + None + } + + async fn execute_model_fetch_execution_plan( + &self, + plan: &ExecutionPlan, + ) -> Result { + self.execution_count.fetch_add(1, Ordering::SeqCst); + self.executed_plans + .lock() + .expect("executed plans mutex") + .push(plan.clone()); + let queued = self + .execution_results + .lock() + .expect("execution results mutex") + .pop_front() + .ok_or_else(|| "missing queued execution result".to_string())?; + if !queued.delay.is_zero() { + tokio::time::sleep(queued.delay).await; + } + queued.result.map(|mut result| { + result.request_id.clone_from(&plan.request_id); + result.candidate_id.clone_from(&plan.candidate_id); + result + }) + } + } + + #[async_trait] + impl CodexCatalogRuntime for TestRuntime { + fn codex_catalog_runtime_state(&self) -> &RuntimeState { + &self.state + } + + fn codex_catalog_fetch_pace_ttl(&self) -> Duration { + self.fetch_pace_ttl + } + + async fn read_codex_catalog_transport_snapshot( + &self, + provider_id: &str, + endpoint_id: &str, + key_id: &str, + ) -> Result, String> { + if self + .unavailable_endpoint_ids + .lock() + .expect("unavailable endpoint ids mutex") + .contains(endpoint_id) + { + return Ok(None); + } + if provider_id != TEST_PROVIDER_ID + || !matches!(endpoint_id, TEST_ENDPOINT_ID | TEST_SECONDARY_ENDPOINT_ID) + { + return Ok(None); + } + let mut transport = self.transport.lock().expect("transport mutex").clone(); + transport.endpoint.id = endpoint_id.to_string(); + transport.key.id = key_id.to_string(); + Ok(Some(transport)) + } + + async fn read_codex_catalog_credential_scope_strong( + &self, + provider_id: &str, + _key_id: &str, + ) -> Result, String> { + if provider_id != TEST_PROVIDER_ID { + return Ok(None); + } + let scope = self + .strong_credential_scope + .lock() + .expect("strong credential scope mutex") + .clone(); + let read_number = self.strong_read_count.fetch_add(1, Ordering::SeqCst) + 1; + let rebind_generation = { + let mut hook = self + .rebind_after_strong_read + .lock() + .expect("strong read rebind hook mutex"); + match hook.as_ref() { + Some((trigger, _)) if *trigger == read_number => { + hook.take().map(|(_, generation)| generation) + } + _ => None, + } + }; + if let Some(generation) = rebind_generation { + self.set_credential_generation(&generation); + } + scope + } + } + + fn sample_codex_transport() -> GatewayProviderTransportSnapshot { + GatewayProviderTransportSnapshot { + provider: GatewayProviderTransportProvider { + id: TEST_PROVIDER_ID.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: TEST_ENDPOINT_ID.to_string(), + provider_id: TEST_PROVIDER_ID.to_string(), + api_format: "openai:responses".to_string(), + api_family: Some("openai".to_string()), + endpoint_kind: Some("responses".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: Some("/responses".to_string()), + config: None, + format_acceptance_config: None, + proxy: None, + }, + key: GatewayProviderTransportKey { + id: TEST_KEY_ID.to_string(), + provider_id: TEST_PROVIDER_ID.to_string(), + name: "primary".to_string(), + auth_type: "oauth".to_string(), + is_active: true, + api_formats: Some(vec!["openai:responses".to_string()]), + auth_type_by_format: None, + allow_auth_channel_mismatch_formats: None, + allowed_models: None, + capabilities: None, + rate_multipliers: None, + global_priority_by_format: None, + expires_at_unix_secs: None, + proxy: None, + fingerprint: None, + upstream_metadata: Some(json!({ + "codex": { + "account_id": "account-test", + "credential_generation": TEST_CREDENTIAL_GENERATION_A, + } + })), + decrypted_api_key: "oauth-placeholder".to_string(), + decrypted_auth_config: None, + }, + } + } + + fn target() -> CodexCatalogTarget { + target_for_key(TEST_KEY_ID) + } + + fn target_for_key(key_id: &str) -> CodexCatalogTarget { + CodexCatalogTarget { + identity: CodexCatalogIdentity::new(TEST_PROVIDER_ID, key_id), + endpoint_ids: vec![TEST_ENDPOINT_ID.to_string()], + credential_scope: Some(test_credential_scope(TEST_CREDENTIAL_GENERATION_A)), + } + } + + fn target_for_generation(generation: &str) -> CodexCatalogTarget { + CodexCatalogTarget { + identity: CodexCatalogIdentity::new(TEST_PROVIDER_ID, TEST_KEY_ID), + endpoint_ids: vec![TEST_ENDPOINT_ID.to_string()], + credential_scope: Some(test_credential_scope(generation)), + } + } + + fn test_credential_scope(generation: &str) -> String { + credential_scope_digest("generation", generation).expect("test credential scope") + } + + fn version(value: &str) -> NormalizedCodexClientVersion { + let normalized = normalize_codex_client_version(Some(value)); + assert!(!normalized.used_fallback(), "test version must be valid"); + normalized + } + + fn codex_model(slug: &str) -> Value { + json!({ + "id": slug, + "slug": slug, + "display_name": format!("Dynamic {slug}"), + "model_messages": {"instructions_template": format!("Instructions for {slug}")}, + "available_in_plans": ["plus"], + "future_capability": {"opaque": true} + }) + } + + fn execution_result( + status_code: u16, + json_body: Option, + etag: Option<&str>, + ) -> ExecutionResult { + ExecutionResult { + request_id: "queued-request".to_string(), + candidate_id: None, + status_code, + headers: etag + .map(|etag| BTreeMap::from([("ETag".to_string(), etag.to_string())])) + .unwrap_or_default(), + response_observation: None, + body: Some(ResponseBody { + json_body, + body_bytes_b64: None, + }), + telemetry: None, + error: None, + } + } + + fn successful_execution(slug: &str, etag: &str) -> QueuedExecution { + QueuedExecution::result(execution_result( + 200, + Some(json!({"models": [codex_model(slug)]})), + Some(etag), + )) + } + + async fn load_one( + runtime: &TestRuntime, + client_version: &NormalizedCodexClientVersion, + ) -> CodexCatalogLoad { + load_codex_catalogs(runtime, &[target()], client_version).await + } + + async fn seed_catalog( + runtime: &TestRuntime, + client_version: &NormalizedCodexClientVersion, + ) -> CodexCatalogSnapshot { + load_one(runtime, client_version) + .await + .snapshot(TEST_PROVIDER_ID, TEST_KEY_ID) + .cloned() + .expect("seed catalog snapshot") + } + + async fn read_raw_snapshot( + runtime: &TestRuntime, + target: &CodexCatalogTarget, + client_version: &str, + ) -> Option { + runtime + .state + .kv_get(&catalog_lkg_key(target, client_version)) + .await + .expect("read raw catalog LKG") + .map(|raw| serde_json::from_str(&raw).expect("decode raw catalog LKG")) + } + + async fn wait_for_execution_count(runtime: &TestRuntime, expected: usize) { + tokio::time::timeout(Duration::from_secs(1), async { + while runtime.execution_count() < expected { + tokio::task::yield_now().await; + } + }) + .await + .expect("model fetch execution should start"); + } + + async fn read_status( + runtime: &TestRuntime, + client_version: &NormalizedCodexClientVersion, + ) -> (String, CodexCatalogStatus) { + let raw = runtime + .state + .kv_get(&catalog_status_key(&target(), client_version.as_str())) + .await + .expect("read catalog status") + .expect("catalog status should exist"); + let status = serde_json::from_str(&raw).expect("decode catalog status"); + (raw, status) + } + + #[test] + fn client_version_normalization_is_strict_and_drops_suffixes() { + let stable = normalize_codex_client_version(Some("0.145.2")); + assert_eq!(stable.as_str(), "0.145.2"); + assert!(!stable.used_fallback()); + + let prerelease = normalize_codex_client_version(Some("0.146.0-beta.2+desktop.7")); + assert_eq!(prerelease.as_str(), "0.146.0"); + assert!(!prerelease.used_fallback()); + } + + #[test] + fn invalid_or_oversized_client_versions_use_bounded_fallback_identity() { + for raw in [ + "not-a-version".to_string(), + "1.2".to_string(), + format!("1.2.3-{}", "x".repeat(CODEX_CLIENT_VERSION_MAX_LEN)), + ] { + let normalized = normalize_codex_client_version(Some(&raw)); + assert_eq!(normalized.as_str(), crate::ai_serving::CODEX_CLIENT_VERSION); + assert!(normalized.used_fallback()); + assert!(!catalog_lkg_key(&target(), normalized.as_str()).contains(&raw)); + } + } + + #[test] + fn model_validation_is_schema_opaque_but_requires_identity() { + let models = vec![serde_json::json!({ + "slug": "gpt-future-dynamic", + "model_messages": {"instructions_template": "Future instructions"}, + "future_capability": {"mode": "opaque"} + })]; + assert!(validate_catalog_models(&models).is_ok()); + assert!(validate_catalog_models(&[serde_json::json!({ + "display_name": "Missing identity" + })]) + .is_err()); + assert!(validate_catalog_models(&[serde_json::json!({ + "slug": " gpt-whitespace-identity " + })]) + .is_err()); + assert!(validate_catalog_models(&[serde_json::json!({ + "slug": "gpt+future@dynamic" + })]) + .is_ok()); + } + + #[test] + fn error_status_is_sanitized_before_persistence_or_logging() { + let secret = "Bearer oauth-secret-token"; + let safe = safe_catalog_error(&format!("connection reset: {secret}")); + assert_eq!(safe, "upstream catalog fetch failed"); + assert!(!safe.contains(secret)); + } + + #[test] + fn local_singleflight_fallback_is_keyed_by_scoped_flight_identity() { + let account_a_key = catalog_flight_key( + &target_for_generation(TEST_CREDENTIAL_GENERATION_A), + "0.145.2", + ); + let account_b_key = catalog_flight_key( + &target_for_generation(TEST_CREDENTIAL_GENERATION_B), + "0.145.2", + ); + let account_a = local_catalog_fallback_flight(&account_a_key); + let account_a_follower = local_catalog_fallback_flight(&account_a_key); + let account_b = local_catalog_fallback_flight(&account_b_key); + + assert!(Arc::ptr_eq(&account_a, &account_a_follower)); + assert!(!Arc::ptr_eq(&account_a, &account_b)); + + drop(account_a_follower); + prune_local_catalog_fallback_flight(&account_a_key, &account_a); + drop(account_a); + drop(account_b); + let cleanup_key = format!("{account_b_key}:cleanup"); + let cleanup = local_catalog_fallback_flight(&cleanup_key); + prune_local_catalog_fallback_flight(&cleanup_key, &cleanup); + drop(cleanup); + let flights = CODEX_CATALOG_LOCAL_FALLBACK_FLIGHTS + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + assert!(!flights.contains_key(&account_a_key)); + assert!(!flights.contains_key(&account_b_key)); + assert!(!flights.contains_key(&cleanup_key)); + } + + #[test] + fn upstream_etag_requires_valid_entity_tag_syntax() { + assert_eq!( + normalize_etag("\"catalog-v1\"").as_deref(), + Some("\"catalog-v1\"") + ); + assert_eq!( + normalize_etag("W/\"catalog-v1\"").as_deref(), + Some("W/\"catalog-v1\"") + ); + assert_eq!(normalize_etag("catalog-v1"), None); + assert_eq!(normalize_etag("\"bad\"quote\""), None); + } + + #[tokio::test] + async fn catalogs_are_isolated_by_normalized_client_version() { + let runtime = TestRuntime::new(vec![ + successful_execution("gpt-dynamic-v145", "\"etag-v145\""), + successful_execution("gpt-dynamic-v146", "\"etag-v146\""), + ]); + let v145 = version("0.145.2"); + let v146 = version("0.146.0"); + + let first = seed_catalog(&runtime, &v145).await; + tokio::time::sleep(Duration::from_millis(25)).await; + let second = seed_catalog(&runtime, &v146).await; + let first_again = seed_catalog(&runtime, &v145).await; + + assert_eq!(first.client_version, "0.145.2"); + assert_eq!(second.client_version, "0.146.0"); + assert_eq!(first.models[0]["slug"], "gpt-dynamic-v145"); + assert_eq!(second.models[0]["slug"], "gpt-dynamic-v146"); + assert_eq!(first_again, first); + assert_eq!(first.etag.as_deref(), Some("\"etag-v145\"")); + assert_eq!(second.etag.as_deref(), Some("\"etag-v146\"")); + assert_eq!(runtime.execution_count(), 2); + + let plans = runtime.executed_plans(); + assert_eq!(plans.len(), 2); + assert!(plans[0].url.ends_with("/models?client_version=0.145.2")); + assert!(plans[1].url.ends_with("/models?client_version=0.146.0")); + assert_eq!( + plans[0].headers.get("user-agent").map(String::as_str), + Some("codex_cli_rs/0.145.2") + ); + assert_eq!( + plans[1].headers.get("user-agent").map(String::as_str), + Some("codex_cli_rs/0.146.0") + ); + + let target = target(); + assert_ne!( + catalog_lkg_key(&target, v145.as_str()), + catalog_lkg_key(&target, v146.as_str()) + ); + } + + #[tokio::test] + async fn credential_rebind_never_returns_previous_scope_lkg_on_503_or_401() { + for status_code in [503, 401] { + let runtime = TestRuntime::new(vec![ + successful_execution("gpt-account-a", "\"etag-account-a\""), + QueuedExecution::result(execution_result( + status_code, + Some(json!({"error": {"message": "account B unavailable"}})), + None, + )), + ]); + let client_version = version("0.145.2"); + let account_a = seed_catalog(&runtime, &client_version).await; + let target_a = target_for_generation(TEST_CREDENTIAL_GENERATION_A); + + runtime.set_credential_generation(TEST_CREDENTIAL_GENERATION_B); + let account_b_load = load_one(&runtime, &client_version).await; + let target_b = target_for_generation(TEST_CREDENTIAL_GENERATION_B); + + assert!(account_b_load + .snapshot(TEST_PROVIDER_ID, TEST_KEY_ID) + .is_none()); + assert_eq!(runtime.execution_count(), 2); + assert_eq!( + read_raw_snapshot(&runtime, &target_a, client_version.as_str()).await, + Some(account_a) + ); + assert!( + read_raw_snapshot(&runtime, &target_b, client_version.as_str()) + .await + .is_none() + ); + assert_ne!( + catalog_lkg_key(&target_a, client_version.as_str()), + catalog_lkg_key(&target_b, client_version.as_str()) + ); + } + } + + #[tokio::test] + async fn previous_scope_fresh_retry_flight_and_pace_do_not_suppress_rebound_scope() { + let runtime = TestRuntime::new(vec![ + successful_execution("gpt-account-a", "\"etag-account-a\""), + successful_execution("gpt-account-b", "\"etag-account-b\""), + ]) + .with_fetch_pace_ttl(Duration::from_secs(2)); + let client_version = version("0.145.2"); + let _ = seed_catalog(&runtime, &client_version).await; + let target_a = target_for_generation(TEST_CREDENTIAL_GENERATION_A); + runtime + .state + .kv_set( + &catalog_retry_key(&target_a, client_version.as_str()), + "account-a-retry".to_string(), + Some(Duration::from_secs(10)), + ) + .await + .expect("seed account A retry marker"); + let account_a_flight = runtime + .state + .lock_try_acquire( + &catalog_flight_key(&target_a, client_version.as_str()), + "account-a-flight-holder", + Duration::from_secs(10), + ) + .await + .expect("acquire account A flight") + .expect("account A flight should be available"); + + runtime.set_credential_generation(TEST_CREDENTIAL_GENERATION_B); + let target_b = target_for_generation(TEST_CREDENTIAL_GENERATION_B); + let account_b_load = + tokio::time::timeout(Duration::from_secs(1), load_one(&runtime, &client_version)) + .await + .expect("account B fetch must not wait on account A markers"); + + assert_eq!( + account_b_load + .snapshot(TEST_PROVIDER_ID, TEST_KEY_ID) + .expect("account B snapshot") + .models[0]["slug"], + "gpt-account-b" + ); + assert_eq!(runtime.execution_count(), 2); + assert_ne!( + catalog_fresh_key(&target_a, client_version.as_str()), + catalog_fresh_key(&target_b, client_version.as_str()) + ); + assert_ne!( + catalog_retry_key(&target_a, client_version.as_str()), + catalog_retry_key(&target_b, client_version.as_str()) + ); + assert_ne!( + catalog_flight_key(&target_a, client_version.as_str()), + catalog_flight_key(&target_b, client_version.as_str()) + ); + assert_ne!( + catalog_fetch_pace_key(&target_a), + catalog_fetch_pace_key(&target_b) + ); + assert!(runtime + .state + .lock_release(&account_a_flight) + .await + .expect("release account A flight")); + } + + #[tokio::test] + async fn token_refresh_with_same_generation_keeps_the_same_lkg_scope() { + let runtime = TestRuntime::new(vec![successful_execution( + "gpt-stable-account", + "\"etag-stable-account\"", + )]); + let client_version = version("0.145.2"); + let seeded = seed_catalog(&runtime, &client_version).await; + let target_before = target_for_generation(TEST_CREDENTIAL_GENERATION_A); + + runtime.refresh_transport_secrets(); + let after_refresh = load_one(&runtime, &client_version).await; + let target_after = target_for_generation(TEST_CREDENTIAL_GENERATION_A); + + assert_eq!( + after_refresh.snapshot(TEST_PROVIDER_ID, TEST_KEY_ID), + Some(&seeded) + ); + assert_eq!(runtime.execution_count(), 1); + assert_eq!( + catalog_lkg_key(&target_before, client_version.as_str()), + catalog_lkg_key(&target_after, client_version.as_str()) + ); + } + + #[tokio::test] + async fn rebind_between_initial_strong_bind_and_lkg_return_fails_closed() { + let runtime = TestRuntime::new(vec![successful_execution( + "gpt-account-a-cached", + "\"etag-account-a-cached\"", + )]); + let client_version = version("0.145.2"); + let seeded = seed_catalog(&runtime, &client_version).await; + let target_a = target_for_generation(TEST_CREDENTIAL_GENERATION_A); + let target_b = target_for_generation(TEST_CREDENTIAL_GENERATION_B); + runtime.rebind_after_strong_read(4, TEST_CREDENTIAL_GENERATION_B); + + let load = load_one(&runtime, &client_version).await; + + assert!(load.snapshot(TEST_PROVIDER_ID, TEST_KEY_ID).is_none()); + assert_eq!(runtime.execution_count(), 1); + assert_eq!( + read_raw_snapshot(&runtime, &target_a, client_version.as_str()).await, + Some(seeded) + ); + assert!( + read_raw_snapshot(&runtime, &target_b, client_version.as_str()) + .await + .is_none() + ); + } + + #[tokio::test] + async fn in_flight_previous_scope_fetch_is_discarded_after_rebind() { + let runtime = TestRuntime::new(vec![ + QueuedExecution::delayed( + Duration::from_millis(150), + execution_result( + 200, + Some(json!({"models": [codex_model("gpt-delayed-account-a")]})), + Some("\"etag-delayed-account-a\""), + ), + ), + successful_execution("gpt-account-b", "\"etag-account-b\""), + ]); + let client_version = version("0.145.2"); + let task_runtime = runtime.clone(); + let task_version = client_version.clone(); + let account_a_fetch = + tokio::spawn(async move { load_one(&task_runtime, &task_version).await }); + wait_for_execution_count(&runtime, 1).await; + + runtime.set_credential_generation(TEST_CREDENTIAL_GENERATION_B); + let account_a_load = account_a_fetch.await.expect("join account A fetch"); + let target_a = target_for_generation(TEST_CREDENTIAL_GENERATION_A); + let target_b = target_for_generation(TEST_CREDENTIAL_GENERATION_B); + + assert!(account_a_load + .snapshot(TEST_PROVIDER_ID, TEST_KEY_ID) + .is_none()); + assert!( + read_lkg_snapshot(&runtime, &target_a, client_version.as_str()) + .await + .is_none() + ); + assert!( + read_lkg_snapshot(&runtime, &target_b, client_version.as_str()) + .await + .is_none() + ); + + let account_b_load = load_one(&runtime, &client_version).await; + assert_eq!( + account_b_load + .snapshot(TEST_PROVIDER_ID, TEST_KEY_ID) + .expect("account B snapshot") + .models[0]["slug"], + "gpt-account-b" + ); + assert_eq!(runtime.execution_count(), 2); + } + + #[tokio::test] + async fn scope_change_during_persistence_removes_old_scope_success() { + let runtime = TestRuntime::new(vec![successful_execution( + "gpt-account-a-raced-persist", + "\"etag-account-a-raced-persist\"", + )]); + let client_version = version("0.145.2"); + runtime.rebind_after_strong_read(2, TEST_CREDENTIAL_GENERATION_B); + + let load = load_one(&runtime, &client_version).await; + let target_a = target_for_generation(TEST_CREDENTIAL_GENERATION_A); + let target_b = target_for_generation(TEST_CREDENTIAL_GENERATION_B); + + assert!(load.snapshot(TEST_PROVIDER_ID, TEST_KEY_ID).is_none()); + assert_eq!(runtime.execution_count(), 1); + for key in [ + catalog_lkg_key(&target_a, client_version.as_str()), + catalog_fresh_key(&target_a, client_version.as_str()), + catalog_status_key(&target_a, client_version.as_str()), + catalog_retry_key(&target_a, client_version.as_str()), + ] { + assert!(runtime + .state + .kv_get(&key) + .await + .expect("read old-scope catalog key") + .is_none()); + } + assert!(!runtime + .state + .score_range_by_min(&catalog_versions_key(&target_a.identity), 0.0) + .await + .expect("read catalog version index") + .contains(&catalog_version_member(&target_a, client_version.as_str()))); + assert!(read_recent_codex_catalog_client_version( + &runtime.state, + TEST_PROVIDER_ID, + TEST_KEY_ID, + target_b.credential_scope().expect("account B scope"), + ) + .await + .is_none()); + } + + #[tokio::test] + async fn stale_transport_scope_fails_closed_before_upstream_execution() { + let runtime = TestRuntime::new(vec![successful_execution( + "must-not-be-fetched", + "\"etag-must-not-be-fetched\"", + )]); + let client_version = version("0.145.2"); + let target_a = target_for_generation(TEST_CREDENTIAL_GENERATION_A); + let target_b = target_for_generation(TEST_CREDENTIAL_GENERATION_B); + runtime.set_strong_credential_scope(Ok(Some(test_credential_scope( + TEST_CREDENTIAL_GENERATION_B, + )))); + + let load = load_one(&runtime, &client_version).await; + + assert!(load.snapshot(TEST_PROVIDER_ID, TEST_KEY_ID).is_none()); + assert_eq!(runtime.execution_count(), 0); + assert!(runtime + .state + .kv_get(&catalog_lkg_key(&target_a, client_version.as_str())) + .await + .expect("read account A LKG") + .is_none()); + assert!(runtime + .state + .kv_get(&catalog_lkg_key(&target_b, client_version.as_str())) + .await + .expect("read account B LKG") + .is_none()); + } + + #[tokio::test] + async fn missing_or_unverifiable_strong_scope_fails_closed_before_cache_or_fetch() { + let runtime = TestRuntime::new(vec![successful_execution( + "gpt-account-a", + "\"etag-account-a\"", + )]); + let client_version = version("0.145.2"); + let seeded = seed_catalog(&runtime, &client_version).await; + let target_a = target_for_generation(TEST_CREDENTIAL_GENERATION_A); + + runtime.set_strong_credential_scope(Ok(None)); + let missing = load_one(&runtime, &client_version).await; + assert!(missing.snapshot(TEST_PROVIDER_ID, TEST_KEY_ID).is_none()); + assert_eq!(runtime.execution_count(), 1); + + runtime.set_strong_credential_scope(Err("strong repository unavailable".to_string())); + let unverifiable = load_one(&runtime, &client_version).await; + assert!(unverifiable + .snapshot(TEST_PROVIDER_ID, TEST_KEY_ID) + .is_none()); + assert_eq!(runtime.execution_count(), 1); + assert_eq!( + read_raw_snapshot(&runtime, &target_a, client_version.as_str()).await, + Some(seeded) + ); + } + + #[tokio::test] + async fn concurrent_cold_requests_for_one_identity_and_version_singleflight() { + let runtime = TestRuntime::new(vec![QueuedExecution::delayed( + Duration::from_millis(100), + execution_result( + 200, + Some(json!({ + "models": [codex_model("gpt-singleflight-dynamic")] + })), + Some("\"etag-singleflight\""), + ), + )]); + let client_version = version("0.145.2"); + + let (left, right) = tokio::join!( + load_one(&runtime, &client_version), + load_one(&runtime, &client_version) + ); + + let left = left + .snapshot(TEST_PROVIDER_ID, TEST_KEY_ID) + .expect("left snapshot"); + let right = right + .snapshot(TEST_PROVIDER_ID, TEST_KEY_ID) + .expect("right snapshot"); + assert_eq!(left, right); + assert_eq!(left.models[0]["slug"], "gpt-singleflight-dynamic"); + assert_eq!(runtime.execution_count(), 1); + assert_eq!(runtime.executed_plans().len(), 1); + } + + #[tokio::test] + async fn expired_fresh_ttl_returns_lkg_and_marks_target_stale_without_fetching() { + let runtime = TestRuntime::new(vec![successful_execution( + "gpt-stale-dynamic", + "\"etag-stale\"", + )]); + let client_version = version("0.145.2"); + let seeded = seed_catalog(&runtime, &client_version).await; + + tokio::time::sleep(CODEX_CATALOG_FRESH_TTL + Duration::from_millis(25)).await; + let load = load_one(&runtime, &client_version).await; + + assert_eq!(load.snapshot(TEST_PROVIDER_ID, TEST_KEY_ID), Some(&seeded)); + assert_eq!(load.stale_targets(), &[target()]); + assert_eq!(runtime.execution_count(), 1); + } + + #[tokio::test] + async fn fresh_and_stale_lkg_reads_do_not_wait_for_the_retention_lock() { + let runtime = TestRuntime::new(vec![successful_execution( + "gpt-retention-lock-dynamic", + "\"etag-retention-lock\"", + )]); + let client_version = version("0.145.2"); + let seeded = seed_catalog(&runtime, &client_version).await; + let target = target(); + let lease = runtime + .state + .lock_try_acquire( + &catalog_retention_lock_key(&target.identity), + "catalog-retention-test-holder", + CODEX_CATALOG_RETENTION_LOCK_TTL, + ) + .await + .expect("acquire retention lock") + .expect("retention lock should be available"); + + let fresh = tokio::time::timeout( + Duration::from_millis(250), + load_one(&runtime, &client_version), + ) + .await + .expect("fresh LKG read must not wait for retention lock"); + assert_eq!(fresh.snapshot(TEST_PROVIDER_ID, TEST_KEY_ID), Some(&seeded)); + assert!(fresh.stale_targets().is_empty()); + + assert!(runtime + .state + .kv_delete(&catalog_fresh_key(&target, client_version.as_str())) + .await + .expect("delete fresh marker")); + let stale = tokio::time::timeout( + Duration::from_millis(250), + load_one(&runtime, &client_version), + ) + .await + .expect("stale LKG read must not wait for retention lock"); + assert_eq!(stale.snapshot(TEST_PROVIDER_ID, TEST_KEY_ID), Some(&seeded)); + assert_eq!(stale.stale_targets(), &[target.clone()]); + assert_eq!(runtime.execution_count(), 1); + + assert!(runtime + .state + .lock_release(&lease) + .await + .expect("release retention lock")); + } + + #[tokio::test] + async fn deleting_fresh_marker_does_not_delete_or_hide_lkg() { + let runtime = TestRuntime::new(vec![successful_execution( + "gpt-marker-dynamic", + "\"etag-marker\"", + )]); + let client_version = version("0.145.2"); + let seeded = seed_catalog(&runtime, &client_version).await; + let target = target(); + + assert!(runtime + .state + .kv_delete(&catalog_fresh_key(&target, client_version.as_str())) + .await + .expect("delete fresh marker")); + let load = load_one(&runtime, &client_version).await; + + assert_eq!(load.snapshot(TEST_PROVIDER_ID, TEST_KEY_ID), Some(&seeded)); + assert_eq!(load.stale_targets(), &[target.clone()]); + assert!(runtime + .state + .kv_get(&catalog_lkg_key(&target, client_version.as_str())) + .await + .expect("read LKG") + .is_some()); + assert_eq!(runtime.execution_count(), 1); + } + + #[tokio::test] + async fn mismatched_fresh_marker_hash_is_treated_as_stale() { + let runtime = TestRuntime::new(vec![successful_execution( + "gpt-marker-hash-dynamic", + "\"etag-marker-hash\"", + )]); + let client_version = version("0.145.2"); + let seeded = seed_catalog(&runtime, &client_version).await; + let target = target(); + runtime + .state + .kv_set( + &catalog_fresh_key(&target, client_version.as_str()), + "not-the-snapshot-hash".to_string(), + Some(CODEX_CATALOG_FRESH_TTL), + ) + .await + .expect("replace fresh marker"); + + let load = load_one(&runtime, &client_version).await; + assert_eq!(load.snapshot(TEST_PROVIDER_ID, TEST_KEY_ID), Some(&seeded)); + assert_eq!(load.stale_targets(), &[target.clone()]); + assert_eq!(runtime.execution_count(), 1); + } + + #[tokio::test] + async fn lkg_content_hash_detects_corrupted_snapshot_payload() { + let runtime = TestRuntime::new(vec![successful_execution( + "gpt-hash-dynamic", + "\"etag-hash\"", + )]); + let client_version = version("0.145.2"); + let _ = seed_catalog(&runtime, &client_version).await; + let target = target(); + let key = catalog_lkg_key(&target, client_version.as_str()); + let raw = runtime + .state + .kv_get(&key) + .await + .expect("read LKG") + .expect("LKG should exist"); + let mut corrupted: CodexCatalogSnapshot = serde_json::from_str(&raw).expect("decode LKG"); + corrupted.models[0]["future_capability"] = json!({"corrupted": true}); + runtime + .state + .kv_set( + &key, + serde_json::to_string(&corrupted).expect("serialize corrupted LKG"), + None, + ) + .await + .expect("replace LKG"); + + assert!( + read_lkg_snapshot(&runtime, &target, client_version.as_str()) + .await + .is_none() + ); + } + + #[tokio::test] + async fn cold_failure_retry_marker_suppresses_immediate_followup_fetch() { + let runtime = TestRuntime::new(vec![QueuedExecution::error("network unavailable")]); + let client_version = version("0.145.2"); + + assert!(load_one(&runtime, &client_version) + .await + .snapshot(TEST_PROVIDER_ID, TEST_KEY_ID) + .is_none()); + assert_eq!(runtime.execution_count(), 1); + assert!(load_one(&runtime, &client_version) + .await + .snapshot(TEST_PROVIDER_ID, TEST_KEY_ID) + .is_none()); + assert_eq!(runtime.execution_count(), 1); + } + + #[tokio::test] + async fn upstream_failures_never_replace_non_empty_lkg_and_status_is_sanitized() { + let runtime = TestRuntime::new(vec![successful_execution( + "gpt-lkg-dynamic", + "\"etag-lkg\"", + )]); + let client_version = version("0.145.2"); + let seeded = seed_catalog(&runtime, &client_version).await; + let target = target(); + + let failures = [ + ( + "5xx", + QueuedExecution::result(execution_result( + 503, + Some(json!({ + "error": { + "message": format!("temporarily unavailable: Bearer {TEST_OAUTH_TOKEN}") + } + })), + None, + )), + Some(503), + ), + ( + "timeout", + QueuedExecution::delayed( + CODEX_CATALOG_FETCH_TIMEOUT + Duration::from_millis(100), + execution_result( + 200, + Some(json!({"models": [codex_model("must-not-be-stored")]})), + Some("\"etag-must-not-be-stored\""), + ), + ), + None, + ), + ( + "empty models", + QueuedExecution::result(execution_result( + 200, + Some(json!({"models": []})), + Some("\"etag-empty\""), + )), + Some(200), + ), + ( + "invalid JSON payload", + QueuedExecution::result(execution_result(200, None, Some("\"etag-invalid\""))), + Some(200), + ), + ]; + + for (label, failure, expected_status) in failures { + runtime.enqueue(failure); + tokio::time::sleep(CODEX_CATALOG_FETCH_PACE_TTL + Duration::from_millis(5)).await; + let _ = runtime + .state + .kv_delete(&catalog_retry_key(&target, client_version.as_str())) + .await; + let _ = runtime + .state + .kv_delete(&catalog_fresh_key(&target, client_version.as_str())) + .await; + refresh_codex_catalog_target(&runtime, &target, &client_version).await; + + let persisted = read_lkg_snapshot(&runtime, &target, client_version.as_str()) + .await + .unwrap_or_else(|| panic!("{label}: LKG should remain available")); + assert_eq!(persisted, seeded, "{label}: LKG was overwritten"); + + let (raw_status, status) = read_status(&runtime, &client_version).await; + assert_eq!( + status.last_http_status, expected_status, + "{label}: upstream status" + ); + assert_eq!( + status.last_success_at_unix_secs, + Some(seeded.fetched_at_unix_secs) + ); + assert_eq!(status.model_count, Some(1)); + assert!(status.used_lkg, "{label}: status should report LKG use"); + assert!( + status.last_error.is_some(), + "{label}: error should be visible" + ); + assert!( + !raw_status.contains(TEST_OAUTH_TOKEN), + "{label}: status leaked OAuth token" + ); + } + + assert_eq!(runtime.execution_count(), 5); + assert_eq!(seeded.etag.as_deref(), Some("\"etag-lkg\"")); + assert!(runtime.executed_plans().iter().all(|plan| { + plan.headers.get("authorization").map(String::as_str) + == Some("Bearer oauth-super-secret-token") + })); + } + + #[tokio::test] + async fn partial_multi_endpoint_success_does_not_replace_lkg() { + let runtime = TestRuntime::new(vec![successful_execution( + "gpt-complete-lkg", + "\"etag-complete\"", + )]); + let client_version = version("0.145.2"); + let seeded = seed_catalog(&runtime, &client_version).await; + let mut multi_endpoint_target = target(); + multi_endpoint_target + .endpoint_ids + .push(TEST_ENDPOINT_ID.to_string()); + runtime.enqueue(successful_execution( + "gpt-partial-replacement", + "\"etag-partial\"", + )); + runtime.enqueue(QueuedExecution::result(execution_result( + 503, + Some(json!({"error": {"message": "temporary endpoint failure"}})), + None, + ))); + runtime + .state + .kv_delete(&catalog_fresh_key( + &multi_endpoint_target, + client_version.as_str(), + )) + .await + .expect("delete fresh marker"); + + tokio::time::sleep(CODEX_CATALOG_FETCH_PACE_TTL + Duration::from_millis(5)).await; + refresh_codex_catalog_target(&runtime, &multi_endpoint_target, &client_version).await; + let persisted = + read_lkg_snapshot(&runtime, &multi_endpoint_target, client_version.as_str()) + .await + .expect("complete LKG should remain"); + assert_eq!(persisted, seeded); + assert_eq!(runtime.execution_count(), 3); + } + + #[tokio::test] + async fn missing_multi_endpoint_transport_does_not_replace_lkg() { + let runtime = TestRuntime::new(vec![ + successful_execution("gpt-primary", "\"etag-primary\""), + successful_execution("gpt-secondary", "\"etag-secondary\""), + ]); + let client_version = version("0.145.2"); + let mut multi_endpoint_target = target(); + multi_endpoint_target + .endpoint_ids + .push(TEST_SECONDARY_ENDPOINT_ID.to_string()); + let seeded = + load_codex_catalogs(&runtime, &[multi_endpoint_target.clone()], &client_version) + .await + .snapshot(TEST_PROVIDER_ID, TEST_KEY_ID) + .cloned() + .expect("seed complete multi-endpoint LKG"); + runtime.mark_endpoint_unavailable(TEST_SECONDARY_ENDPOINT_ID); + runtime + .state + .kv_delete(&catalog_fresh_key( + &multi_endpoint_target, + client_version.as_str(), + )) + .await + .expect("delete fresh marker"); + + tokio::time::sleep(CODEX_CATALOG_FETCH_PACE_TTL + Duration::from_millis(5)).await; + refresh_codex_catalog_target(&runtime, &multi_endpoint_target, &client_version).await; + let persisted = + read_lkg_snapshot(&runtime, &multi_endpoint_target, client_version.as_str()) + .await + .expect("complete LKG should remain"); + assert_eq!(persisted, seeded); + assert_eq!(runtime.execution_count(), 2); + let (_, status) = read_status(&runtime, &client_version).await; + assert!(status.used_lkg); + assert!(status.last_error.is_some()); + } + + #[tokio::test] + async fn catalog_status_sanitizes_transport_errors_containing_credentials() { + let runtime = TestRuntime::new(vec![QueuedExecution::error(format!( + "401 Unauthorized: Bearer {TEST_OAUTH_TOKEN}" + ))]); + let client_version = version("0.145.2"); + + let load = load_one(&runtime, &client_version).await; + assert!(load.snapshot(TEST_PROVIDER_ID, TEST_KEY_ID).is_none()); + + let (raw_status, status) = read_status(&runtime, &client_version).await; + assert_eq!( + status.last_error.as_deref(), + Some("upstream catalog authentication failed (status 401)") + ); + assert!(!status.used_lkg); + assert!(!raw_status.contains(TEST_OAUTH_TOKEN)); + assert!(!raw_status.to_ascii_lowercase().contains("bearer")); + } + + #[tokio::test] + async fn version_index_retains_only_eight_most_recent_catalogs() { + let executions = (0..10) + .map(|index| { + successful_execution( + &format!("gpt-version-{index}"), + &format!("\"etag-{index}\""), + ) + }) + .collect(); + let runtime = TestRuntime::new(executions); + let target = target(); + + for index in 0..10 { + let client_version = version(&format!("0.200.{index}")); + let snapshot = seed_catalog(&runtime, &client_version).await; + assert_eq!(snapshot.models[0]["slug"], format!("gpt-version-{index}")); + tokio::time::sleep(Duration::from_millis(25)).await; + } + + let versions = runtime + .state + .score_range_by_min(&catalog_versions_key(&target.identity), 0.0) + .await + .expect("read catalog versions"); + assert_eq!(versions.len(), CODEX_CATALOG_MAX_VERSIONS_PER_KEY); + assert!(!versions.contains(&catalog_version_member(&target, "0.200.0"))); + assert!(!versions.contains(&catalog_version_member(&target, "0.200.1"))); + for index in 2..10 { + let client_version = format!("0.200.{index}"); + assert!(versions.contains(&catalog_version_member(&target, &client_version))); + assert!(runtime + .state + .kv_get(&catalog_lkg_key(&target, &client_version)) + .await + .expect("read retained LKG") + .is_some()); + } + for trimmed in ["0.200.0", "0.200.1"] { + assert!(runtime + .state + .kv_get(&catalog_lkg_key(&target, trimmed)) + .await + .expect("read trimmed LKG") + .is_none()); + assert!(runtime + .state + .kv_get(&catalog_status_key(&target, trimmed)) + .await + .expect("read trimmed status") + .is_none()); + } + assert_eq!(runtime.execution_count(), 10); + assert_eq!( + read_recent_codex_catalog_client_version( + &runtime.state, + TEST_PROVIDER_ID, + TEST_KEY_ID, + target.credential_scope().expect("test credential scope"), + ) + .await + .as_deref(), + Some("0.200.9") + ); + } + + #[tokio::test] + async fn version_retention_is_bounded_across_all_credential_scopes_for_one_key() { + let runtime = TestRuntime::new( + (0..10) + .map(|index| { + successful_execution( + &format!("gpt-scope-version-{index}"), + &format!("\"etag-scope-version-{index}\""), + ) + }) + .collect(), + ); + let target_a = target_for_generation(TEST_CREDENTIAL_GENERATION_A); + let target_b = target_for_generation(TEST_CREDENTIAL_GENERATION_B); + + for index in 0..5 { + seed_catalog(&runtime, &version(&format!("0.300.{index}"))).await; + tokio::time::sleep(Duration::from_millis(25)).await; + } + runtime.set_credential_generation(TEST_CREDENTIAL_GENERATION_B); + for index in 5..10 { + seed_catalog(&runtime, &version(&format!("0.300.{index}"))).await; + tokio::time::sleep(Duration::from_millis(25)).await; + } + + let retained = runtime + .state + .score_range_by_min(&catalog_versions_key(&target_a.identity), 0.0) + .await + .expect("read cross-scope catalog versions"); + assert_eq!(retained.len(), CODEX_CATALOG_MAX_VERSIONS_PER_KEY); + for trimmed in ["0.300.0", "0.300.1"] { + assert!(!retained.contains(&catalog_version_member(&target_a, trimmed))); + assert!(runtime + .state + .kv_get(&catalog_lkg_key(&target_a, trimmed)) + .await + .expect("read trimmed account A LKG") + .is_none()); + } + for retained_a in ["0.300.2", "0.300.3", "0.300.4"] { + assert!(retained.contains(&catalog_version_member(&target_a, retained_a))); + assert!(runtime + .state + .kv_get(&catalog_lkg_key(&target_a, retained_a)) + .await + .expect("read retained account A LKG") + .is_some()); + } + for index in 5..10 { + let retained_b = format!("0.300.{index}"); + assert!(retained.contains(&catalog_version_member(&target_b, &retained_b))); + assert!(runtime + .state + .kv_get(&catalog_lkg_key(&target_b, &retained_b)) + .await + .expect("read retained account B LKG") + .is_some()); + } + assert!(read_recent_codex_catalog_client_version( + &runtime.state, + TEST_PROVIDER_ID, + TEST_KEY_ID, + target_a.credential_scope().expect("account A scope"), + ) + .await + .is_none()); + assert_eq!( + read_recent_codex_catalog_client_version( + &runtime.state, + TEST_PROVIDER_ID, + TEST_KEY_ID, + target_b.credential_scope().expect("account B scope"), + ) + .await + .as_deref(), + Some("0.300.9") + ); + } + + #[tokio::test] + async fn failed_client_versions_are_also_bounded_by_version_retention() { + let runtime = TestRuntime::new( + (0..10) + .map(|_| QueuedExecution::error("network unavailable")) + .collect(), + ); + let target = target(); + + for index in 0..10 { + let client_version = version(&format!("0.220.{index}")); + assert!(load_one(&runtime, &client_version) + .await + .snapshot(TEST_PROVIDER_ID, TEST_KEY_ID) + .is_none()); + tokio::time::sleep(Duration::from_millis(25)).await; + } + + let versions = runtime + .state + .score_range_by_min(&catalog_failed_versions_key(&target.identity), 0.0) + .await + .expect("read failed catalog versions"); + assert_eq!(versions.len(), CODEX_CATALOG_MAX_FAILED_VERSIONS_PER_KEY); + for trimmed in ["0.220.0", "0.220.1"] { + assert!(runtime + .state + .kv_get(&catalog_status_key(&target, trimmed)) + .await + .expect("read trimmed failure status") + .is_none()); + } + assert_eq!(runtime.execution_count(), 10); + } + + #[tokio::test] + async fn failed_version_churn_does_not_evict_non_empty_lkg_snapshots() { + let mut executions = (0..CODEX_CATALOG_MAX_VERSIONS_PER_KEY) + .map(|index| { + successful_execution( + &format!("gpt-stable-{index}"), + &format!("\"stable-etag-{index}\""), + ) + }) + .collect::>(); + executions + .extend((0..10).map(|_| QueuedExecution::error("network unavailable during churn"))); + let runtime = TestRuntime::new(executions); + let target = target(); + + for index in 0..CODEX_CATALOG_MAX_VERSIONS_PER_KEY { + seed_catalog(&runtime, &version(&format!("0.240.{index}"))).await; + tokio::time::sleep(Duration::from_millis(25)).await; + } + for index in 0..10 { + assert!(load_one(&runtime, &version(&format!("0.241.{index}"))) + .await + .snapshot(TEST_PROVIDER_ID, TEST_KEY_ID) + .is_none()); + tokio::time::sleep(Duration::from_millis(25)).await; + } + + let retained_lkg_versions = runtime + .state + .score_range_by_min(&catalog_versions_key(&target.identity), 0.0) + .await + .expect("read retained LKG versions"); + assert_eq!( + retained_lkg_versions.len(), + CODEX_CATALOG_MAX_VERSIONS_PER_KEY + ); + for index in 0..CODEX_CATALOG_MAX_VERSIONS_PER_KEY { + let client_version = format!("0.240.{index}"); + assert!( + retained_lkg_versions.contains(&catalog_version_member(&target, &client_version)) + ); + assert!(runtime + .state + .kv_get(&catalog_lkg_key(&target, &client_version)) + .await + .expect("read stable LKG") + .is_some()); + } + assert_eq!( + runtime + .state + .score_range_by_min(&catalog_failed_versions_key(&target.identity), 0.0) + .await + .expect("read retained failed versions") + .len(), + CODEX_CATALOG_MAX_FAILED_VERSIONS_PER_KEY + ); + assert_eq!( + runtime.execution_count(), + CODEX_CATALOG_MAX_VERSIONS_PER_KEY + 10 + ); + } + + #[tokio::test] + async fn concurrent_cold_versions_are_paced_per_provider_key() { + let mut executions = vec![QueuedExecution::delayed( + CODEX_CATALOG_FETCH_PACE_TTL + Duration::from_millis(30), + execution_result( + 503, + Some(json!({"error": {"message": "network unavailable"}})), + None, + ), + )]; + executions.extend((1..12).map(|_| QueuedExecution::error("network unavailable"))); + let runtime = TestRuntime::new(executions).with_fetch_pace_ttl(Duration::from_secs(2)); + let versions = (0..12) + .map(|index| version(&format!("0.230.{index}"))) + .collect::>(); + futures_util::future::join_all( + versions + .iter() + .map(|client_version| load_one(&runtime, client_version)), + ) + .await; + + assert_eq!(runtime.execution_count(), 1); + let indexed_failures = runtime + .state + .score_range_by_min(&catalog_failed_versions_key(&target().identity), 0.0) + .await + .expect("read paced failure versions"); + assert!(indexed_failures.len() <= 1); + } + + #[tokio::test] + async fn cold_fetch_budget_never_returns_a_partial_catalog_as_complete() { + let target_count = CODEX_CATALOG_MAX_COLD_FETCHES_PER_LOAD + 1; + let runtime = TestRuntime::new( + (0..target_count) + .map(|index| { + successful_execution( + &format!("gpt-budget-{index}"), + &format!("\"etag-budget-{index}\""), + ) + }) + .collect(), + ); + let targets = (0..target_count) + .map(|index| target_for_key(&format!("key-budget-{index}"))) + .collect::>(); + let client_version = version("0.231.0"); + + let first = load_codex_catalogs(&runtime, &targets, &client_version).await; + assert!(!first.is_complete()); + assert_eq!( + first.snapshots.len(), + CODEX_CATALOG_MAX_COLD_FETCHES_PER_LOAD + ); + assert_eq!( + runtime.execution_count(), + CODEX_CATALOG_MAX_COLD_FETCHES_PER_LOAD + ); + + let second = load_codex_catalogs(&runtime, &targets, &client_version).await; + assert!(second.is_complete()); + assert_eq!(second.snapshots.len(), target_count); + assert_eq!(runtime.execution_count(), target_count); + for target in &targets { + assert!(second + .snapshot(&target.identity.provider_id, &target.identity.key_id) + .is_some()); + } + } + + #[tokio::test] + async fn multiple_large_source_catalogs_cannot_exceed_the_aggregate_body_budget() { + let oversized_pair_payload = "x".repeat(CODEX_CATALOG_MAX_AGGREGATE_BODY_BYTES / 2); + let large_execution = |slug: &str, etag: &str| { + let mut model = codex_model(slug); + model["future_capability"]["large_opaque_payload"] = + Value::String(oversized_pair_payload.clone()); + QueuedExecution::result(execution_result( + 200, + Some(json!({"models": [model]})), + Some(etag), + )) + }; + let runtime = TestRuntime::new(vec![ + large_execution("gpt-large-source-a", "\"etag-large-a\""), + large_execution("gpt-large-source-b", "\"etag-large-b\""), + ]); + let targets = vec![target_for_key("key-large-a"), target_for_key("key-large-b")]; + + let load = load_codex_catalogs(&runtime, &targets, &version("0.232.0")).await; + + assert!(!load.is_complete()); + assert!(load.snapshots.is_empty()); + assert!(load.stale_targets().is_empty()); + assert_eq!(runtime.execution_count(), targets.len()); + } + + #[tokio::test] + async fn legacy_unversioned_cache_is_never_used_as_a_versioned_catalog() { + let runtime = TestRuntime::new(vec![QueuedExecution::error("network unavailable")]); + runtime + .state + .kv_set( + &format!("upstream_models:{TEST_PROVIDER_ID}:{TEST_KEY_ID}"), + serde_json::to_string(&vec![codex_model("gpt-legacy-cache")]) + .expect("serialize legacy cache"), + None, + ) + .await + .expect("seed legacy cache"); + + let load = load_one(&runtime, &version(crate::ai_serving::CODEX_CLIENT_VERSION)).await; + assert!(load.snapshot(TEST_PROVIDER_ID, TEST_KEY_ID).is_none()); + assert_eq!(runtime.execution_count(), 1); + } + + #[tokio::test] + async fn recent_version_tracks_normalized_last_seen_version_even_after_cold_failure() { + let runtime = TestRuntime::new(vec![QueuedExecution::error("network unavailable")]); + let client_version = version("0.211.0-beta.4+desktop.9"); + + let load = load_one(&runtime, &client_version).await; + assert!(load.snapshot(TEST_PROVIDER_ID, TEST_KEY_ID).is_none()); + assert_eq!(client_version.as_str(), "0.211.0"); + assert_eq!( + read_recent_codex_catalog_client_version( + &runtime.state, + TEST_PROVIDER_ID, + TEST_KEY_ID, + target().credential_scope().expect("test credential scope"), + ) + .await + .as_deref(), + Some("0.211.0") + ); + } + + #[tokio::test] + async fn invalid_version_fallback_does_not_overwrite_recent_real_client_version() { + let runtime = TestRuntime::new(vec![ + successful_execution("gpt-real-version", "\"etag-real\""), + successful_execution("gpt-fallback-version", "\"etag-fallback\""), + ]); + seed_catalog(&runtime, &version("0.212.3")).await; + tokio::time::sleep(Duration::from_millis(25)).await; + + let fallback = normalize_codex_client_version(Some("malicious-not-semver")); + assert!(fallback.used_fallback()); + seed_catalog(&runtime, &fallback).await; + + assert_eq!( + read_recent_codex_catalog_client_version( + &runtime.state, + TEST_PROVIDER_ID, + TEST_KEY_ID, + target().credential_scope().expect("test credential scope"), + ) + .await + .as_deref(), + Some("0.212.3") + ); + } +} diff --git a/apps/aether-gateway/src/model_fetch/mod.rs b/apps/aether-gateway/src/model_fetch/mod.rs index 2499b8f01..545078399 100644 --- a/apps/aether-gateway/src/model_fetch/mod.rs +++ b/apps/aether-gateway/src/model_fetch/mod.rs @@ -1,8 +1,15 @@ +mod catalog; mod runtime; #[cfg(test)] mod tests; pub(crate) use aether_model_fetch::ModelFetchRunSummary; +pub(crate) use catalog::{ + codex_catalog_credential_scope_from_stored_key, codex_catalog_targets, load_codex_catalogs, + normalize_codex_client_version, read_recent_codex_catalog_client_version, + refresh_codex_catalog_target, CodexCatalogLoad, CodexCatalogRuntime, CodexCatalogTarget, + NormalizedCodexClientVersion, +}; pub(crate) use runtime::state::ModelFetchRuntimeState; pub(crate) use runtime::{ perform_model_fetch_for_key, perform_model_fetch_for_keys, perform_model_fetch_once, diff --git a/apps/aether-gateway/src/model_fetch/runtime.rs b/apps/aether-gateway/src/model_fetch/runtime.rs index 748381e92..b68c836c2 100644 --- a/apps/aether-gateway/src/model_fetch/runtime.rs +++ b/apps/aether-gateway/src/model_fetch/runtime.rs @@ -7,7 +7,7 @@ use aether_data_contracts::repository::provider_catalog::{ StoredProviderCatalogKey, StoredProviderCatalogProvider, }; use aether_model_fetch::{ - apply_model_filters, fetch_models_from_transports, json_string_list, + apply_model_filters, fetch_models_from_transports_for_client_version, json_string_list, model_catalog_upstream_metadata, model_fetch_interval_minutes, model_fetch_startup_delay_seconds, model_fetch_startup_enabled, preset_models_for_provider, selected_models_fetch_endpoints, sync_provider_model_whitelist_associations, @@ -315,7 +315,25 @@ async fn fetch_and_persist_key_models( return Ok(KeyFetchDisposition::Skipped); } - let result = match fetch_models_from_transports(state, &transports).await { + let codex_client_version = if target + .provider + .provider_type + .trim() + .eq_ignore_ascii_case("codex") + { + state + .read_recent_codex_catalog_client_version(&target.provider.id, &target.key.id) + .await + } else { + None + }; + let result = match fetch_models_from_transports_for_client_version( + state, + &transports, + codex_client_version.as_deref(), + ) + .await + { Ok(result) => result, Err(err) => { persist_key_fetch_failure(state, &target.key, now_unix_secs, err.clone()).await?; diff --git a/apps/aether-gateway/src/model_fetch/runtime/state.rs b/apps/aether-gateway/src/model_fetch/runtime/state.rs index 81d23a5e0..e90d720fa 100644 --- a/apps/aether-gateway/src/model_fetch/runtime/state.rs +++ b/apps/aether-gateway/src/model_fetch/runtime/state.rs @@ -43,6 +43,14 @@ pub(crate) trait ModelFetchRuntimeState: plan: &ExecutionPlan, ) -> Result; + async fn read_recent_codex_catalog_client_version( + &self, + _provider_id: &str, + _key_id: &str, + ) -> Option { + None + } + async fn update_provider_catalog_key_model_fetch_state( &self, key_id: &str, diff --git a/apps/aether-gateway/src/request_candidate_queue.rs b/apps/aether-gateway/src/request_candidate_queue.rs index 2456abbb5..bc61aa2db 100644 --- a/apps/aether-gateway/src/request_candidate_queue.rs +++ b/apps/aether-gateway/src/request_candidate_queue.rs @@ -3838,7 +3838,9 @@ mod tests { .expect("overflow should resume after repository recovery") .expect("overflow task should complete"); tokio::time::timeout(Duration::from_secs(2), async { - while runtime.metrics.pending_current.load(Ordering::Acquire) != 0 { + while runtime.metrics.pending_current.load(Ordering::Acquire) != 0 + || runtime.priority_admission.available_permits() != 2 + { tokio::task::yield_now().await; } }) diff --git a/apps/aether-gateway/src/state/catalog.rs b/apps/aether-gateway/src/state/catalog.rs index 3aeed16c9..5b3ee60ec 100644 --- a/apps/aether-gateway/src/state/catalog.rs +++ b/apps/aether-gateway/src/state/catalog.rs @@ -480,6 +480,16 @@ impl AppState { .map_err(|err| GatewayError::Internal(err.to_string())) } + pub(crate) async fn list_provider_catalog_keys_by_ids_strong( + &self, + key_ids: &[String], + ) -> Result, GatewayError> { + self.data + .list_provider_catalog_keys_by_ids_strong(key_ids) + .await + .map_err(|err| GatewayError::Internal(err.to_string())) + } + pub(crate) async fn list_provider_catalog_key_page( &self, query: &provider_catalog::ProviderCatalogKeyListQuery, diff --git a/apps/aether-gateway/src/state/integrations.rs b/apps/aether-gateway/src/state/integrations.rs index a2b06a081..36ce42648 100644 --- a/apps/aether-gateway/src/state/integrations.rs +++ b/apps/aether-gateway/src/state/integrations.rs @@ -25,7 +25,7 @@ use tracing::{debug, warn}; use super::{AppState, GatewayError}; use crate::clock::current_unix_secs; -use crate::model_fetch::ModelFetchRuntimeState; +use crate::model_fetch::{CodexCatalogRuntime, ModelFetchRuntimeState}; use crate::provider_transport::{GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth}; use crate::request_candidate_runtime::{ RequestCandidateRuntimeCapabilityReader, RequestCandidateRuntimeReader, @@ -34,6 +34,8 @@ use crate::request_candidate_runtime::{ use crate::scheduler::state::SchedulerRuntimeState; use crate::{execution_runtime, provider_transport}; +const MODEL_FETCH_RESPONSE_BODY_LIMIT_BYTES: usize = 8 * 1024 * 1024; + impl AppState { pub(crate) async fn hydrate_antigravity_project_metadata_for_transport( &self, @@ -274,12 +276,90 @@ impl ModelFetchTransportRuntime for AppState { &self, plan: &ExecutionPlan, ) -> Result { - execution_runtime::execute_execution_runtime_sync_plan(self, None, plan) + let bounded_plan = execution_runtime::transport::with_upstream_response_body_limit( + plan, + MODEL_FETCH_RESPONSE_BODY_LIMIT_BYTES, + ); + execution_runtime::execute_execution_runtime_sync_plan(self, None, &bounded_plan) .await .map_err(GatewayError::into_message) } } +#[async_trait] +impl CodexCatalogRuntime for AppState { + fn codex_catalog_runtime_state(&self) -> &aether_runtime_state::RuntimeState { + self.runtime_state.as_ref() + } + + async fn read_codex_catalog_transport_snapshot( + &self, + provider_id: &str, + endpoint_id: &str, + key_id: &str, + ) -> Result, String> { + self.read_provider_transport_snapshot(provider_id, endpoint_id, key_id) + .await + .map_err(GatewayError::into_message) + } + + async fn read_codex_catalog_credential_scope_strong( + &self, + provider_id: &str, + key_id: &str, + ) -> Result, String> { + let Some(key) = self + .list_provider_catalog_keys_by_ids_strong(&[key_id.to_string()]) + .await + .map_err(GatewayError::into_message)? + .into_iter() + .find(|key| key.id == key_id && key.provider_id == provider_id && key.is_active) + else { + return Ok(None); + }; + + if let Some(scope) = + crate::model_fetch::codex_catalog_credential_scope_from_stored_key(&key, None, None) + { + return Ok(Some(scope)); + } + + let decrypted_auth_config = match key.encrypted_auth_config.as_deref() { + Some(ciphertext) => Some( + crate::handlers::shared::decrypt_catalog_secret_with_fallbacks( + self.encryption_key(), + ciphertext, + ) + .ok_or_else(|| { + "Codex catalog auth config could not be verified for credential fencing" + .to_string() + })?, + ), + None => None, + }; + let decrypted_api_key = match key.encrypted_api_key.as_deref() { + Some(ciphertext) => Some( + crate::handlers::shared::decrypt_catalog_secret_with_fallbacks( + self.encryption_key(), + ciphertext, + ) + .ok_or_else(|| { + "Codex catalog API key could not be verified for credential fencing".to_string() + })?, + ), + None => None, + }; + + Ok( + crate::model_fetch::codex_catalog_credential_scope_from_stored_key( + &key, + decrypted_auth_config.as_deref(), + decrypted_api_key.as_deref(), + ), + ) + } +} + #[async_trait] impl ModelFetchRuntimeState for AppState { fn has_provider_catalog_data_reader(&self) -> bool { @@ -320,6 +400,29 @@ impl ModelFetchRuntimeState for AppState { execution_runtime::execute_execution_runtime_sync_plan(self, None, plan).await } + async fn read_recent_codex_catalog_client_version( + &self, + provider_id: &str, + key_id: &str, + ) -> Option { + let credential_scope = + ::read_codex_catalog_credential_scope_strong( + self, + provider_id, + key_id, + ) + .await + .ok() + .flatten()?; + crate::model_fetch::read_recent_codex_catalog_client_version( + self.runtime_state.as_ref(), + provider_id, + key_id, + &credential_scope, + ) + .await + } + async fn update_provider_catalog_key_model_fetch_state( &self, key_id: &str, diff --git a/apps/aether-gateway/src/tests/control/admin/oauth.rs b/apps/aether-gateway/src/tests/control/admin/oauth.rs index 804d868eb..8f330ca32 100644 --- a/apps/aether-gateway/src/tests/control/admin/oauth.rs +++ b/apps/aether-gateway/src/tests/control/admin/oauth.rs @@ -8266,24 +8266,21 @@ async fn gateway_manual_oauth_refresh_prefers_fresher_transport_auth_config_over "Bearer cached-codex-access-token" ); - let mut updated_key = provider_catalog_repository - .list_keys_by_ids(&["key-codex-oauth-stale-cache".to_string()]) - .await - .expect("keys should list") - .into_iter() - .next() - .expect("key should exist"); - updated_key.encrypted_auth_config = Some( - encrypt_python_fernet_plaintext( - DEVELOPMENT_ENCRYPTION_KEY, - r#"{"provider_type":"codex","refresh_token":"fresh-codex-refresh-token","email":"alice@example.com","account_id":"acct-codex-123","plan_type":"plus","expires_at":1,"updated_at":4102444810}"#, + let fresh_auth_config = encrypt_python_fernet_plaintext( + DEVELOPMENT_ENCRYPTION_KEY, + r#"{"provider_type":"codex","refresh_token":"fresh-codex-refresh-token","email":"alice@example.com","account_id":"acct-codex-123","plan_type":"plus","expires_at":1,"updated_at":4102444810}"#, + ) + .expect("updated auth config ciphertext should build"); + assert!(provider_catalog_repository + .update_key_oauth_runtime_state( + "key-codex-oauth-stale-cache", + None, + None, + Some(&fresh_auth_config), + Some(4_102_444_810), ) - .expect("updated auth config ciphertext should build"), - ); - provider_catalog_repository - .update_key(&updated_key) .await - .expect("key should update"); + .expect("OAuth runtime state should update")); let gateway = build_router_with_state(app_state); let (gateway_url, gateway_handle) = start_server(gateway).await; diff --git a/apps/aether-gateway/src/tests/frontdoor/ai.rs b/apps/aether-gateway/src/tests/frontdoor/ai.rs index 5f19d6589..72d88382b 100644 --- a/apps/aether-gateway/src/tests/frontdoor/ai.rs +++ b/apps/aether-gateway/src/tests/frontdoor/ai.rs @@ -7,11 +7,16 @@ use super::{ use crate::image_capabilities::openai_image_gateway_max_generation_count; use crate::tests::{ any, build_router_with_state, build_state_with_execution_runtime_override, json, start_server, - to_bytes, AppState, Arc, Body, Json, Mutex, Request, Router, StatusCode, EXECUTION_PATH_HEADER, - EXECUTION_PATH_LOCAL_AI_PUBLIC, EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS, + to_bytes, wait_until, AppState, Arc, Body, Json, Mutex, Request, Router, StatusCode, + EXECUTION_PATH_HEADER, EXECUTION_PATH_LOCAL_AI_PUBLIC, + EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS, }; +use aether_contracts::{ExecutionResult, ExecutionTelemetry, ResponseBody}; +use aether_crypto::encrypt_python_fernet_plaintext; use aether_data::repository::global_models::InMemoryGlobalModelReadRepository; +use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository; use aether_data::DataLayerError; +use aether_data_contracts::repository::auth::AuthApiKeyWriteRepository; use aether_data_contracts::repository::candidate_selection::{ MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow, StoredPoolKeyCandidateRowsByKeyIdsQuery, StoredPoolKeyCandidateRowsQuery, @@ -20,13 +25,21 @@ use aether_data_contracts::repository::candidate_selection::{ use aether_data_contracts::repository::global_models::{ StoredAdminGlobalModel, UpdateAdminGlobalModelRecord, }; +use aether_data_contracts::repository::provider_catalog::{ + ProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey, + StoredProviderCatalogProvider, +}; use async_trait::async_trait; use axum::response::IntoResponse; use std::collections::HashMap; use std::future::pending; -use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; -fn codex_models_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot { +fn codex_models_snapshot( + api_key_id: &str, + user_id: &str, + allowed_models: &[&str], +) -> StoredAuthApiKeySnapshot { StoredAuthApiKeySnapshot::new( user_id.to_string(), "alice".to_string(), @@ -37,7 +50,7 @@ fn codex_models_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySna false, Some(json!(["codex"])), Some(json!(["openai:responses"])), - Some(json!(["frontier-sol", "broken-luna"])), + Some(json!(allowed_models)), api_key_id.to_string(), Some("codex-models".to_string()), true, @@ -48,7 +61,7 @@ fn codex_models_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySna Some(4_102_444_800), Some(json!(["codex"])), Some(json!(["openai:responses"])), - Some(json!(["frontier-sol", "broken-luna"])), + Some(json!(allowed_models)), ) .expect("Codex models auth snapshot should build") } @@ -104,6 +117,7 @@ fn complete_codex_model_card(source_model_name: &str) -> serde_json::Value { "upgrade": null, "base_instructions": "Use the current Codex instructions.", "model_messages": null, + "available_in_plans": ["plus", "pro"], "support_verbosity": true, "default_verbosity": "low", "apply_patch_tool_type": "freeform", @@ -115,6 +129,105 @@ fn complete_codex_model_card(source_model_name: &str) -> serde_json::Value { }) } +fn codex_catalog_provider(provider_id: &str) -> StoredProviderCatalogProvider { + StoredProviderCatalogProvider::new( + provider_id.to_string(), + "codex".to_string(), + Some("https://chatgpt.com".to_string()), + "codex".to_string(), + ) + .expect("Codex provider should build") +} + +fn codex_catalog_endpoint(provider_id: &str, endpoint_id: &str) -> StoredProviderCatalogEndpoint { + StoredProviderCatalogEndpoint::new( + endpoint_id.to_string(), + provider_id.to_string(), + "openai:responses".to_string(), + Some("openai".to_string()), + Some("responses".to_string()), + true, + ) + .expect("Codex endpoint should build") + .with_transport_fields( + "https://chatgpt.example/backend-api/codex".to_string(), + None, + None, + None, + None, + None, + None, + None, + ) + .expect("Codex endpoint transport should build") +} + +fn codex_catalog_key( + provider_id: &str, + key_id: &str, + allowed_models: &[&str], +) -> StoredProviderCatalogKey { + let mut key = StoredProviderCatalogKey::new( + key_id.to_string(), + provider_id.to_string(), + "manual".to_string(), + "bearer".to_string(), + None, + true, + ) + .expect("Codex key should build") + .with_transport_fields( + Some(json!(["openai:responses"])), + encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "oauth-upstream-secret") + .expect("Codex test token should encrypt"), + None, + None, + None, + Some(json!(allowed_models)), + None, + None, + None, + ) + .expect("Codex key transport should build"); + key.auto_fetch_models = false; + key.locked_models = Some(json!(["manual-locked-model"])); + key.model_include_patterns = Some(json!(["gpt-future-*"])); + key.model_exclude_patterns = Some(json!(["gpt-future-denied"])); + key +} + +fn codex_catalog_execution_result( + plan: &aether_contracts::ExecutionPlan, + status_code: u16, + body: serde_json::Value, + etag: Option<&str>, +) -> ExecutionResult { + let mut headers = std::collections::BTreeMap::from([( + "content-type".to_string(), + "application/json".to_string(), + )]); + if let Some(etag) = etag { + headers.insert("ETag".to_string(), etag.to_string()); + } + ExecutionResult { + request_id: plan.request_id.clone(), + candidate_id: plan.candidate_id.clone(), + status_code, + headers, + response_observation: None, + body: Some(ResponseBody { + json_body: Some(body), + body_bytes_b64: None, + }), + telemetry: Some(ExecutionTelemetry { + ttfb_ms: Some(1), + elapsed_ms: Some(2), + upstream_bytes: None, + }), + error: None, + } +} + fn gemini_operation_status_label(status: VideoTaskStatus) -> &'static str { match status { VideoTaskStatus::Pending => "Pending", @@ -449,102 +562,618 @@ async fn gateway_handles_public_openai_models_without_hitting_fallback_probe() { fallback_probe_handle.abort(); } -#[tokio::test] -async fn gateway_serves_codex_model_cards_for_versioned_models_requests() { - let codex_row = - sample_codex_models_candidate_row("provider-codex-models", "frontier-sol", "gpt-5.6-sol"); - let incomplete_codex_row = sample_codex_models_candidate_row( - "provider-codex-incomplete", - "broken-luna", - "gpt-5.6-luna", - ); - let candidate_repository = - Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![ - codex_row.clone(), - incomplete_codex_row.clone(), - sample_models_candidate_row( - "provider-openai-responses", - "openai", - "openai:responses", - "custom-responses-model", - 20, +#[test] +fn gateway_versioned_models_fail_closed_when_cached_auth_becomes_unusable_or_missing() { + std::thread::Builder::new() + .name("codex-model-catalog-auth-race".to_string()) + .stack_size(16 * 1024 * 1024) + .spawn(|| { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("Codex auth race test runtime should build") + .block_on(run_versioned_models_auth_race_scenario()); + }) + .expect("Codex auth race test thread should spawn") + .join() + .expect("Codex auth race test thread should finish"); +} + +async fn run_versioned_models_auth_race_scenario() { + let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( + Some(hash_api_key("sk-codex-models-auth-race")), + codex_models_snapshot( + "key-codex-models-auth-race", + "user-codex-models-auth-race", + &["future-alias"], + ), + )])); + let candidate_repository = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed( + Vec::new(), + )); + let gateway = build_router_with_state( + AppState::new() + .expect("gateway should build") + .with_data_state_for_tests( + crate::data::GatewayDataState::with_minimal_candidate_selection_and_auth_for_tests( + candidate_repository, + auth_repository.clone(), + ), ), - ])); + ); + let (gateway_url, gateway_handle) = start_server(gateway).await; + let client = reqwest::Client::new(); + let models_url = format!("{gateway_url}/v1/models?client_version=0.145.2"); + + let warm_response = client + .get(&models_url) + .header("authorization", "Bearer sk-codex-models-auth-race") + .send() + .await + .expect("initial versioned models request should succeed"); + assert_eq!(warm_response.status(), StatusCode::OK); + + assert!(auth_repository + .set_user_api_key_locked( + "user-codex-models-auth-race", + "key-codex-models-auth-race", + true, + ) + .await + .expect("locking the API key should succeed")); + let locked_response = client + .get(&models_url) + .header("authorization", "Bearer sk-codex-models-auth-race") + .send() + .await + .expect("locked versioned models request should complete"); + assert_eq!(locked_response.status(), StatusCode::UNAUTHORIZED); + + assert!(auth_repository + .delete_user_api_key("user-codex-models-auth-race", "key-codex-models-auth-race",) + .await + .expect("deleting the API key should succeed")); + let missing_response = client + .get(&models_url) + .header("authorization", "Bearer sk-codex-models-auth-race") + .send() + .await + .expect("missing versioned models request should complete"); + assert_eq!(missing_response.status(), StatusCode::UNAUTHORIZED); + + gateway_handle.abort(); +} + +#[test] +fn gateway_serves_codex_model_cards_for_versioned_models_requests() { + std::thread::Builder::new() + .name("codex-model-catalog-frontdoor".to_string()) + .stack_size(16 * 1024 * 1024) + .spawn(|| { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("Codex frontdoor test runtime should build") + .block_on(run_versioned_codex_model_cards_frontdoor_scenario()); + }) + .expect("Codex frontdoor test thread should spawn") + .join() + .expect("Codex frontdoor test thread should finish"); +} + +async fn run_versioned_codex_model_cards_frontdoor_scenario() { + const PROVIDER_ID: &str = "provider-codex-models"; + const CATALOG_KEY_ID: &str = "key-provider-codex-models"; + const CATALOG_ENDPOINT_ID: &str = "endpoint-provider-codex-models"; + const SOURCE_MODELS: &[&str] = &[ + "gpt-future-dynamic", + "gpt-future-legacy", + "gpt-future-second", + "gpt-hidden-direct", + ]; + const GLOBAL_MODELS: &[&str] = &["future-alias", "legacy-alias"]; + + let mut codex_rows = vec![ + sample_codex_models_candidate_row(PROVIDER_ID, "future-alias", "gpt-future-dynamic"), + sample_codex_models_candidate_row(PROVIDER_ID, "legacy-alias", "gpt-future-legacy"), + sample_codex_models_candidate_row(PROVIDER_ID, "second-alias", "gpt-future-second"), + sample_codex_models_candidate_row(PROVIDER_ID, "hidden-alias", "gpt-hidden-direct"), + ]; + for row in &mut codex_rows { + row.key_allowed_models = Some( + SOURCE_MODELS + .iter() + .map(|value| value.to_string()) + .collect(), + ); + } + let mut all_candidate_rows = codex_rows.clone(); + all_candidate_rows.push(sample_models_candidate_row( + "provider-openai-responses", + "openai", + "openai:responses", + "custom-responses-model", + 20, + )); + let candidate_repository = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed( + all_candidate_rows, + )); let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![ ( Some(hash_api_key("sk-codex-models")), - codex_models_snapshot("key-codex-models", "user-codex-models"), + codex_models_snapshot("key-codex-models", "user-codex-models", GLOBAL_MODELS), ), ( Some(hash_api_key("sk-standard-models")), unrestricted_models_snapshot("key-standard-models", "user-standard-models"), ), + ( + Some(hash_api_key("sk-codex-legacy-only")), + codex_models_snapshot( + "key-codex-legacy-only", + "user-codex-legacy-only", + &["legacy-alias"], + ), + ), + ( + Some(hash_api_key("sk-codex-hidden-mixed")), + codex_models_snapshot( + "key-codex-hidden-mixed", + "user-codex-hidden-mixed", + &["future-alias", "hidden-alias"], + ), + ), + ( + Some(hash_api_key("sk-codex-second-mixed")), + codex_models_snapshot( + "key-codex-second-mixed", + "user-codex-second-mixed", + &["future-alias", "second-alias"], + ), + ), ])); - let state = AppState::new() - .expect("gateway should build") + let original_catalog_key = codex_catalog_key(PROVIDER_ID, CATALOG_KEY_ID, SOURCE_MODELS); + assert!(!original_catalog_key.auto_fetch_models); + let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( + vec![codex_catalog_provider(PROVIDER_ID)], + vec![codex_catalog_endpoint(PROVIDER_ID, CATALOG_ENDPOINT_ID)], + vec![original_catalog_key.clone()], + )); + let rows_before = candidate_repository + .list_for_exact_api_format("openai:responses") + .await + .expect("candidate rows should load before request"); + + let catalog_generation = Arc::new(AtomicUsize::new(0)); + let catalog_hits = Arc::new(AtomicUsize::new(0)); + let captured_plans = Arc::new(Mutex::new(Vec::<(String, Option)>::new())); + let generation_for_runtime = Arc::clone(&catalog_generation); + let hits_for_runtime = Arc::clone(&catalog_hits); + let plans_for_runtime = Arc::clone(&captured_plans); + let execution_runtime = Router::new().route( + "/v1/execute/sync", + any(move |request: Request| { + let generation_for_request = Arc::clone(&generation_for_runtime); + let hits_for_request = Arc::clone(&hits_for_runtime); + let plans_for_request = Arc::clone(&plans_for_runtime); + async move { + let (_parts, body) = request.into_parts(); + let raw_body = to_bytes(body, usize::MAX) + .await + .expect("execution runtime request body should read"); + let plan: aether_contracts::ExecutionPlan = + serde_json::from_slice(&raw_body).expect("execution runtime plan should parse"); + plans_for_request + .lock() + .expect("plans mutex") + .push((plan.url.clone(), plan.headers.get("user-agent").cloned())); + if plan.url.contains("/models?") { + hits_for_request.fetch_add(1, Ordering::SeqCst); + let generation = generation_for_request.load(Ordering::SeqCst); + if generation == 1 { + return Json(codex_catalog_execution_result( + &plan, + 503, + json!({"error":{"message":"temporary catalog outage"}}), + None, + )); + } + + let mut current = complete_codex_model_card("gpt-future-dynamic"); + let current_object = current.as_object_mut().expect("current card object"); + current_object.remove("base_instructions"); + current_object.insert( + "model_messages".to_string(), + json!({"instructions_template":"Use future dynamic instructions."}), + ); + current_object.insert( + "future_capability".to_string(), + json!({"mode":"opaque-current"}), + ); + + let mut legacy = complete_codex_model_card("gpt-future-legacy"); + legacy["base_instructions"] = json!("Use legacy future instructions."); + legacy["future_capability"] = json!({"mode":"opaque-legacy"}); + + let mut models = vec![current, legacy]; + if generation >= 2 { + let mut second = complete_codex_model_card("gpt-future-second"); + second + .as_object_mut() + .expect("second card object") + .remove("base_instructions"); + second["model_messages"] = + json!({"instructions_template":"Use second future instructions."}); + second["future_capability"] = json!({"mode":"added-without-code-change"}); + models.push(second); + models.push(complete_codex_model_card("gpt-future-unmapped")); + } + let etag = if generation >= 2 { + "\"catalog-etag-v2\"" + } else { + "\"catalog-etag-v1\"" + }; + return Json(codex_catalog_execution_result( + &plan, + 200, + json!({"models": models, "future_top_level": true}), + Some(etag), + )); + } + + Json(codex_catalog_execution_result( + &plan, + 200, + json!({ + "id": "resp-future-dynamic", + "object": "response", + "model": "gpt-future-dynamic", + "output": [], + "usage": { + "input_tokens": 1, + "output_tokens": 2, + "total_tokens": 3 + } + }), + Some("\"catalog-etag-v2\""), + )) + } + }), + ); + let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await; + let state = build_state_with_execution_runtime_override(execution_runtime_url) .with_data_state_for_tests( crate::data::GatewayDataState::with_minimal_candidate_selection_and_auth_for_tests( - candidate_repository, + candidate_repository.clone(), auth_repository, - ), + ) + .attach_provider_catalog_repository_for_tests(provider_catalog_repository.clone()) + .with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY), ); - state - .runtime_kv_setex( - &format!( - "upstream_models:{}:{}", - codex_row.provider_id, codex_row.key_id - ), - &serde_json::to_string(&vec![complete_codex_model_card("gpt-5.6-sol")]) - .expect("model cache should serialize"), - 60, + + let configured_rows = state + .list_minimal_candidate_selection_rows_for_api_format("openai:responses") + .await + .expect("configured Codex candidate rows should be readable"); + assert_eq!( + configured_rows + .iter() + .filter(|row| row.provider_type == "codex") + .count(), + codex_rows.len() + ); + let resolved_auth = + aether_data_contracts::repository::auth::ResolvedAuthApiKeySnapshot::from_stored( + codex_models_snapshot("key-codex-models", "user-codex-models", GLOBAL_MODELS), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs(), + ); + let eligible_rows = crate::handlers::public::filter_eligible_model_rows( + configured_rows.clone(), + Some(&resolved_auth), + "openai:responses", + ); + assert_eq!( + eligible_rows + .iter() + .filter(|row| row.provider_type == "codex") + .count(), + GLOBAL_MODELS.len(), + "Codex fixture rows must survive the same provider/model/key authorization filters as the route" + ); + let actual_auth = state + .data + .read_auth_api_key_snapshot( + "user-codex-models", + "key-codex-models", + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs(), ) .await - .expect("model cache should seed"); - state - .runtime_kv_setex( - &format!( - "upstream_models:{}:{}", - incomplete_codex_row.provider_id, incomplete_codex_row.key_id - ), - &serde_json::to_string(&vec![json!({ - "id": "gpt-5.6-luna", - "slug": "gpt-5.6-luna", - "display_name": "GPT-5.6-Luna" - })]) - .expect("incomplete model cache should serialize"), - 60, + .expect("Codex auth snapshot read should succeed") + .expect("Codex auth snapshot should exist"); + assert_eq!( + crate::handlers::public::filter_eligible_model_rows( + configured_rows.clone(), + Some(&actual_auth), + "openai:responses", + ) + .iter() + .filter(|row| row.provider_type == "codex") + .count(), + GLOBAL_MODELS.len(), + "stored Codex auth snapshot must preserve every authorized manual mapping" + ); + assert!( + ::read_codex_catalog_transport_snapshot( + &state, + PROVIDER_ID, + CATALOG_ENDPOINT_ID, + CATALOG_KEY_ID, ) .await - .expect("incomplete model cache should seed"); + .expect("Codex catalog transport lookup should succeed") + .is_some(), + "Codex catalog transport must be available even when auto_fetch_models is disabled" + ); let gateway = build_router_with_state(state); let (gateway_url, gateway_handle) = start_server(gateway).await; let client = reqwest::Client::new(); let codex_response = client - .get(format!("{gateway_url}/v1/models?client_version=0.144.1")) + .get(format!( + "{gateway_url}/v1/models?client_version=0.145.2-beta.7%2Bdesktop.9" + )) .header("authorization", "Bearer sk-codex-models") .send() .await .expect("Codex models request should succeed"); assert_eq!(codex_response.status(), StatusCode::OK); + let codex_etag = codex_response + .headers() + .get(http::header::ETAG) + .and_then(|value| value.to_str().ok()) + .map(ToOwned::to_owned); let codex_payload: serde_json::Value = codex_response .json() .await .expect("Codex models body should parse"); - assert_eq!(codex_payload["models"].as_array().map(Vec::len), Some(1)); - assert_eq!(codex_payload["models"][0]["slug"], "frontier-sol"); assert_eq!( - codex_payload["models"][0]["supported_reasoning_levels"][5]["effort"], - "ultra" + catalog_hits.load(Ordering::SeqCst), + 1, + "cold Codex request must reach the upstream catalog once; payload={codex_payload}" + ); + assert_eq!(codex_etag.as_deref(), Some("\"catalog-etag-v1\"")); + assert_eq!(codex_payload["models"].as_array().map(Vec::len), Some(2)); + let current_card = codex_payload["models"] + .as_array() + .and_then(|models| models.iter().find(|model| model["slug"] == "future-alias")) + .expect("current model card should be projected"); + assert_eq!( + current_card["model_messages"]["instructions_template"], + "Use future dynamic instructions." ); assert_eq!( - codex_payload["models"][0]["future_capability"], - json!({"enabled": true}) + current_card["future_capability"], + json!({"mode":"opaque-current"}) ); - assert!(codex_payload["models"][0].get("id").is_none()); - assert!(codex_payload["models"][0].get("api_formats").is_none()); + assert_eq!(current_card["available_in_plans"], json!(["plus", "pro"])); + assert!(current_card.get("base_instructions").is_none()); + assert!(current_card.get("id").is_none()); + assert!(current_card.get("api_formats").is_none()); + let legacy_card = codex_payload["models"] + .as_array() + .and_then(|models| models.iter().find(|model| model["slug"] == "legacy-alias")) + .expect("legacy model card should be projected"); + assert_eq!( + legacy_card["base_instructions"], + "Use legacy future instructions." + ); + assert_eq!( + legacy_card["future_capability"], + json!({"mode":"opaque-legacy"}) + ); + assert!(codex_payload["models"] + .as_array() + .is_some_and( + |models| models.iter().all(|model| model["slug"] != "hidden-alias" + && model["slug"] != "second-alias" + && model["slug"] != "gpt-future-unmapped") + )); assert!(codex_payload.get("object").is_none()); + assert_eq!(catalog_hits.load(Ordering::SeqCst), 1); + + let restricted_response = client + .get(format!("{gateway_url}/v1/models?client_version=0.145.2")) + .header("authorization", "Bearer sk-codex-legacy-only") + .send() + .await + .expect("restricted Codex models request should succeed"); + assert_eq!(restricted_response.status(), StatusCode::OK); + let restricted_payload: serde_json::Value = restricted_response + .json() + .await + .expect("restricted Codex models body should parse"); + assert_eq!( + restricted_payload["models"].as_array().map(|models| models + .iter() + .map(|model| model["slug"].as_str().unwrap_or_default()) + .collect::>()), + Some(vec!["legacy-alias"]) + ); + assert_eq!(catalog_hits.load(Ordering::SeqCst), 1); + + let incomplete_authorized_response = client + .get(format!("{gateway_url}/v1/models?client_version=0.145.2")) + .header("authorization", "Bearer sk-codex-hidden-mixed") + .send() + .await + .expect("incomplete authorized Codex catalog request should succeed"); + assert_eq!(incomplete_authorized_response.status(), StatusCode::OK); + assert!(incomplete_authorized_response + .headers() + .get(http::header::ETAG) + .is_none()); + let incomplete_authorized_payload: serde_json::Value = incomplete_authorized_response + .json() + .await + .expect("incomplete authorized Codex body should parse"); + assert_eq!( + incomplete_authorized_payload["models"] + .as_array() + .map(Vec::len), + Some(0), + "a partial non-empty remote catalog would hide the client's bundled fallback models" + ); + assert_eq!(catalog_hits.load(Ordering::SeqCst), 1); + + let pending_second_response = client + .get(format!("{gateway_url}/v1/models?client_version=0.145.2")) + .header("authorization", "Bearer sk-codex-second-mixed") + .send() + .await + .expect("not-yet-published authorized model request should succeed"); + assert_eq!(pending_second_response.status(), StatusCode::OK); + let pending_second_payload: serde_json::Value = pending_second_response + .json() + .await + .expect("not-yet-published authorized model body should parse"); + assert_eq!( + pending_second_payload["models"].as_array().map(Vec::len), + Some(0) + ); + assert_eq!(catalog_hits.load(Ordering::SeqCst), 1); + + let captured_catalog_plan = captured_plans + .lock() + .expect("plans mutex") + .iter() + .find(|(url, _)| url.contains("/models?")) + .cloned() + .expect("catalog execution plan should be captured"); + assert_eq!( + captured_catalog_plan.0, + "https://chatgpt.example/backend-api/codex/models?client_version=0.145.2" + ); + assert_eq!( + captured_catalog_plan.1.as_deref(), + Some("codex_cli_rs/0.145.2") + ); + + let fresh_response = client + .get(format!("{gateway_url}/v1/models?client_version=0.145.2")) + .header("authorization", "Bearer sk-codex-models") + .send() + .await + .expect("fresh Codex models request should succeed"); + assert_eq!(fresh_response.status(), StatusCode::OK); + assert_eq!(catalog_hits.load(Ordering::SeqCst), 1); + + tokio::time::sleep(std::time::Duration::from_millis(1_100)).await; + catalog_generation.store(1, Ordering::SeqCst); + let stale_started = std::time::Instant::now(); + let stale_response = client + .get(format!("{gateway_url}/v1/models?client_version=0.145.2")) + .header("authorization", "Bearer sk-codex-models") + .send() + .await + .expect("stale Codex models request should succeed"); + assert_eq!(stale_response.status(), StatusCode::OK); + assert!(stale_started.elapsed() < std::time::Duration::from_millis(400)); + let stale_payload: serde_json::Value = stale_response + .json() + .await + .expect("stale body should parse"); + assert!(stale_payload["models"] + .as_array() + .is_some_and(|models| models.iter().any(|model| model["slug"] == "future-alias"))); + wait_until(1_000, || catalog_hits.load(Ordering::SeqCst) >= 2).await; + + let failed_refresh_lkg_response = client + .get(format!("{gateway_url}/v1/models?client_version=0.145.2")) + .header("authorization", "Bearer sk-codex-models") + .send() + .await + .expect("failed refresh should keep serving LKG"); + let failed_refresh_lkg_payload: serde_json::Value = failed_refresh_lkg_response + .json() + .await + .expect("failed refresh LKG body should parse"); + assert!(failed_refresh_lkg_payload["models"] + .as_array() + .is_some_and(|models| models.iter().any(|model| model["slug"] == "future-alias"))); + assert_eq!(catalog_hits.load(Ordering::SeqCst), 2); + + tokio::time::sleep(std::time::Duration::from_millis(150)).await; + catalog_generation.store(2, Ordering::SeqCst); + let refresh_trigger = client + .get(format!("{gateway_url}/v1/models?client_version=0.145.2")) + .header("authorization", "Bearer sk-codex-models") + .send() + .await + .expect("recovered refresh trigger should succeed"); + assert_eq!(refresh_trigger.status(), StatusCode::OK); + wait_until(1_000, || catalog_hits.load(Ordering::SeqCst) >= 3).await; + + let updated_response = client + .get(format!("{gateway_url}/v1/models?client_version=0.145.2")) + .header("authorization", "Bearer sk-codex-second-mixed") + .send() + .await + .expect("updated catalog request should succeed"); + assert_eq!( + updated_response + .headers() + .get(http::header::ETAG) + .and_then(|value| value.to_str().ok()), + Some("\"catalog-etag-v2\"") + ); + let updated_payload: serde_json::Value = updated_response + .json() + .await + .expect("updated body should parse"); + assert!(updated_payload["models"] + .as_array() + .is_some_and(|models| models.iter().any(|model| { + model["slug"] == "second-alias" + && model["future_capability"] == json!({"mode":"added-without-code-change"}) + }))); + assert!(updated_payload["models"] + .as_array() + .is_some_and(|models| models.iter().all(|model| { + model["slug"] != "hidden-alias" && model["slug"] != "gpt-future-unmapped" + }))); + + let stored_keys = provider_catalog_repository + .list_keys_by_ids(&[CATALOG_KEY_ID.to_string()]) + .await + .expect("catalog key should remain readable"); + assert_eq!(stored_keys.as_slice(), &[original_catalog_key]); + let rows_after = candidate_repository + .list_for_exact_api_format("openai:responses") + .await + .expect("candidate rows should load after request"); + assert_eq!(rows_after, rows_before); + + let inference_response = client + .post(format!("{gateway_url}/v1/responses")) + .header(http::header::CONTENT_TYPE, "application/json") + .header("authorization", "Bearer sk-codex-models") + .body(r#"{"model":"future-alias","input":"hello","store":false}"#) + .send() + .await + .expect("manual alias inference should succeed"); + assert_eq!(inference_response.status(), StatusCode::OK); + let inference_payload: serde_json::Value = inference_response + .json() + .await + .expect("inference body should parse"); + assert_eq!(inference_payload["model"], "gpt-future-dynamic"); let standard_response = client .get(format!("{gateway_url}/v1/models")) @@ -562,6 +1191,7 @@ async fn gateway_serves_codex_model_cards_for_versioned_models_requests() { assert!(standard_payload.get("models").is_none()); gateway_handle.abort(); + execution_runtime_handle.abort(); } #[tokio::test] @@ -695,7 +1325,7 @@ async fn gateway_returns_empty_openai_models_when_candidate_rows_stall() { let (gateway_url, gateway_handle) = start_server(gateway).await; let response = reqwest::Client::builder() - .timeout(std::time::Duration::from_millis(500)) + .timeout(std::time::Duration::from_secs(3)) .build() .expect("client should build") .get(format!("{gateway_url}/v1/models")) @@ -755,7 +1385,7 @@ async fn gateway_returns_not_found_for_openai_model_detail_when_candidate_rows_s let (gateway_url, gateway_handle) = start_server(gateway).await; let response = reqwest::Client::builder() - .timeout(std::time::Duration::from_millis(500)) + .timeout(std::time::Duration::from_secs(3)) .build() .expect("client should build") .get(format!("{gateway_url}/v1/models/gpt-stalled")) diff --git a/crates/aether-ai/formats/src/api.rs b/crates/aether-ai/formats/src/api.rs index 3ff098b1d..d969f1071 100644 --- a/crates/aether-ai/formats/src/api.rs +++ b/crates/aether-ai/formats/src/api.rs @@ -188,9 +188,9 @@ pub use crate::formats::{ apply_codex_openai_special_headers, apply_openai_responses_compact_special_body_edits, build_codex_model_catalog_metadata, parse_codex_auth_identity, - resolve_codex_responses_model_capabilities, CodexAuthIdentity, - CodexResponsesModelCapabilities, CODEX_OPENAI_IMAGE_DEFAULT_MODEL, - CODEX_OPENAI_IMAGE_DEFAULT_OUTPUT_FORMAT, + project_codex_catalog_model_card, resolve_codex_responses_model_capabilities, + CodexAuthIdentity, CodexResponsesModelCapabilities, + CODEX_OPENAI_IMAGE_DEFAULT_MODEL, CODEX_OPENAI_IMAGE_DEFAULT_OUTPUT_FORMAT, CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_MODEL, CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_PROMPT, CODEX_OPENAI_IMAGE_INTERNAL_MODEL, }, diff --git a/crates/aether-ai/formats/src/formats/openai/responses/codex.rs b/crates/aether-ai/formats/src/formats/openai/responses/codex.rs index a55ed0d7f..745840528 100644 --- a/crates/aether-ai/formats/src/formats/openai/responses/codex.rs +++ b/crates/aether-ai/formats/src/formats/openai/responses/codex.rs @@ -236,6 +236,40 @@ pub fn build_codex_model_catalog_metadata(cards: &[Value]) -> Value { }) } +/// Projects an upstream Codex catalog card onto an authorized Aether model name. +/// +/// Catalog cards are versioned by Codex and may gain fields that Aether does not know about. Keep +/// the matching card opaque: only the Aether-internal identity fields are removed and `slug` is +/// replaced with the downstream model name. A non-empty, exact `id` or `slug` match is the only +/// schema requirement. +pub fn project_codex_catalog_model_card( + catalog_models: &[Value], + source_model: &str, + global_model: &str, +) -> Option { + if source_model.trim().is_empty() || global_model.trim().is_empty() { + return None; + } + + let mut card = catalog_models + .iter() + .filter_map(Value::as_object) + .find(|card| { + ["id", "slug"].iter().any(|field| { + card.get(*field) + .and_then(Value::as_str) + .filter(|identity| !identity.trim().is_empty()) + .is_some_and(|identity| identity == source_model) + }) + })? + .clone(); + + card.remove("id"); + card.remove("api_formats"); + card.insert("slug".to_string(), Value::String(global_model.to_string())); + Some(Value::Object(card)) +} + fn codex_execution_model_card(card: &Value) -> Value { const EXCLUDED_FIELDS: [&str; 3] = ["base_instructions", "model_messages", "available_in_plans"]; @@ -1989,7 +2023,7 @@ mod tests { apply_codex_openai_responses_special_body_edits_with_source_model_and_capabilities, apply_codex_openai_special_headers, apply_openai_responses_compact_special_body_edits, build_codex_model_catalog_metadata, bundled_codex_model_cards, effective_codex_model_cards, - resolve_codex_responses_model_capabilities, + project_codex_catalog_model_card, resolve_codex_responses_model_capabilities, validate_codex_openai_responses_compact_request_contract, CODEX_CLIENT_ORIGINATOR, CODEX_CLIENT_USER_AGENT, CODEX_OPENAI_IMAGE_INTERNAL_MODEL, CODEX_OPENAI_RESPONSES_UNSUPPORTED_BODY_FIELDS, CODEX_RESPONSES_LITE_HEADER, @@ -2091,6 +2125,105 @@ mod tests { ); } + #[test] + fn projects_current_codex_catalog_card_as_opaque_json() { + let card = json!({ + "id": "upstream-internal-id", + "slug": "gpt-future-dynamic", + "api_formats": ["openai:responses"], + "model_messages": { + "instructions_template": "Current instructions for {{ personality }}" + }, + "available_in_plans": ["plus", "pro"], + "future_capability": { + "mode": "native", + "unknown_nested_field": [1, 2, 3] + } + }); + + let projected = + project_codex_catalog_model_card(&[card], "gpt-future-dynamic", "future-model-alias") + .expect("current Codex card should project"); + + assert_eq!(projected["slug"], "future-model-alias"); + assert!(projected.get("id").is_none()); + assert!(projected.get("api_formats").is_none()); + assert_eq!( + projected["model_messages"]["instructions_template"], + "Current instructions for {{ personality }}" + ); + assert_eq!(projected["available_in_plans"], json!(["plus", "pro"])); + assert_eq!( + projected["future_capability"], + json!({ + "mode": "native", + "unknown_nested_field": [1, 2, 3] + }) + ); + } + + #[test] + fn projects_legacy_codex_catalog_card_by_id_without_known_schema_fields() { + let card = json!({ + "id": "gpt-future-dynamic", + "base_instructions": "Legacy instructions", + "available_in_plans": ["team"], + "future_capability": true + }); + + let projected = project_codex_catalog_model_card( + &[json!("not an object"), card], + "gpt-future-dynamic", + "legacy-model-alias", + ) + .expect("legacy Codex card should project by id"); + + assert_eq!( + projected, + json!({ + "slug": "legacy-model-alias", + "base_instructions": "Legacy instructions", + "available_in_plans": ["team"], + "future_capability": true + }) + ); + } + + #[test] + fn codex_catalog_projection_requires_a_non_empty_exact_identity() { + let cards = [ + json!(null), + json!({}), + json!({"id": "", "slug": " "}), + json!({"slug": "gpt-future-dynamic-preview"}), + ]; + + assert!(project_codex_catalog_model_card( + &cards, + "gpt-future-dynamic", + "future-model-alias" + ) + .is_none()); + assert!(project_codex_catalog_model_card( + &[json!({"slug": "gpt-future-dynamic"})], + "", + "future-model-alias" + ) + .is_none()); + assert!(project_codex_catalog_model_card( + &[json!({"slug": "gpt-future-dynamic"})], + " gpt-future-dynamic ", + "future-model-alias" + ) + .is_none()); + assert!(project_codex_catalog_model_card( + &[json!({"slug": "gpt-future-dynamic"})], + "gpt-future-dynamic", + " " + ) + .is_none()); + } + #[test] fn codex_responses_lite_header_is_omitted_for_server_side_compaction() { let mut headers = std::collections::BTreeMap::from([( diff --git a/crates/aether-ai/formats/src/formats/openai/responses/history.rs b/crates/aether-ai/formats/src/formats/openai/responses/history.rs index 8c6ca56d7..f229f720b 100644 --- a/crates/aether-ai/formats/src/formats/openai/responses/history.rs +++ b/crates/aether-ai/formats/src/formats/openai/responses/history.rs @@ -397,8 +397,8 @@ mod tests { use super::{ expand_previous_response_for_chat, hydrate_response_history, - record_converted_response_history, response_history_storage_key, response_history_store, - ResponseHistoryStore, + record_converted_response_history, response_history_key, response_history_storage_key, + response_history_store, }; fn conversion_report_context(original_request_body: serde_json::Value) -> serde_json::Value { @@ -496,9 +496,13 @@ mod tests { assert!(!record.storage_key.contains("distributed-history-key-a")); assert!(!record.storage_key.contains("resp_distributed_history_1")); - *response_history_store() + response_history_store() .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()) = ResponseHistoryStore::default(); + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .remove(&response_history_key( + "resp_distributed_history_1", + Some("distributed-history-key-a"), + )); let continuation = json!({ "model": "deepseek-v4-flash", "previous_response_id": "resp_distributed_history_1", diff --git a/crates/aether-ai/formats/src/lib.rs b/crates/aether-ai/formats/src/lib.rs index 452733881..b32be13d2 100644 --- a/crates/aether-ai/formats/src/lib.rs +++ b/crates/aether-ai/formats/src/lib.rs @@ -44,9 +44,10 @@ pub use formats::openai::request_contract::{ }; pub use formats::openai::responses::codex::{ build_codex_model_catalog_metadata, bundled_codex_model_cards, effective_codex_model_cards, - parse_codex_auth_identity, resolve_codex_responses_model_capabilities, CodexAuthIdentity, - CodexResponsesModelCapabilities, CODEX_CLIENT_ORIGINATOR, CODEX_CLIENT_USER_AGENT, - CODEX_CLIENT_VERSION, CODEX_MODEL_CATALOG_METADATA_FIELD, CODEX_RESPONSES_LITE_HEADER, + parse_codex_auth_identity, project_codex_catalog_model_card, + resolve_codex_responses_model_capabilities, CodexAuthIdentity, CodexResponsesModelCapabilities, + CODEX_CLIENT_ORIGINATOR, CODEX_CLIENT_USER_AGENT, CODEX_CLIENT_VERSION, + CODEX_MODEL_CATALOG_METADATA_FIELD, CODEX_RESPONSES_LITE_HEADER, }; pub use formats::openai::responses::request::{ validate_openai_responses_request_contract, OpenAiResponsesRequestContractViolation, diff --git a/crates/aether-data/contracts/src/repository/provider_catalog/types.rs b/crates/aether-data/contracts/src/repository/provider_catalog/types.rs index 7039a5878..5e1f277de 100644 --- a/crates/aether-data/contracts/src/repository/provider_catalog/types.rs +++ b/crates/aether-data/contracts/src/repository/provider_catalog/types.rs @@ -797,6 +797,18 @@ pub trait ProviderCatalogReadRepository: Send + Sync { key_ids: &[String], ) -> Result, crate::DataLayerError>; + /// Reads credential-bearing key records without an intervening repository cache. + /// + /// Callers use this only for security-sensitive generation fences where a short-lived cached + /// key record could bind data to credentials that an administrator has already replaced. + /// Repository implementations that do not add a read cache can use the default behavior. + async fn list_keys_by_ids_strong( + &self, + key_ids: &[String], + ) -> Result, crate::DataLayerError> { + self.list_keys_by_ids(key_ids).await + } + async fn list_keys_by_provider_ids( &self, provider_ids: &[String], diff --git a/crates/aether-model-fetch/src/lib.rs b/crates/aether-model-fetch/src/lib.rs index 6ce0e69c6..e86e8fa63 100644 --- a/crates/aether-model-fetch/src/lib.rs +++ b/crates/aether-model-fetch/src/lib.rs @@ -12,22 +12,24 @@ pub use config::{ }; pub use logic::{ aggregate_models_for_cache, apply_model_filters, build_models_fetch_url, - deepseek_anthropic_models_fetch_uses_openai_auth, endpoint_supports_rust_models_fetch, - extract_error_message, json_string_list, merge_upstream_metadata, - model_catalog_upstream_metadata, parse_models_response, parse_models_response_page, - parse_windsurf_model_configs_response, preset_models_for_provider, + build_models_fetch_url_for_client_version, deepseek_anthropic_models_fetch_uses_openai_auth, + endpoint_supports_rust_models_fetch, extract_error_message, json_string_list, + merge_upstream_metadata, model_catalog_upstream_metadata, parse_models_response, + parse_models_response_page, parse_windsurf_model_configs_response, preset_models_for_provider, provider_type_uses_preset_models, select_models_fetch_endpoint, selected_models_fetch_endpoints, upstream_metadata_namespace_updates, ModelFetchRunSummary, ModelsFetchPage, ModelsFetchSuccess, }; pub use strategy::{ - fetch_models_from_transports, ModelFetchStrategy, ModelFetchStrategyKind, ModelsFetchOutcome, - SelectedModelFetchStrategy, + fetch_models_from_transports, fetch_models_from_transports_for_client_version, + ModelFetchStrategy, ModelFetchStrategyKind, ModelsFetchOutcome, SelectedModelFetchStrategy, }; pub use transport::{ build_antigravity_fetch_available_models_plan, build_antigravity_load_code_assist_plan, build_gemini_cli_load_code_assist_plan, build_kiro_list_available_models_plan, - build_models_fetch_execution_plan, build_standard_models_fetch_execution_plan, + build_models_fetch_execution_plan, build_models_fetch_execution_plan_for_client_version, + build_standard_models_fetch_execution_plan, + build_standard_models_fetch_execution_plan_for_client_version, build_vertex_models_fetch_execution_plan, build_windsurf_model_configs_execution_plan, ModelFetchTransportRuntime, }; diff --git a/crates/aether-model-fetch/src/logic.rs b/crates/aether-model-fetch/src/logic.rs index a7d425129..b353f97cc 100644 --- a/crates/aether-model-fetch/src/logic.rs +++ b/crates/aether-model-fetch/src/logic.rs @@ -21,6 +21,9 @@ const MODEL_FETCH_FORMAT_PRIORITY: &[&[&str]] = &[ &["gemini:generate_content"], ]; +pub(crate) const CODEX_MODELS_MAX_ITEMS: usize = 512; +pub(crate) const CODEX_MODELS_MAX_JSON_BYTES: usize = 8 * 1024 * 1024; + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct ModelFetchRunSummary { pub attempted: usize, @@ -66,6 +69,15 @@ pub fn build_models_fetch_url( provider_type: &str, endpoint_api_format: &str, base_url: &str, +) -> Option<(String, String)> { + build_models_fetch_url_for_client_version(provider_type, endpoint_api_format, base_url, None) +} + +pub fn build_models_fetch_url_for_client_version( + provider_type: &str, + endpoint_api_format: &str, + base_url: &str, + codex_client_version: Option<&str>, ) -> Option<(String, String)> { let api_format = normalize_api_format(endpoint_api_format); if !endpoint_supports_rust_models_fetch(&api_format) { @@ -73,7 +85,7 @@ pub fn build_models_fetch_url( } let provider_type = provider_type.trim().to_ascii_lowercase(); let url = if provider_type == "codex" && api_format.starts_with("openai:") { - build_codex_models_url(base_url) + build_codex_models_url(base_url, codex_client_version) } else if api_format.starts_with("openai:") { build_v1_models_url(base_url) } else if api_format.starts_with("claude:") { @@ -173,6 +185,142 @@ pub fn parse_models_response_page( }) } +/// Parses the Codex `/models` response without applying the generic cache projection. +/// +/// Codex model cards are versioned protocol data. They must remain opaque so future fields and +/// instruction representations survive catalog caching and downstream projection. Invalid entries +/// reject the whole response instead of being skipped and accidentally replacing a complete LKG +/// with a partial directory. +pub(crate) fn parse_codex_models_response_page(body: &Value) -> Result { + let serialized = serde_json::to_vec(body) + .map_err(|_| "Codex models response could not be serialized".to_string())?; + if serialized.len() > CODEX_MODELS_MAX_JSON_BYTES { + return Err(format!( + "Codex models response exceeds {CODEX_MODELS_MAX_JSON_BYTES} bytes" + )); + } + + let items = body + .get("models") + .and_then(Value::as_array) + .ok_or_else(|| "Codex models response is missing models array".to_string())?; + if items.is_empty() { + return Err("Codex models response contains no models".to_string()); + } + if items.len() > CODEX_MODELS_MAX_ITEMS { + return Err(format!( + "Codex models response exceeds {CODEX_MODELS_MAX_ITEMS} models" + )); + } + + let cached_models = merge_codex_models_preserving_cards(items)?; + let fetched_model_ids = cached_models + .iter() + .filter_map(codex_model_identity) + .map(ToOwned::to_owned) + .collect(); + + Ok(ModelsFetchPage { + fetched_model_ids, + cached_models, + has_more: false, + next_after_id: None, + }) +} + +/// Merges opaque Codex model cards without silently selecting one of two conflicting cards. +/// +/// Both `id` and `slug` are mapping identities. Exact duplicate JSON cards can occur when the +/// same catalog is fetched through multiple endpoint transports and are collapsed. If any valid +/// identity is reused by a different card, the response is ambiguous and must not replace a +/// last-known-good catalog. +pub(crate) fn merge_codex_models_preserving_cards(models: &[Value]) -> Result, String> { + let mut merged = Vec::::with_capacity(models.len()); + let mut index_by_identity = BTreeMap::::new(); + + for model in models { + let identities = codex_model_identities(model)?; + let mut duplicate_index = None; + for identity in &identities { + let Some(existing_index) = index_by_identity.get(*identity).copied() else { + continue; + }; + if merged.get(existing_index) != Some(model) { + return Err(format!( + "Codex models response contains conflicting cards for identity '{identity}'" + )); + } + if duplicate_index.is_some_and(|index| index != existing_index) { + return Err(format!( + "Codex models response contains conflicting cards for identity '{identity}'" + )); + } + duplicate_index = Some(existing_index); + } + + if let Some(existing_index) = duplicate_index { + for identity in identities { + index_by_identity + .entry(identity.to_string()) + .or_insert(existing_index); + } + continue; + } + + let model_index = merged.len(); + merged.push(model.clone()); + for identity in identities { + index_by_identity.insert(identity.to_string(), model_index); + } + } + + Ok(merged) +} + +fn codex_model_identities(model: &Value) -> Result, String> { + let object = model + .as_object() + .ok_or_else(|| "Codex models response contains a non-object model card".to_string())?; + let mut identities = Vec::with_capacity(2); + for field in ["id", "slug"] { + let Some(identity) = object + .get(field) + .and_then(Value::as_str) + .filter(|value| *value == value.trim()) + .filter(|value| valid_codex_model_identity(value)) + else { + continue; + }; + if !identities.contains(&identity) { + identities.push(identity); + } + } + if identities.is_empty() { + return Err("Codex models response contains a card without a valid id or slug".to_string()); + } + Ok(identities) +} + +pub(crate) fn codex_model_identity(model: &Value) -> Option<&str> { + let object = model.as_object()?; + ["id", "slug"].iter().find_map(|field| { + object + .get(*field) + .and_then(Value::as_str) + .filter(|value| *value == value.trim()) + .filter(|value| valid_codex_model_identity(value)) + }) +} + +fn valid_codex_model_identity(value: &str) -> bool { + !value.is_empty() + && value.len() <= 256 + && value == value.trim() + && value + .chars() + .all(|character| !character.is_whitespace() && !character.is_control()) +} + pub fn parse_windsurf_model_configs_response( body: &Value, updated_at_unix_secs: u64, @@ -678,9 +826,17 @@ fn build_deepseek_anthropic_models_url(base_url: &str) -> Option { Some(url) } -fn build_codex_models_url(base_url: &str) -> Option { +fn build_codex_models_url(base_url: &str, client_version: Option<&str>) -> Option { if let Some(url) = build_bigmodel_coding_models_url(base_url) { - return Some(url); + return Some( + client_version + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(|client_version| { + replace_or_append_query_param(&url, "client_version", client_version) + }) + .unwrap_or(url), + ); } let (trimmed_base_url, query) = split_url_query(base_url); @@ -692,34 +848,77 @@ fn build_codex_models_url(base_url: &str) -> Option { || trimmed_base_url.ends_with("/codex") || trimmed_base_url.ends_with("/models"); if !is_codex_backend && openai_compatible_base_includes_unversioned_api_root(base_url) { - return build_openai_compatible_models_url(base_url); + let url = build_openai_compatible_models_url(base_url)?; + return Some( + client_version + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(|client_version| { + replace_or_append_query_param(&url, "client_version", client_version) + }) + .unwrap_or(url), + ); } let mut url = if trimmed_base_url.ends_with("/models") { trimmed_base_url.to_string() } else { format!("{trimmed_base_url}/models") }; - let mut has_client_version = false; - if let Some(query) = query.filter(|value| !value.trim().is_empty()) { - has_client_version = query.split('&').any(|part| { - part.split_once('=') + let explicit_client_version = client_version + .map(str::trim) + .filter(|value| !value.is_empty()); + let mut query_parts = query + .filter(|value| !value.trim().is_empty()) + .map(|value| value.split('&').map(ToOwned::to_owned).collect::>()) + .unwrap_or_default(); + let has_client_version = query_parts.iter().any(|part| { + part.split_once('=') + .map(|(key, _)| key) + .unwrap_or(part) + .trim() + .eq_ignore_ascii_case("client_version") + }); + if let Some(client_version) = explicit_client_version { + query_parts.retain(|part| { + !part + .split_once('=') .map(|(key, _)| key) .unwrap_or(part) .trim() .eq_ignore_ascii_case("client_version") }); - url.push('?'); - url.push_str(query); + query_parts.push(format!("client_version={client_version}")); + } else if !has_client_version { + query_parts.push(format!( + "client_version={}", + aether_ai_formats::CODEX_CLIENT_VERSION + )); } - if !has_client_version { - let separator = if url.contains('?') { '&' } else { '?' }; - url.push(separator); - url.push_str("client_version="); - url.push_str(aether_ai_formats::CODEX_CLIENT_VERSION); + if !query_parts.is_empty() { + url.push('?'); + url.push_str(&query_parts.join("&")); } Some(url) } +fn replace_or_append_query_param(url: &str, name: &str, value: &str) -> String { + let (base, query) = split_url_query(url); + let mut query_parts = query + .filter(|query| !query.trim().is_empty()) + .map(|query| query.split('&').map(ToOwned::to_owned).collect::>()) + .unwrap_or_default(); + query_parts.retain(|part| { + !part + .split_once('=') + .map(|(key, _)| key) + .unwrap_or(part) + .trim() + .eq_ignore_ascii_case(name) + }); + query_parts.push(format!("{name}={value}")); + format!("{base}?{}", query_parts.join("&")) +} + fn build_gemini_models_url(base_url: &str) -> Option { let (trimmed_base_url, base_query) = split_url_query(base_url); let trimmed_base_url = trimmed_base_url.trim_end_matches('/'); @@ -793,12 +992,13 @@ fn windsurf_json_f64(value: &Value) -> Option { fn collect_cached_model_ids(models: &[Value]) -> Vec { let mut ids = Vec::new(); for model in models { - let Some(model_id) = model - .get("id") - .and_then(Value::as_str) - .map(str::trim) - .filter(|value| !value.is_empty()) - else { + let Some(model_id) = codex_model_identity(model).or_else(|| { + model + .get("id") + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + }) else { continue; }; ids.push(model_id.to_string()); @@ -880,8 +1080,9 @@ mod tests { use super::{ aggregate_models_for_cache, apply_model_filters, build_gemini_models_url, - build_models_fetch_url, merge_upstream_metadata, parse_models_response, - parse_models_response_page, preset_models_for_provider, selected_models_fetch_endpoints, + build_models_fetch_url, build_models_fetch_url_for_client_version, merge_upstream_metadata, + parse_codex_models_response_page, parse_models_response, parse_models_response_page, + preset_models_for_provider, selected_models_fetch_endpoints, }; fn sample_endpoint( @@ -1036,6 +1237,56 @@ mod tests { ); } + #[test] + fn build_models_fetch_url_uses_explicit_codex_client_version() { + assert_eq!( + build_models_fetch_url_for_client_version( + "codex", + "openai:responses", + "https://chatgpt.com/backend-api/codex", + Some("0.145.2"), + ), + Some(( + "https://chatgpt.com/backend-api/codex/models?client_version=0.145.2".to_string(), + "openai:responses".to_string() + )) + ); + } + + #[test] + fn explicit_codex_client_version_replaces_stale_base_query_value() { + assert_eq!( + build_models_fetch_url_for_client_version( + "codex", + "openai:responses", + "https://chatgpt.com/backend-api/codex?feature=on&client_version=0.144.1", + Some("0.145.2"), + ), + Some(( + "https://chatgpt.com/backend-api/codex/models?feature=on&client_version=0.145.2" + .to_string(), + "openai:responses".to_string() + )) + ); + } + + #[test] + fn explicit_codex_client_version_is_forwarded_through_compatible_proxy_roots() { + assert_eq!( + build_models_fetch_url_for_client_version( + "codex", + "openai:responses", + "https://proxy.example.com/api?feature=on&client_version=0.144.1", + Some("0.145.2"), + ), + Some(( + "https://proxy.example.com/api/models?feature=on&client_version=0.145.2" + .to_string(), + "openai:responses".to_string() + )) + ); + } + #[test] fn build_models_fetch_url_supports_bigmodel_coding_paas_root() { assert_eq!( @@ -1060,6 +1311,19 @@ mod tests { "openai:responses".to_string() )) ); + assert_eq!( + build_models_fetch_url_for_client_version( + "codex", + "openai:responses", + "https://open.bigmodel.cn/api/coding/paas/v4?tenant=demo&client_version=0.144.1", + Some("0.145.2"), + ), + Some(( + "https://open.bigmodel.cn/api/coding/paas/v4/models?tenant=demo&client_version=0.145.2" + .to_string(), + "openai:responses".to_string() + )) + ); } #[test] @@ -1185,6 +1449,97 @@ mod tests { assert_eq!(cached["api_formats"], json!(["openai:responses"])); } + #[test] + fn strict_codex_parser_preserves_opaque_cards_without_cache_projection() { + let card = json!({ + "id": "gpt-future-dynamic", + "slug": "gpt-future-dynamic", + "api_format": "future-protocol-field", + "model_messages": {"instructions_template": "Future instructions"}, + "available_in_plans": ["plus"], + "future_capability": {"opaque": true} + }); + let parsed = parse_codex_models_response_page(&json!({"models": [card.clone()]})) + .expect("opaque Codex card should parse"); + + assert_eq!(parsed.fetched_model_ids, vec!["gpt-future-dynamic"]); + assert_eq!(parsed.cached_models, vec![card]); + } + + #[test] + fn strict_codex_parser_rejects_empty_models_array() { + let error = parse_codex_models_response_page(&json!({"models": []})) + .expect_err("empty Codex catalog must fail"); + + assert!(error.contains("no models")); + } + + #[test] + fn strict_codex_parser_merges_only_exact_duplicate_cards() { + let card = json!({ + "id": "gpt-future-duplicate", + "slug": "gpt-future-duplicate", + "model_messages": {"instructions_template": "Opaque instructions"}, + "future_capability": {"opaque": true} + }); + let parsed = parse_codex_models_response_page(&json!({ + "models": [card.clone(), card.clone()] + })) + .expect("exact duplicate Codex cards should merge"); + + assert_eq!(parsed.fetched_model_ids, vec!["gpt-future-duplicate"]); + assert_eq!(parsed.cached_models, vec![card]); + } + + #[test] + fn strict_codex_parser_rejects_same_or_cross_identity_conflicts() { + let conflicts = [ + json!({ + "models": [ + {"id": "gpt-conflict", "slug": "gpt-conflict", "future": 1}, + {"id": "gpt-conflict", "slug": "gpt-conflict", "future": 2} + ] + }), + json!({ + "models": [ + {"id": "gpt-id-one", "slug": "gpt-cross-identity", "future": 1}, + {"id": "gpt-cross-identity", "slug": "gpt-slug-two", "future": 2} + ] + }), + ]; + + for body in conflicts { + let error = parse_codex_models_response_page(&body) + .expect_err("ambiguous Codex identities must fail"); + assert!(error.contains("conflicting cards")); + } + } + + #[test] + fn strict_codex_parser_rejects_non_object_or_synthetic_identity_cards() { + for body in [ + json!({"models": ["gpt-future-dynamic"]}), + json!({"models": [{"model": "gpt-future-dynamic"}]}), + json!({"models": [{"name": "gpt-future-dynamic"}]}), + json!({"models": [{"slug": " gpt-future-dynamic "}]}), + ] { + assert!(parse_codex_models_response_page(&body).is_err()); + } + } + + #[test] + fn strict_codex_parser_does_not_impose_an_ascii_symbol_allowlist_on_identities() { + let card = json!({ + "slug": "gpt+future@dynamic", + "model_messages": {"instructions_template": "Future instructions"} + }); + let parsed = parse_codex_models_response_page(&json!({"models": [card.clone()]})) + .expect("future identity punctuation should remain opaque"); + + assert_eq!(parsed.fetched_model_ids, vec!["gpt+future@dynamic"]); + assert_eq!(parsed.cached_models, vec![card]); + } + #[test] fn parse_models_response_page_reads_claude_pagination_state() { let parsed = parse_models_response_page( diff --git a/crates/aether-model-fetch/src/strategy.rs b/crates/aether-model-fetch/src/strategy.rs index bc00a0850..bfcb582c3 100644 --- a/crates/aether-model-fetch/src/strategy.rs +++ b/crates/aether-model-fetch/src/strategy.rs @@ -20,14 +20,16 @@ use serde_json::{json, Value}; use sha2::Sha256; use crate::logic::{ - aggregate_models_for_cache, extract_error_message, parse_models_response_page, - parse_windsurf_model_configs_response, preset_models_for_provider, + aggregate_models_for_cache, codex_model_identity, extract_error_message, + merge_codex_models_preserving_cards, parse_codex_models_response_page, + parse_models_response_page, parse_windsurf_model_configs_response, preset_models_for_provider, }; use crate::transport::{ build_antigravity_fetch_available_models_plan, build_antigravity_load_code_assist_plan, build_gemini_cli_load_code_assist_plan, build_kiro_list_available_models_plan, - build_standard_models_fetch_execution_plan, build_vertex_models_fetch_execution_plan, - build_windsurf_model_configs_execution_plan, ModelFetchTransportRuntime, + build_standard_models_fetch_execution_plan_for_client_version, + build_vertex_models_fetch_execution_plan, build_windsurf_model_configs_execution_plan, + ModelFetchTransportRuntime, }; const ANTIGRAVITY_SANDBOX_BASE_URL: &str = "https://daily-cloudcode-pa.sandbox.googleapis.com"; @@ -48,6 +50,46 @@ pub struct ModelsFetchOutcome { pub errors: Vec, pub has_success: bool, pub upstream_metadata: Option, + pub etag: Option, + pub upstream_status: Option, +} + +#[derive(Debug)] +struct ConsistentValue { + value: Option, + observed: bool, + consistent: bool, +} + +impl Default for ConsistentValue { + fn default() -> Self { + Self { + value: None, + observed: false, + consistent: true, + } + } +} + +impl ConsistentValue { + fn observe(&mut self, candidate: Option) { + if !self.observed { + self.consistent = candidate.is_some(); + self.value = candidate; + self.observed = true; + return; + } + if self.value.as_ref() != candidate.as_ref() { + self.consistent = false; + self.value = None; + } + } + + fn finish(self) -> Option { + (self.observed && self.consistent) + .then_some(self.value) + .flatten() + } } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -87,9 +129,17 @@ impl ModelFetchStrategy for SelectedModelFetchStrategy { pub async fn fetch_models_from_transports( runtime: &(impl ModelFetchTransportRuntime + ?Sized), transports: &[GatewayProviderTransportSnapshot], +) -> Result { + fetch_models_from_transports_for_client_version(runtime, transports, None).await +} + +pub async fn fetch_models_from_transports_for_client_version( + runtime: &(impl ModelFetchTransportRuntime + ?Sized), + transports: &[GatewayProviderTransportSnapshot], + codex_client_version: Option<&str>, ) -> Result { let strategy = select_model_fetch_strategy(transports)?; - execute_model_fetch_strategy(runtime, transports, strategy).await + execute_model_fetch_strategy(runtime, transports, strategy, codex_client_version).await } fn select_model_fetch_strategy( @@ -158,6 +208,7 @@ async fn execute_model_fetch_strategy( runtime: &(impl ModelFetchTransportRuntime + ?Sized), transports: &[GatewayProviderTransportSnapshot], strategy: SelectedModelFetchStrategy, + codex_client_version: Option<&str>, ) -> Result { let Some(first_transport) = transports.first() else { return Err("No transport snapshots available for models fetch".to_string()); @@ -170,7 +221,13 @@ async fn execute_model_fetch_strategy( true, )), ModelFetchStrategyKind::StandardTransport => { - fetch_standard_models(runtime, transports, strategy.provider_id()).await + fetch_standard_models( + runtime, + transports, + strategy.provider_id(), + codex_client_version, + ) + .await } ModelFetchStrategyKind::Vertex => fetch_vertex_models(runtime, transports).await, ModelFetchStrategyKind::Antigravity => { @@ -193,60 +250,110 @@ async fn fetch_standard_models( runtime: &(impl ModelFetchTransportRuntime + ?Sized), transports: &[GatewayProviderTransportSnapshot], provider_type: &str, + codex_client_version: Option<&str>, ) -> Result { let mut all_models = Vec::new(); let mut errors = Vec::new(); let mut has_success = false; + let mut etag = ConsistentValue::default(); + let mut upstream_status = ConsistentValue::default(); for transport in transports { - match fetch_standard_models_for_transport(runtime, transport).await { + match fetch_standard_models_for_transport(runtime, transport, codex_client_version).await { Ok(outcome) => { all_models.extend(outcome.cached_models); has_success |= outcome.has_success; + if outcome.has_success { + etag.observe(outcome.etag); + upstream_status.observe(outcome.upstream_status); + } + } + Err((err, status)) => { + upstream_status.observe(status); + errors.push(format!("{}: {err}", transport.endpoint.api_format.trim())); } - Err(err) => errors.push(format!("{}: {err}", transport.endpoint.api_format.trim())), } } - let merged_models = aggregate_models_for_cache(&all_models); + let is_codex = provider_type.trim().eq_ignore_ascii_case("codex"); + let merged_models = if is_codex { + merge_codex_models_preserving_cards(&all_models)? + } else { + aggregate_models_for_cache(&all_models) + }; + let codex_model_ids = is_codex.then(|| collect_codex_model_ids(&merged_models)); let upstream_metadata = crate::logic::model_catalog_upstream_metadata(provider_type, &merged_models); - Ok(build_success_outcome(merged_models, upstream_metadata, has_success).with_errors(errors)) + let mut outcome = build_success_outcome(merged_models, upstream_metadata, has_success); + if let Some(model_ids) = codex_model_ids { + outcome.fetched_model_ids = model_ids; + } + Ok(outcome + .with_errors(errors) + .with_etag(etag.finish()) + .with_upstream_status(upstream_status.finish())) } async fn fetch_standard_models_for_transport( runtime: &(impl ModelFetchTransportRuntime + ?Sized), transport: &GatewayProviderTransportSnapshot, -) -> Result { + codex_client_version: Option<&str>, +) -> Result)> { let mut all_models = Vec::new(); let mut seen_ids = BTreeSet::new(); let mut next_after_id = None; let mut has_success = false; + let mut etag = ConsistentValue::default(); + let mut upstream_status = ConsistentValue::default(); + let is_codex = transport + .provider + .provider_type + .trim() + .eq_ignore_ascii_case("codex"); for _ in 0..20 { - let plan = build_standard_models_fetch_execution_plan( + let plan = build_standard_models_fetch_execution_plan_for_client_version( runtime, transport, next_after_id.as_deref(), + codex_client_version, ) - .await?; - let result = runtime.execute_model_fetch_execution_plan(&plan).await?; - let body_json = execution_result_json_body(&result)?; - let parsed = parse_models_response_page(&transport.endpoint.api_format, &body_json)?; + .await + .map_err(|err| (err, None))?; + let result = runtime + .execute_model_fetch_execution_plan(&plan) + .await + .map_err(|err| (err, None))?; + upstream_status.observe(Some(result.status_code)); + let body_json = + execution_result_json_body(&result).map_err(|err| (err, Some(result.status_code)))?; + let parsed = if is_codex { + parse_codex_models_response_page(&body_json) + } else { + parse_models_response_page(&transport.endpoint.api_format, &body_json) + } + .map_err(|err| (err, Some(result.status_code)))?; + etag.observe(execution_result_header(&result, "etag")); has_success = true; - for model in parsed.cached_models { - let Some(model_id) = model - .get("id") - .and_then(Value::as_str) - .map(str::trim) - .filter(|value| !value.is_empty()) - else { - continue; - }; - if !seen_ids.insert(model_id.to_string()) { - continue; + if is_codex { + // Preserve every opaque card until the catalog-wide merge can distinguish exact + // duplicates from conflicting `id`/`slug` identities across endpoint transports. + all_models.extend(parsed.cached_models); + } else { + for model in parsed.cached_models { + let Some(model_id) = model + .get("id") + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + else { + continue; + }; + if !seen_ids.insert(model_id.to_string()) { + continue; + } + all_models.push(model); } - all_models.push(model); } let Some(next_cursor) = parsed @@ -260,7 +367,9 @@ async fn fetch_standard_models_for_transport( next_after_id = Some(next_cursor); } - Ok(build_success_outcome(all_models, None, has_success)) + Ok(build_success_outcome(all_models, None, has_success) + .with_etag(etag.finish()) + .with_upstream_status(upstream_status.finish())) } async fn fetch_antigravity_models( @@ -319,6 +428,8 @@ async fn fetch_antigravity_models( errors, has_success: false, upstream_metadata: None, + etag: None, + upstream_status: None, }) } @@ -499,6 +610,8 @@ async fn fetch_vertex_api_key_models( errors: vec!["vertex_ai(api_key): missing api key".to_string()], has_success: false, upstream_metadata: None, + etag: None, + upstream_status: None, }); } @@ -556,6 +669,8 @@ async fn fetch_vertex_api_key_models( errors, has_success, upstream_metadata: None, + etag: None, + upstream_status: None, }) } @@ -571,6 +686,8 @@ async fn fetch_vertex_service_account_models( errors: vec!["vertex_ai(service_account): missing auth_config".to_string()], has_success: false, upstream_metadata: None, + etag: None, + upstream_status: None, }); }; let token = exchange_vertex_service_account_token(runtime, &transports[0], auth_config).await?; @@ -639,6 +756,8 @@ async fn fetch_vertex_service_account_models( errors, has_success, upstream_metadata: None, + etag: None, + upstream_status: None, }) } @@ -822,8 +941,18 @@ fn execution_result_json_body_allow_empty(result: &ExecutionResult) -> Result String { +fn execution_result_header(result: &ExecutionResult, name: &str) -> Option { result + .headers + .iter() + .find(|(header_name, _)| header_name.eq_ignore_ascii_case(name)) + .map(|(_, value)| value.trim()) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) +} + +fn execution_result_error_message(result: &ExecutionResult) -> String { + let detail = result .body .as_ref() .and_then(|body| body.json_body.as_ref()) @@ -833,8 +962,14 @@ fn execution_result_error_message(result: &ExecutionResult) -> String { let message = error.message.trim(); (!message.is_empty()).then_some(message.to_string()) }) - }) - .unwrap_or_else(|| format!("HTTP {}: upstream request failed", result.status_code)) + }); + match detail { + Some(detail) if !(200..300).contains(&result.status_code) => { + format!("HTTP {}: {detail}", result.status_code) + } + Some(detail) => detail, + None => format!("HTTP {}: upstream request failed", result.status_code), + } } fn parse_antigravity_models_response(body: &Value) -> Result<(Vec, Option), String> { @@ -1238,6 +1373,8 @@ fn build_success_outcome( errors: Vec::new(), has_success, upstream_metadata, + etag: None, + upstream_status: None, } } @@ -1260,6 +1397,16 @@ fn collect_model_ids(models: &[Value]) -> Vec { ids } +fn collect_codex_model_ids(models: &[Value]) -> Vec { + let mut seen = BTreeSet::new(); + models + .iter() + .filter_map(codex_model_identity) + .filter(|model_id| seen.insert((*model_id).to_string())) + .map(ToOwned::to_owned) + .collect() +} + fn transport_auth_config(transport: &GatewayProviderTransportSnapshot) -> Option { transport .key @@ -1397,6 +1544,8 @@ fn now_unix_secs() -> u64 { trait OutcomeExt { fn with_errors(self, errors: Vec) -> Self; + fn with_etag(self, etag: Option) -> Self; + fn with_upstream_status(self, upstream_status: Option) -> Self; } impl OutcomeExt for ModelsFetchOutcome { @@ -1404,6 +1553,16 @@ impl OutcomeExt for ModelsFetchOutcome { self.errors = errors; self } + + fn with_etag(mut self, etag: Option) -> Self { + self.etag = etag; + self + } + + fn with_upstream_status(mut self, upstream_status: Option) -> Self { + self.upstream_status = upstream_status; + self + } } #[cfg(test)] @@ -1423,8 +1582,8 @@ mod tests { build_vertex_google_list_url, build_vertex_service_account_list_url, select_model_fetch_strategy, ModelFetchStrategy, ModelFetchStrategyKind, }; - use crate::fetch_models_from_transports; use crate::transport::ModelFetchTransportRuntime; + use crate::{fetch_models_from_transports, fetch_models_from_transports_for_client_version}; type RouteResult = Result<(u16, Value), String>; type ModelFetchRoute = (String, RouteResult); @@ -1433,6 +1592,7 @@ mod tests { executed_urls: Arc>>, response_body: Value, status_code: u16, + response_headers: BTreeMap, } struct RoutingTestRuntime { @@ -1474,7 +1634,7 @@ mod tests { request_id: plan.request_id.clone(), candidate_id: plan.candidate_id.clone(), status_code: self.status_code, - headers: BTreeMap::new(), + headers: self.response_headers.clone(), response_observation: None, body: Some(ResponseBody { json_body: Some(self.response_body.clone()), @@ -1666,6 +1826,16 @@ mod tests { transport } + fn sample_codex_transport_for_base( + endpoint_id: &str, + base_url: &str, + ) -> GatewayProviderTransportSnapshot { + let mut transport = sample_codex_transport(); + transport.endpoint.id = endpoint_id.to_string(); + transport.endpoint.base_url = base_url.to_string(); + transport + } + fn sample_kiro_transport() -> GatewayProviderTransportSnapshot { let mut transport = sample_custom_aiplatform_transport(); transport.provider.provider_type = "kiro".to_string(); @@ -1801,6 +1971,7 @@ mod tests { }] }), status_code: 200, + response_headers: BTreeMap::new(), }; let outcome = fetch_models_from_transports(&runtime, &[sample_custom_aiplatform_transport()]) @@ -2004,24 +2175,40 @@ mod tests { "models": [{ "id": "gpt-5.6-future", "slug": "gpt-5.6-future", + "api_format": "opaque-future-field", "default_reasoning_level": "high", "supported_reasoning_levels": [{"effort": "high"}], "future_capability": {"mode": "preserve-me"} }] }), status_code: 200, + response_headers: BTreeMap::from([( + "ETag".to_string(), + "\"codex-models-0.145.2\"".to_string(), + )]), }; - let outcome = fetch_models_from_transports(&runtime, &[sample_codex_transport()]) - .await - .expect("models fetch should succeed"); + let outcome = fetch_models_from_transports_for_client_version( + &runtime, + &[sample_codex_transport()], + Some("0.145.2"), + ) + .await + .expect("models fetch should succeed"); let urls = executed_urls.lock().expect("executed_urls lock"); assert_eq!( urls.as_slice(), - &["https://chatgpt.com/backend-api/codex/models?client_version=0.144.1"] + &["https://chatgpt.com/backend-api/codex/models?client_version=0.145.2"] ); + assert_eq!(outcome.etag.as_deref(), Some("\"codex-models-0.145.2\"")); + assert_eq!(outcome.upstream_status, Some(200)); assert_eq!(outcome.fetched_model_ids, vec!["gpt-5.6-future"]); assert_eq!(outcome.cached_models.len(), 1); + assert_eq!( + outcome.cached_models[0]["api_format"], + "opaque-future-field" + ); + assert!(outcome.cached_models[0].get("api_formats").is_none()); let card = &outcome .upstream_metadata .as_ref() @@ -2030,6 +2217,159 @@ mod tests { assert_eq!(card["future_capability"]["mode"], "preserve-me"); } + #[tokio::test] + async fn codex_transport_reports_slug_only_ids_without_rewriting_opaque_cards() { + let executed_urls = Arc::new(Mutex::new(Vec::new())); + let card = json!({ + "slug": "gpt-slug-only-future", + "model_messages": {"instructions_template": "Slug-only instructions"}, + "future_capability": {"opaque": true} + }); + let runtime = TestRuntime { + executed_urls, + response_body: json!({"models": [card.clone()]}), + status_code: 200, + response_headers: BTreeMap::new(), + }; + + let outcome = fetch_models_from_transports_for_client_version( + &runtime, + &[sample_codex_transport()], + Some("0.145.2"), + ) + .await + .expect("slug-only Codex card should fetch"); + + assert_eq!(outcome.fetched_model_ids, vec!["gpt-slug-only-future"]); + assert_eq!(outcome.cached_models, vec![card]); + assert!(outcome.cached_models[0].get("id").is_none()); + } + + #[tokio::test] + async fn codex_transport_merges_exact_duplicate_cards_across_endpoints() { + let executed_urls = Arc::new(Mutex::new(Vec::new())); + let card = json!({ + "id": "gpt-exact-duplicate", + "slug": "gpt-exact-duplicate", + "model_messages": {"instructions_template": "Opaque instructions"}, + "future_capability": {"opaque": true} + }); + let runtime = RoutingTestRuntime { + executed_urls, + routes: vec![ + ( + "first.example.com/backend-api/codex/models".to_string(), + Ok((200, json!({"models": [card.clone()]}))), + ), + ( + "second.example.com/backend-api/codex/models".to_string(), + Ok((200, json!({"models": [card.clone()]}))), + ), + ], + }; + let transports = vec![ + sample_codex_transport_for_base( + "endpoint-first", + "https://first.example.com/backend-api/codex", + ), + sample_codex_transport_for_base( + "endpoint-second", + "https://second.example.com/backend-api/codex", + ), + ]; + + let outcome = + fetch_models_from_transports_for_client_version(&runtime, &transports, Some("0.145.2")) + .await + .expect("exact duplicate endpoint catalogs should merge"); + + assert!(outcome.has_success); + assert!(outcome.errors.is_empty()); + assert_eq!(outcome.fetched_model_ids, vec!["gpt-exact-duplicate"]); + assert_eq!(outcome.cached_models, vec![card]); + } + + #[tokio::test] + async fn codex_transport_rejects_cross_identity_conflicts_across_endpoints() { + let executed_urls = Arc::new(Mutex::new(Vec::new())); + let runtime = RoutingTestRuntime { + executed_urls, + routes: vec![ + ( + "first.example.com/backend-api/codex/models".to_string(), + Ok(( + 200, + json!({ + "models": [{ + "id": "gpt-id-one", + "slug": "gpt-cross-identity", + "future_capability": {"source": "first"} + }] + }), + )), + ), + ( + "second.example.com/backend-api/codex/models".to_string(), + Ok(( + 200, + json!({ + "models": [{ + "id": "gpt-cross-identity", + "slug": "gpt-slug-two", + "future_capability": {"source": "second"} + }] + }), + )), + ), + ], + }; + let transports = vec![ + sample_codex_transport_for_base( + "endpoint-first", + "https://first.example.com/backend-api/codex", + ), + sample_codex_transport_for_base( + "endpoint-second", + "https://second.example.com/backend-api/codex", + ), + ]; + + let error = + fetch_models_from_transports_for_client_version(&runtime, &transports, Some("0.145.2")) + .await + .expect_err("conflicting endpoint catalogs must fail"); + + assert!(error.contains("conflicting cards")); + assert!(error.contains("gpt-cross-identity")); + } + + #[tokio::test] + async fn codex_transport_reports_non_success_upstream_status() { + let executed_urls = Arc::new(Mutex::new(Vec::new())); + let runtime = TestRuntime { + executed_urls, + response_body: json!({ + "error": { "message": "temporarily unavailable" } + }), + status_code: 503, + response_headers: BTreeMap::new(), + }; + + let outcome = fetch_models_from_transports_for_client_version( + &runtime, + &[sample_codex_transport()], + Some("0.145.2"), + ) + .await + .expect("models fetch should return an observable failed outcome"); + + assert!(!outcome.has_success); + assert_eq!(outcome.upstream_status, Some(503)); + assert_eq!(outcome.etag, None); + assert_eq!(outcome.errors.len(), 1); + assert!(outcome.errors[0].contains("HTTP 503: temporarily unavailable")); + } + #[tokio::test] async fn gemini_cli_load_code_assist_preserves_paid_tier_credits() { let executed_urls = Arc::new(Mutex::new(Vec::new())); @@ -2053,6 +2393,7 @@ mod tests { } }), status_code: 200, + response_headers: BTreeMap::new(), }; let outcome = fetch_models_from_transports(&runtime, &[sample_gemini_cli_transport()]) .await @@ -2189,6 +2530,7 @@ mod tests { ] }), status_code: 200, + response_headers: BTreeMap::new(), }; let outcome = fetch_models_from_transports(&runtime, &[sample_kiro_transport()]) .await @@ -2252,6 +2594,7 @@ mod tests { } }), status_code: 200, + response_headers: BTreeMap::new(), }; let outcome = fetch_models_from_transports(&runtime, &[sample_windsurf_transport()]) .await diff --git a/crates/aether-model-fetch/src/transport.rs b/crates/aether-model-fetch/src/transport.rs index 94a7e6857..e49a5f343 100644 --- a/crates/aether-model-fetch/src/transport.rs +++ b/crates/aether-model-fetch/src/transport.rs @@ -23,7 +23,9 @@ use aether_provider_transport::{ use async_trait::async_trait; use serde_json::json; -use crate::{build_models_fetch_url, deepseek_anthropic_models_fetch_uses_openai_auth}; +use crate::{ + build_models_fetch_url_for_client_version, deepseek_anthropic_models_fetch_uses_openai_auth, +}; const CLAUDE_CLI_USER_AGENT: &str = "claude-code/1.0.1"; const GEMINI_CLI_USER_AGENT: &str = "GeminiCLI/0.1.5 (Windows; AMD64)"; @@ -75,7 +77,21 @@ pub async fn build_models_fetch_execution_plan( runtime: &(impl ModelFetchTransportRuntime + ?Sized), transport: &GatewayProviderTransportSnapshot, ) -> Result { - build_standard_models_fetch_execution_plan(runtime, transport, None).await + build_models_fetch_execution_plan_for_client_version(runtime, transport, None).await +} + +pub async fn build_models_fetch_execution_plan_for_client_version( + runtime: &(impl ModelFetchTransportRuntime + ?Sized), + transport: &GatewayProviderTransportSnapshot, + codex_client_version: Option<&str>, +) -> Result { + build_standard_models_fetch_execution_plan_for_client_version( + runtime, + transport, + None, + codex_client_version, + ) + .await } struct ModelFetchExecutionPlanRequest { @@ -93,6 +109,18 @@ pub async fn build_standard_models_fetch_execution_plan( runtime: &(impl ModelFetchTransportRuntime + ?Sized), transport: &GatewayProviderTransportSnapshot, after_id: Option<&str>, +) -> Result { + build_standard_models_fetch_execution_plan_for_client_version( + runtime, transport, after_id, None, + ) + .await +} + +pub async fn build_standard_models_fetch_execution_plan_for_client_version( + runtime: &(impl ModelFetchTransportRuntime + ?Sized), + transport: &GatewayProviderTransportSnapshot, + after_id: Option<&str>, + codex_client_version: Option<&str>, ) -> Result { let api_format = transport.endpoint.api_format.trim().to_ascii_lowercase(); let provider_api_format = api_format.clone(); @@ -101,7 +129,8 @@ pub async fn build_standard_models_fetch_execution_plan( provider_type == "codex" && api_format.starts_with("openai:"); let is_deepseek_anthropic_models_fetch = api_format.starts_with("claude:") && deepseek_anthropic_models_fetch_uses_openai_auth(&transport.endpoint.base_url); - let mut headers = standard_models_fetch_headers(&api_format, &provider_type); + let mut headers = + standard_models_fetch_headers(&api_format, &provider_type, codex_client_version); if is_codex_openai_models_fetch { headers.insert("accept".to_string(), "application/json".to_string()); } @@ -110,6 +139,10 @@ pub async fn build_standard_models_fetch_execution_plan( headers.insert("accept".to_string(), "application/json".to_string()); } let mut protected_headers = Vec::::new(); + if is_codex_openai_models_fetch { + protected_headers.push("user-agent".to_string()); + protected_headers.push("originator".to_string()); + } if api_format.starts_with("openai:") || api_format.starts_with("claude:") { let resolved_auth = if is_deepseek_anthropic_models_fetch { @@ -155,7 +188,7 @@ pub async fn build_standard_models_fetch_execution_plan( headers = apply_fetch_header_rules(transport, headers, &protected_headers)?; } - let upstream_url = build_standard_models_fetch_url(transport, after_id)?; + let upstream_url = build_standard_models_fetch_url(transport, after_id, codex_client_version)?; build_execution_plan( runtime, transport, @@ -426,7 +459,8 @@ pub async fn build_vertex_models_fetch_execution_plan( api_format: &str, auth_header: Option<(String, String)>, ) -> Result { - let mut headers = standard_models_fetch_headers(api_format, &transport.provider.provider_type); + let mut headers = + standard_models_fetch_headers(api_format, &transport.provider.provider_type, None); let mut protected_headers = Vec::::new(); if let Some((name, value)) = auth_header { insert_non_empty_auth_header(&mut headers, &mut protected_headers, &name, &value); @@ -588,23 +622,34 @@ fn apply_fetch_header_rules( fn standard_models_fetch_headers( api_format: &str, provider_type: &str, + codex_client_version: Option<&str>, ) -> BTreeMap { let api_format = aether_ai_formats::normalize_api_format_alias(api_format); let provider_type = provider_type.trim().to_ascii_lowercase(); - match api_format.as_str() { - "openai:responses" | "openai:responses:compact" => { - let mut headers = BTreeMap::from([( + if provider_type == "codex" && api_format.starts_with("openai:") { + let client_version = codex_client_version + .map(str::trim) + .filter(|value| !value.is_empty()) + .unwrap_or(aether_ai_formats::CODEX_CLIENT_VERSION); + return BTreeMap::from([ + ( "user-agent".to_string(), - aether_ai_formats::CODEX_CLIENT_USER_AGENT.to_string(), - )]); - if provider_type == "codex" { - headers.insert( - "originator".to_string(), - aether_ai_formats::CODEX_CLIENT_ORIGINATOR.to_string(), - ); - } - headers - } + format!( + "{}/{client_version}", + aether_ai_formats::CODEX_CLIENT_ORIGINATOR + ), + ), + ( + "originator".to_string(), + aether_ai_formats::CODEX_CLIENT_ORIGINATOR.to_string(), + ), + ]); + } + match api_format.as_str() { + "openai:responses" | "openai:responses:compact" => BTreeMap::from([( + "user-agent".to_string(), + aether_ai_formats::CODEX_CLIENT_USER_AGENT.to_string(), + )]), "claude:messages" => { let mut headers = BTreeMap::from([( "anthropic-version".to_string(), @@ -632,6 +677,7 @@ fn standard_models_fetch_headers( fn build_standard_models_fetch_url( transport: &GatewayProviderTransportSnapshot, after_id: Option<&str>, + codex_client_version: Option<&str>, ) -> Result { let api_format = transport.endpoint.api_format.trim().to_ascii_lowercase(); if api_format.starts_with("gemini:") { @@ -648,19 +694,21 @@ fn build_standard_models_fetch_url( }) .ok_or_else(|| "Gemini models fetch requires an API key".to_string())?; - let (url, _) = build_models_fetch_url( + let (url, _) = build_models_fetch_url_for_client_version( &transport.provider.provider_type, &transport.endpoint.api_format, &transport.endpoint.base_url, + codex_client_version, ) .ok_or_else(|| "Rust models fetch does not support this provider format yet".to_string())?; return Ok(append_query_param(url, "key", &secret)); } - let (mut url, _) = build_models_fetch_url( + let (mut url, _) = build_models_fetch_url_for_client_version( &transport.provider.provider_type, &transport.endpoint.api_format, &transport.endpoint.base_url, + codex_client_version, ) .ok_or_else(|| "Rust models fetch does not support this provider format yet".to_string())?; @@ -727,9 +775,9 @@ mod tests { use super::{ build_antigravity_fetch_available_models_plan, build_antigravity_load_code_assist_plan, build_gemini_cli_load_code_assist_plan, build_kiro_list_available_models_plan, - build_models_fetch_execution_plan, build_standard_models_fetch_execution_plan, - build_vertex_models_fetch_execution_plan, ModelFetchTransportRuntime, - ANTIGRAVITY_REQUEST_USER_AGENT, + build_models_fetch_execution_plan, build_models_fetch_execution_plan_for_client_version, + build_standard_models_fetch_execution_plan, build_vertex_models_fetch_execution_plan, + ModelFetchTransportRuntime, ANTIGRAVITY_REQUEST_USER_AGENT, }; struct TestRuntime { @@ -922,7 +970,7 @@ mod tests { ), proxy: None, }; - let mut transport = sample_transport("codex", "openai:responses", "oauth"); + let mut transport = sample_transport("codex", "openai:chat", "oauth"); transport.endpoint.base_url = "https://chatgpt.com/backend-api/codex".to_string(); transport.endpoint.header_rules = Some(json!([ {"op": "set", "name": "chatgpt-account-id", "value": "spoofed-account"}, @@ -959,9 +1007,67 @@ mod tests { plan.headers.get("originator").map(String::as_str), Some("codex_cli_rs") ); + assert_eq!( + plan.headers.get("user-agent").map(String::as_str), + Some(aether_ai_formats::CODEX_CLIENT_USER_AGENT) + ); assert!(!plan.headers.contains_key("version")); } + #[tokio::test] + async fn builds_codex_models_fetch_plan_with_explicit_client_version() { + let runtime = TestRuntime { + oauth_auth: Some( + aether_provider_transport::LocalResolvedOAuthRequestAuth::Header { + name: "authorization".to_string(), + value: "Bearer access-token".to_string(), + }, + ), + proxy: None, + }; + let mut transport = sample_transport("codex", "openai:responses", "oauth"); + transport.endpoint.base_url = "https://chatgpt.com/backend-api/codex".to_string(); + transport.endpoint.header_rules = Some(json!([ + {"op": "set", "name": "user-agent", "value": "codex_cli_rs/0.1.0"}, + {"op": "remove", "name": "originator"} + ])); + transport.key.decrypted_auth_config = + Some(r#"{"account_id":"account-1","chatgpt_account_is_fedramp":true}"#.to_string()); + + let plan = build_models_fetch_execution_plan_for_client_version( + &runtime, + &transport, + Some("0.145.2"), + ) + .await + .expect("plan"); + + assert_eq!( + plan.url, + "https://chatgpt.com/backend-api/codex/models?client_version=0.145.2" + ); + assert_eq!( + plan.headers.get("user-agent").map(String::as_str), + Some("codex_cli_rs/0.145.2") + ); + assert_eq!( + plan.headers.get("originator").map(String::as_str), + Some("codex_cli_rs") + ); + assert_eq!( + plan.headers.get("originator").map(String::as_str), + Some("codex_cli_rs") + ); + assert_eq!( + plan.headers.get("chatgpt-account-id").map(String::as_str), + Some("account-1") + ); + assert_eq!( + plan.headers.get("x-openai-fedramp").map(String::as_str), + Some("true") + ); + } + #[tokio::test] async fn builds_claude_models_fetch_plan_with_pagination() { let runtime = TestRuntime {