From a39048eccea368a49e7a36dafb673bef55271046 Mon Sep 17 00:00:00 2001 From: elky Date: Tue, 1 Sep 2026 15:33:40 +0800 Subject: [PATCH] feat(codex): stabilize identity across retries --- .../src/ai_serving/codex_context.rs | 249 ++++++++ apps/aether-gateway/src/ai_serving/mod.rs | 1 + .../src/ai_serving/planner/decision_input.rs | 91 ++- .../planner/standard/family/payload.rs | 2 +- .../standard/openai/chat/decision/request.rs | 2 +- .../src/ai_serving/transport.rs | 1 + .../src/client_session_affinity.rs | 489 ++++++++++++++- apps/aether-gateway/src/handlers/proxy/mod.rs | 1 + .../handlers/proxy/websocket/live/planner.rs | 18 +- .../handlers/proxy/websocket/live/registry.rs | 5 + .../proxy/websocket/responses/client.rs | 23 +- .../proxy/websocket/responses/quota.rs | 9 +- .../proxy/websocket/responses/session.rs | 11 +- .../proxy/websocket/responses/turn_state.rs | 13 + .../src/formats/openai/responses/codex.rs | 142 ++++- .../src/provider/providers/generic.rs | 265 ++++++++- .../src/provider/providers/mod.rs | 3 +- .../transport/src/codex_fingerprint.rs | 556 +++++++++++++++++- crates/aether-provider/transport/src/lib.rs | 6 +- 19 files changed, 1731 insertions(+), 156 deletions(-) create mode 100644 apps/aether-gateway/src/ai_serving/codex_context.rs diff --git a/apps/aether-gateway/src/ai_serving/codex_context.rs b/apps/aether-gateway/src/ai_serving/codex_context.rs new file mode 100644 index 000000000..ac45689ac --- /dev/null +++ b/apps/aether-gateway/src/ai_serving/codex_context.rs @@ -0,0 +1,249 @@ +use std::sync::{Arc, OnceLock}; +use std::time::{SystemTime, UNIX_EPOCH}; + +use aether_provider_transport::CodexFingerprintConvergenceContext; +use http::{request::Parts, HeaderMap}; +use serde_json::Value; +use uuid::Uuid; + +use crate::client_session_affinity::codex_request_signals_from_request; + +#[derive(Debug, Clone)] +pub(crate) struct CodexFingerprintContextSlot(Arc>); + +impl Default for CodexFingerprintContextSlot { + fn default() -> Self { + Self(Arc::new(OnceLock::new())) + } +} + +impl CodexFingerprintContextSlot { + fn resolve( + &self, + headers: &HeaderMap, + body_json: &Value, + ) -> CodexFingerprintConvergenceContext { + self.0 + .get_or_init(|| { + build_codex_fingerprint_context(headers, body_json, Uuid::now_v7().to_string()) + }) + .clone() + } +} + +pub(crate) fn resolve_codex_fingerprint_context( + parts: &Parts, + body_json: &Value, +) -> CodexFingerprintConvergenceContext { + if let Some(context) = parts + .extensions + .get::() + .cloned() + { + return context; + } + if let Some(slot) = parts.extensions.get::() { + return slot.resolve(&parts.headers, body_json); + } + build_codex_fingerprint_context(&parts.headers, body_json, Uuid::now_v7().to_string()) +} + +pub(crate) fn install_codex_fingerprint_context_slot(parts: &mut Parts) { + if parts + .extensions + .get::() + .is_none() + && parts + .extensions + .get::() + .is_none() + { + parts + .extensions + .insert(CodexFingerprintContextSlot::default()); + } +} + +pub(crate) fn ensure_codex_fingerprint_context( + parts: &mut Parts, + body_json: &Value, +) -> CodexFingerprintConvergenceContext { + let context = resolve_codex_fingerprint_context(parts, body_json); + if parts + .extensions + .get::() + .is_none() + { + parts.extensions.remove::(); + parts.extensions.insert(context.clone()); + } + context +} + +pub(crate) fn attach_codex_logical_turn_context( + parts: &mut Parts, + body_json: &Value, + logical_turn_id: &str, +) -> CodexFingerprintConvergenceContext { + let context = + build_codex_fingerprint_context(&parts.headers, body_json, logical_turn_id.to_string()); + parts.extensions.remove::(); + parts.extensions.insert(context.clone()); + context +} + +pub(crate) fn restore_codex_logical_turn_context( + parts: &mut Parts, + context: &CodexFingerprintConvergenceContext, +) { + parts.extensions.remove::(); + parts.extensions.insert(context.clone()); +} + +fn build_codex_fingerprint_context( + headers: &HeaderMap, + body_json: &Value, + logical_turn_id: String, +) -> CodexFingerprintConvergenceContext { + let signals = codex_request_signals_from_request(headers, Some(body_json)); + let mut context = + CodexFingerprintConvergenceContext::new(logical_turn_id, current_unix_millis()); + + if let Some(turn_id) = signals.turn_id { + context = context.with_original_turn_id(turn_id); + } + if let Some(session_id) = signals.thread_id.or(signals.session_id) { + context = context.with_original_client_session_id(session_id); + } + if let Some(prompt_cache_key) = signals.prompt_cache_key { + context = context.with_original_prompt_cache_key(prompt_cache_key); + } + + context +} + +fn current_unix_millis() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_millis() + .try_into() + .unwrap_or(u64::MAX) +} + +#[cfg(test)] +mod tests { + use http::HeaderValue; + use serde_json::json; + + use super::*; + + #[test] + fn request_signals_are_captured_once_for_the_logical_turn() { + let request = http::Request::builder() + .header("thread-id", "header-thread") + .body(()) + .expect("request should build"); + let (mut parts, _) = request.into_parts(); + let body = json!({ + "prompt_cache_key": "client-cache", + "client_metadata": { + "turn_id": "client-turn", + "thread_id": "body-thread" + } + }); + + let context = attach_codex_logical_turn_context(&mut parts, &body, "logical-turn"); + + assert_eq!(context.logical_turn_id(), "logical-turn"); + assert_eq!(context.original_turn_id(), Some("client-turn")); + assert_eq!(context.original_client_session_id(), Some("header-thread")); + assert_eq!(context.original_prompt_cache_key(), Some("client-cache")); + assert_eq!( + parts.extensions.get::(), + Some(&context) + ); + } + + #[test] + fn restored_context_wins_over_retry_request_signals() { + let original = CodexFingerprintConvergenceContext::new("logical-turn", 1234) + .with_original_turn_id("original-turn") + .with_original_client_session_id("original-thread") + .with_original_prompt_cache_key("original-cache"); + let request = http::Request::builder() + .body(()) + .expect("request should build"); + let (mut parts, _) = request.into_parts(); + parts + .headers + .insert("thread-id", HeaderValue::from_static("retry-thread")); + restore_codex_logical_turn_context(&mut parts, &original); + + let resolved = resolve_codex_fingerprint_context( + &parts, + &json!({ + "prompt_cache_key": "retry-cache", + "client_metadata": {"turn_id": "retry-turn"} + }), + ); + + assert_eq!(resolved, original); + assert_eq!(resolved.turn_started_at_unix_ms(), 1234); + } + + #[test] + fn generated_context_is_persisted_for_http_replanning() { + let request = http::Request::builder() + .header("session-id", "client-session") + .body(()) + .expect("request should build"); + let (mut parts, _) = request.into_parts(); + let body = json!({ + "prompt_cache_key": "client-cache", + "client_metadata": {"turn_id": "client-turn"} + }); + + let first = ensure_codex_fingerprint_context(&mut parts, &body); + let second = resolve_codex_fingerprint_context( + &parts, + &json!({ + "prompt_cache_key": "retry-cache", + "client_metadata": {"turn_id": "retry-turn"} + }), + ); + + assert_eq!(second, first); + assert_eq!(second.original_turn_id(), Some("client-turn")); + assert_eq!(second.original_prompt_cache_key(), Some("client-cache")); + } + + #[test] + fn installed_slot_reuses_context_across_cloned_parts() { + let request = http::Request::builder() + .body(()) + .expect("request should build"); + let (mut parts, _) = request.into_parts(); + install_codex_fingerprint_context_slot(&mut parts); + let cloned_parts = parts.clone(); + + let first = resolve_codex_fingerprint_context( + &parts, + &json!({ + "prompt_cache_key": "first-cache", + "client_metadata": {"turn_id": "first-turn"} + }), + ); + let second = resolve_codex_fingerprint_context( + &cloned_parts, + &json!({ + "prompt_cache_key": "second-cache", + "client_metadata": {"turn_id": "second-turn"} + }), + ); + + assert_eq!(second, first); + assert_eq!(second.original_turn_id(), Some("first-turn")); + assert_eq!(second.original_prompt_cache_key(), Some("first-cache")); + } +} diff --git a/apps/aether-gateway/src/ai_serving/mod.rs b/apps/aether-gateway/src/ai_serving/mod.rs index cca447e05..8462645fd 100644 --- a/apps/aether-gateway/src/ai_serving/mod.rs +++ b/apps/aether-gateway/src/ai_serving/mod.rs @@ -1,5 +1,6 @@ mod adaptation; pub(crate) mod api; +pub(crate) mod codex_context; mod finalize; mod planner; mod pure; 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 93fbdf708..d16b97399 100644 --- a/apps/aether-gateway/src/ai_serving/planner/decision_input.rs +++ b/apps/aether-gateway/src/ai_serving/planner/decision_input.rs @@ -2,6 +2,7 @@ use std::collections::BTreeMap; use std::time::Duration; use aether_ai_serving::{run_ai_authenticated_decision_input, AiAuthenticatedDecisionInputPort}; +use aether_provider_transport::CodexFingerprintConvergenceContext; use aether_routing_core::{ rank_vector_for_candidate, CandidateKind, ResolvedRoutingPolicy, RoutingCandidateFacts, RoutingCandidateTrace, RoutingDecisionTrace, RoutingPoolExpansionTrace, RoutingRulePhase, @@ -55,7 +56,7 @@ pub(crate) struct LocalRequestedModelDecisionInput { pub(crate) client_surface: Option, pub(crate) gateway_credential_carrier: Option, pub(crate) client_session_affinity: Option, - pub(crate) original_client_session_id: Option, + pub(crate) codex_fingerprint_context: Option, pub(crate) routing_policy: Option, pub(crate) routing_trace_seed: Option, pub(crate) routing_context: Option, @@ -376,13 +377,25 @@ fn apply_codex_oauth_fingerprint_convergence_to_decision( else { return; }; - crate::ai_serving::transport::apply_codex_oauth_fingerprint_convergence( - transport, - provider_api_format, - input.original_client_session_id.as_deref(), - &mut decision.provider_request_headers, - provider_request_body, - ); + let Some(context) = input.codex_fingerprint_context.as_ref() else { + return; + }; + let applied = + crate::ai_serving::transport::apply_codex_oauth_fingerprint_convergence_with_context( + transport, + provider_api_format, + context, + &mut decision.provider_request_headers, + provider_request_body, + ); + if applied { + decision.prompt_cache_key = provider_request_body + .get("prompt_cache_key") + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned); + } } struct GatewayAuthenticatedDecisionInputPort<'a> { @@ -471,7 +484,7 @@ pub(crate) fn build_local_requested_model_decision_input( client_surface: None, gateway_credential_carrier: None, client_session_affinity: None, - original_client_session_id: None, + codex_fingerprint_context: None, routing_policy: None, routing_trace_seed: None, routing_context: None, @@ -486,7 +499,8 @@ pub(crate) async fn attach_routing_policy_to_local_requested_model_input( body_json: &Value, client_api_format: &str, ) -> Result<(), GatewayError> { - input.original_client_session_id = original_client_session_id_from_headers(&parts.headers); + input.codex_fingerprint_context = + Some(crate::ai_serving::codex_context::resolve_codex_fingerprint_context(parts, body_json)); let explicit_group = routing_header_value_str(&parts.headers, ROUTING_GROUP_HEADER); let selected_group = match state.routing_group_read_repository() { Some(repository) => { @@ -737,12 +751,6 @@ pub(crate) async fn attach_routing_policy_to_local_requested_model_input( Ok(()) } -fn original_client_session_id_from_headers(headers: &HeaderMap) -> Option { - routing_header_value_str(headers, "session-id") - .or_else(|| routing_header_value_str(headers, "session_id")) - .or_else(|| routing_header_value_str(headers, "x-session-id")) -} - fn try_attach_static_default_routing_policy_to_input( input: &mut LocalRequestedModelDecisionInput, parts: &http::request::Parts, @@ -1106,38 +1114,6 @@ mod tests { GatewayProviderTransportProvider, }; - #[test] - fn original_client_session_id_accepts_live_header_as_fallback() { - let headers = HeaderMap::from_iter([( - HeaderName::from_static("x-session-id"), - HeaderValue::from_static("live-thread-1"), - )]); - - assert_eq!( - original_client_session_id_from_headers(&headers).as_deref(), - Some("live-thread-1") - ); - } - - #[test] - fn original_client_session_id_prefers_responses_headers_over_live_fallback() { - let headers = HeaderMap::from_iter([ - ( - HeaderName::from_static("session-id"), - HeaderValue::from_static("responses-session"), - ), - ( - HeaderName::from_static("x-session-id"), - HeaderValue::from_static("live-thread"), - ), - ]); - - assert_eq!( - original_client_session_id_from_headers(&headers).as_deref(), - Some("responses-session") - ); - } - #[test] fn explicit_routing_selection_cache_key_is_principal_specific() { let first = routing_group_selection_cache_key( @@ -1350,7 +1326,7 @@ mod tests { client_surface: None, gateway_credential_carrier: None, client_session_affinity: None, - original_client_session_id: None, + codex_fingerprint_context: None, routing_policy: None, routing_trace_seed: None, model_directive_policy: Default::default(), @@ -1599,7 +1575,7 @@ mod tests { client_surface: None, gateway_credential_carrier: None, client_session_affinity: None, - original_client_session_id: None, + codex_fingerprint_context: None, routing_policy: None, routing_trace_seed: None, model_directive_policy: Default::default(), @@ -1669,7 +1645,7 @@ mod tests { client_surface: None, gateway_credential_carrier: None, client_session_affinity: None, - original_client_session_id: None, + codex_fingerprint_context: None, routing_policy: None, routing_trace_seed: None, routing_context: None, @@ -1754,7 +1730,13 @@ mod tests { }); let mut with_mutation = sample_decision_input(); for input in [&mut no_context, &mut empty_mutation, &mut with_mutation] { - input.original_client_session_id = Some("client-session-1".to_string()); + input.codex_fingerprint_context = Some( + CodexFingerprintConvergenceContext::new( + uuid::Uuid::new_v4().to_string(), + 1_756_668_000_000, + ) + .with_original_client_session_id("client-session-1".to_string()), + ); } let mut stable_identity = None; @@ -1801,6 +1783,10 @@ mod tests { .provider_request_body .as_ref() .expect("request body"); + assert_eq!( + decision.prompt_cache_key.as_deref(), + body.get("prompt_cache_key").and_then(Value::as_str) + ); assert_eq!( body["prompt_cache_key"], "172c39e6-c0a0-5a70-8b63-e0f8e0d185a3" @@ -2005,6 +1991,7 @@ mod tests { let body = decision.provider_request_body.as_ref().expect("body"); assert!(body.get("prompt_cache_key").is_none()); + assert!(decision.prompt_cache_key.is_none()); assert!(body.get("client_metadata").is_none()); assert!(!decision.provider_request_headers.contains_key("session-id")); assert!(!decision.provider_request_headers.contains_key("thread-id")); diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/family/payload.rs b/apps/aether-gateway/src/ai_serving/planner/standard/family/payload.rs index d213404a9..a025c1349 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/family/payload.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/family/payload.rs @@ -377,7 +377,7 @@ mod tests { client_surface: None, gateway_credential_carrier: None, client_session_affinity: None, - original_client_session_id: None, + codex_fingerprint_context: None, routing_policy: None, routing_trace_seed: None, routing_context: None, diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs index 599a1e763..066eb8e0a 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs @@ -2182,7 +2182,7 @@ mod tests { client_surface: None, gateway_credential_carrier: None, client_session_affinity: None, - original_client_session_id: None, + codex_fingerprint_context: None, routing_policy: None, routing_trace_seed: None, routing_context: None, diff --git a/apps/aether-gateway/src/ai_serving/transport.rs b/apps/aether-gateway/src/ai_serving/transport.rs index 42af31e02..6883df7bb 100644 --- a/apps/aether-gateway/src/ai_serving/transport.rs +++ b/apps/aether-gateway/src/ai_serving/transport.rs @@ -60,6 +60,7 @@ pub(crate) mod windsurf { pub(crate) use aether_provider_transport::{ append_transport_diagnostics_to_value, apply_codex_oauth_fingerprint_convergence, + apply_codex_oauth_fingerprint_convergence_with_context, 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, diff --git a/apps/aether-gateway/src/client_session_affinity.rs b/apps/aether-gateway/src/client_session_affinity.rs index 66e009184..e74cb1e50 100644 --- a/apps/aether-gateway/src/client_session_affinity.rs +++ b/apps/aether-gateway/src/client_session_affinity.rs @@ -23,6 +23,21 @@ pub(crate) struct ClientSessionScope { pub(crate) source: ClientSessionSignalSource, } +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub(crate) struct CodexRequestSignals { + pub(crate) session_id: Option, + pub(crate) thread_id: Option, + pub(crate) turn_id: Option, + pub(crate) prompt_cache_key: Option, +} + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +struct CodexTurnMetadataSignals { + session_id: Option, + thread_id: Option, + turn_id: Option, +} + impl ClientSessionScope { fn new( client_family: impl Into, @@ -136,6 +151,13 @@ pub(crate) fn client_session_scope_from_request( .or_else(|| extract_scope_from_other_specific_adapters(&request, client_family.as_str())) } +pub(crate) fn codex_request_signals_from_request( + headers: &http::HeaderMap, + body_json: Option<&Value>, +) -> CodexRequestSignals { + extract_codex_request_signals(&ClientSessionRequest { headers, body_json }) +} + fn codex_search_session_scope(request: &ClientSessionRequest<'_>) -> Option { let session_id = request .body_json? @@ -408,29 +430,7 @@ impl ClientSessionScopeAdapter for CodexSessionScopeAdapter { } fn extract_scope(&self, request: &ClientSessionRequest<'_>) -> Option { - header_value_str(request.headers, "session-id") - .or_else(|| header_value_str(request.headers, "thread-id")) - .or_else(|| header_value_str(request.headers, "session_id")) - .or_else(|| header_value_str(request.headers, "conversation_id")) - .map(|root_session| { - ClientSessionScope::new( - self.family(), - root_session, - None, - header_value_str(request.headers, "chatgpt-account-id"), - ClientSessionSignalSource::Header, - ) - }) - .or_else(|| { - let body_session = GenericSessionScopeAdapter.extract_scope(request)?; - Some(ClientSessionScope::new( - self.family(), - body_session.session_id, - body_session.agent_id, - header_value_str(request.headers, "chatgpt-account-id"), - body_session.source, - )) - }) + codex_request_session_scope_from_request(request) } } @@ -778,6 +778,162 @@ fn explicit_aether_session_scope( )) } +fn extract_codex_request_signals(request: &ClientSessionRequest<'_>) -> CodexRequestSignals { + let body_client_metadata = request + .body_json + .and_then(|body| body.get("client_metadata")) + .and_then(Value::as_object); + let body_turn_metadata = codex_turn_metadata_signals( + body_client_metadata.and_then(|metadata| metadata.get("x-codex-turn-metadata")), + ); + let header_turn_metadata = header_value_str(request.headers, "x-codex-turn-metadata") + .map(|raw| parse_codex_turn_metadata(&raw)) + .unwrap_or_default(); + + let native_thread_id = header_value_str(request.headers, "thread-id") + .or_else(|| { + body_client_metadata + .and_then(|metadata| value_at_map_path(metadata, "thread_id")) + .map(ToOwned::to_owned) + }) + .or_else(|| body_turn_metadata.thread_id.clone()); + let turn_id = body_client_metadata + .and_then(|metadata| value_at_map_path(metadata, "turn_id")) + .map(ToOwned::to_owned) + .or_else(|| body_turn_metadata.turn_id.clone()) + .or_else(|| { + request + .body_json + .and_then(|body| value_at_path(body, &["turn_id"])) + .map(ToOwned::to_owned) + }) + .or(header_turn_metadata.turn_id); + let prompt_cache_key = request + .body_json + .and_then(|body| value_at_path(body, &["prompt_cache_key"])) + .map(ToOwned::to_owned); + let session_id = + codex_request_session_scope(request, &body_turn_metadata).map(|scope| scope.session_id); + let thread_id = native_thread_id.or_else(|| session_id.clone()); + + CodexRequestSignals { + session_id, + thread_id, + turn_id, + prompt_cache_key, + } +} + +fn codex_request_session_scope_from_request( + request: &ClientSessionRequest<'_>, +) -> Option { + let body_turn_metadata = codex_turn_metadata_signals( + request + .body_json + .and_then(|body| body.get("client_metadata")) + .and_then(Value::as_object) + .and_then(|metadata| metadata.get("x-codex-turn-metadata")), + ); + codex_request_session_scope(request, &body_turn_metadata) +} + +fn codex_request_session_scope( + request: &ClientSessionRequest<'_>, + body_turn_metadata: &CodexTurnMetadataSignals, +) -> Option { + if let Some(scope) = explicit_aether_session_scope(request, CodexSessionScopeAdapter.family()) { + return Some(scope); + } + + if let Some(root_session) = header_value_str(request.headers, "session-id") + .or_else(|| header_value_str(request.headers, "thread-id")) + .or_else(|| header_value_str(request.headers, "session_id")) + .or_else(|| header_value_str(request.headers, "conversation_id")) + .or_else(|| header_value_str(request.headers, "x-session-id")) + { + return Some(codex_session_scope( + request, + root_session, + None, + ClientSessionSignalSource::Header, + )); + } + + let body_client_metadata = request + .body_json + .and_then(|body| body.get("client_metadata")) + .and_then(Value::as_object); + if let Some(root_session) = body_client_metadata + .and_then(|metadata| value_at_map_path(metadata, "session_id")) + .or_else(|| { + body_client_metadata.and_then(|metadata| value_at_map_path(metadata, "thread_id")) + }) + .map(ToOwned::to_owned) + .or_else(|| body_turn_metadata.session_id.clone()) + .or_else(|| body_turn_metadata.thread_id.clone()) + { + return Some(codex_session_scope( + request, + root_session, + None, + ClientSessionSignalSource::Body, + )); + } + + let generic = GenericSessionScopeAdapter.extract_scope(request)?; + Some(codex_session_scope( + request, + generic.session_id, + generic.agent_id, + generic.source, + )) +} + +fn codex_session_scope( + request: &ClientSessionRequest<'_>, + session_id: String, + agent_id: Option, + source: ClientSessionSignalSource, +) -> ClientSessionScope { + ClientSessionScope::new( + CodexSessionScopeAdapter.family(), + session_id, + agent_id, + header_value_str(request.headers, "chatgpt-account-id"), + source, + ) +} + +fn codex_turn_metadata_signals(value: Option<&Value>) -> CodexTurnMetadataSignals { + match value { + Some(Value::Object(metadata)) => codex_turn_metadata_signals_from_map(metadata), + Some(Value::String(raw)) => parse_codex_turn_metadata(raw), + _ => CodexTurnMetadataSignals::default(), + } +} + +fn parse_codex_turn_metadata(raw: &str) -> CodexTurnMetadataSignals { + serde_json::from_str::>(raw) + .map(|metadata| codex_turn_metadata_signals_from_map(&metadata)) + .unwrap_or_default() +} + +fn codex_turn_metadata_signals_from_map(metadata: &Map) -> CodexTurnMetadataSignals { + CodexTurnMetadataSignals { + session_id: value_at_map_path(metadata, "session_id").map(ToOwned::to_owned), + thread_id: value_at_map_path(metadata, "thread_id").map(ToOwned::to_owned), + turn_id: value_at_map_path(metadata, "turn_id").map(ToOwned::to_owned), + } +} + +fn value_at_map_path<'a>(object: &'a Map, key: &str) -> Option<&'a str> { + object + .get(key) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) +} + fn normalize_session_key( account_hint: Option<&str>, root_session: &str, @@ -831,12 +987,25 @@ mod tests { client_session_affinity_from_api_request, client_session_affinity_from_report_context_value, client_session_affinity_from_request, client_session_affinity_report_context_value, client_session_scope_from_request, - ClientSessionSignalSource, AETHER_AGENT_ID_HEADER, AETHER_SESSION_ID_HEADER, + codex_request_signals_from_request, ClientSessionSignalSource, AETHER_AGENT_ID_HEADER, + AETHER_SESSION_ID_HEADER, }; use aether_scheduler_core::ClientSessionAffinity; - use http::{HeaderMap, HeaderValue}; + use http::{HeaderMap, HeaderName, HeaderValue}; use serde_json::json; + fn request_headers(values: &[(&str, &str)]) -> HeaderMap { + values + .iter() + .map(|(name, value)| { + ( + HeaderName::from_bytes(name.as_bytes()).expect("valid test header name"), + HeaderValue::from_bytes(value.as_bytes()).expect("valid test header value"), + ) + }) + .collect() + } + #[test] fn unknown_adapter_extracts_body_session_and_agent() { let body = json!({ @@ -933,6 +1102,276 @@ mod tests { ); } + #[test] + fn codex_request_signals_apply_session_precedence() { + let cases = vec![ + ( + request_headers(&[ + (AETHER_SESSION_ID_HEADER, "aether-session"), + ("session-id", "header-session"), + ]), + json!({"client_metadata": {"session_id": "body-session"}}), + "aether-session", + ClientSessionSignalSource::ExplicitAetherHeader, + ), + ( + request_headers(&[ + ("session-id", "header-session"), + ("thread-id", "header-thread"), + ("session_id", "header-session-underscore"), + ("conversation_id", "header-conversation"), + ]), + json!({"client_metadata": {"session_id": "body-session"}}), + "header-session", + ClientSessionSignalSource::Header, + ), + ( + request_headers(&[ + ("thread-id", "header-thread"), + ("session_id", "header-session-underscore"), + ("conversation_id", "header-conversation"), + ]), + json!({"client_metadata": {"session_id": "body-session"}}), + "header-thread", + ClientSessionSignalSource::Header, + ), + ( + request_headers(&[ + ("session_id", "header-session-underscore"), + ("conversation_id", "header-conversation"), + ]), + json!({"client_metadata": {"session_id": "body-session"}}), + "header-session-underscore", + ClientSessionSignalSource::Header, + ), + ( + request_headers(&[("conversation_id", "header-conversation")]), + json!({"client_metadata": {"session_id": "body-session"}}), + "header-conversation", + ClientSessionSignalSource::Header, + ), + ( + HeaderMap::new(), + json!({ + "prompt_cache_key": "prompt-cache", + "client_metadata": { + "session_id": "body-session", + "thread_id": "body-thread", + "x-codex-turn-metadata": { + "session_id": "nested-session", + "thread_id": "nested-thread" + } + } + }), + "body-session", + ClientSessionSignalSource::Body, + ), + ( + HeaderMap::new(), + json!({ + "prompt_cache_key": "prompt-cache", + "client_metadata": { + "thread_id": "body-thread", + "x-codex-turn-metadata": {"session_id": "nested-session"} + } + }), + "body-thread", + ClientSessionSignalSource::Body, + ), + ( + HeaderMap::new(), + json!({ + "prompt_cache_key": "prompt-cache", + "client_metadata": { + "x-codex-turn-metadata": { + "session_id": "nested-session", + "thread_id": "nested-thread" + } + } + }), + "nested-session", + ClientSessionSignalSource::Body, + ), + ( + HeaderMap::new(), + json!({ + "prompt_cache_key": "prompt-cache", + "client_metadata": { + "x-codex-turn-metadata": json!({ + "thread_id": "nested-thread" + }).to_string() + } + }), + "nested-thread", + ClientSessionSignalSource::Body, + ), + ( + HeaderMap::new(), + json!({ + "prompt_cache_key": "prompt-cache", + "conversation_id": "generic-conversation" + }), + "prompt-cache", + ClientSessionSignalSource::Body, + ), + ( + HeaderMap::new(), + json!({"metadata": {"session_id": "generic-session"}}), + "generic-session", + ClientSessionSignalSource::Body, + ), + ]; + + for (headers, body, expected_session_id, expected_source) in cases { + let signals = codex_request_signals_from_request(&headers, Some(&body)); + assert_eq!(signals.session_id.as_deref(), Some(expected_session_id)); + + let mut codex_headers = headers; + codex_headers.insert( + http::header::USER_AGENT, + HeaderValue::from_static("codex_cli_rs/0.144.1"), + ); + let scope = client_session_scope_from_request(&codex_headers, Some(&body)) + .expect("Codex scope should reuse the native signal precedence"); + assert_eq!(scope.client_family, "codex"); + assert_eq!(scope.session_id, expected_session_id); + assert_eq!(scope.source, expected_source); + } + } + + #[test] + fn codex_request_signals_extract_thread_and_prompt_cache_independently() { + let body = json!({ + "prompt_cache_key": "prompt-cache", + "client_metadata": { + "thread_id": "body-thread", + "x-codex-turn-metadata": {"thread_id": "nested-thread"} + } + }); + let headers = request_headers(&[("thread-id", "header-thread")]); + let header_signals = codex_request_signals_from_request(&headers, Some(&body)); + assert_eq!(header_signals.thread_id.as_deref(), Some("header-thread")); + assert_eq!( + header_signals.prompt_cache_key.as_deref(), + Some("prompt-cache") + ); + + let body_signals = codex_request_signals_from_request(&HeaderMap::new(), Some(&body)); + assert_eq!(body_signals.thread_id.as_deref(), Some("body-thread")); + + let nested_body = json!({ + "client_metadata": { + "x-codex-turn-metadata": json!({ + "thread_id": "nested-thread" + }).to_string() + } + }); + let nested_signals = + codex_request_signals_from_request(&HeaderMap::new(), Some(&nested_body)); + assert_eq!(nested_signals.thread_id.as_deref(), Some("nested-thread")); + + let session_only_body = json!({"client_metadata": {"session_id": "body-session"}}); + let session_only_signals = + codex_request_signals_from_request(&HeaderMap::new(), Some(&session_only_body)); + assert_eq!( + session_only_signals.thread_id.as_deref(), + Some("body-session") + ); + } + + #[test] + fn codex_request_signals_use_live_session_header() { + let headers = request_headers(&[("x-session-id", "live-session")]); + let signals = codex_request_signals_from_request(&headers, None); + + assert_eq!(signals.session_id.as_deref(), Some("live-session")); + assert_eq!(signals.thread_id.as_deref(), Some("live-session")); + } + + #[test] + fn codex_request_signals_prefer_responses_session_header_over_live_session_header() { + let headers = request_headers(&[ + ("session-id", "responses-session"), + ("x-session-id", "live-session"), + ]); + let signals = codex_request_signals_from_request(&headers, None); + + assert_eq!(signals.session_id.as_deref(), Some("responses-session")); + assert_eq!(signals.thread_id.as_deref(), Some("responses-session")); + } + + #[test] + fn codex_request_signals_apply_turn_precedence() { + let headers = request_headers(&[("x-codex-turn-metadata", r#"{"turn_id":"header-turn"}"#)]); + let direct_body = json!({ + "turn_id": "top-level-turn", + "client_metadata": { + "turn_id": "body-turn", + "x-codex-turn-metadata": {"turn_id": "nested-turn"} + } + }); + assert_eq!( + codex_request_signals_from_request(&headers, Some(&direct_body)) + .turn_id + .as_deref(), + Some("body-turn") + ); + + let nested_object_body = json!({ + "turn_id": "top-level-turn", + "client_metadata": { + "x-codex-turn-metadata": {"turn_id": "nested-object-turn"} + } + }); + assert_eq!( + codex_request_signals_from_request(&headers, Some(&nested_object_body)) + .turn_id + .as_deref(), + Some("nested-object-turn") + ); + + let nested_string_body = json!({ + "turn_id": "top-level-turn", + "client_metadata": { + "x-codex-turn-metadata": json!({ + "turn_id": "nested-string-turn" + }).to_string() + } + }); + assert_eq!( + codex_request_signals_from_request(&headers, Some(&nested_string_body)) + .turn_id + .as_deref(), + Some("nested-string-turn") + ); + + let top_level_body = json!({ + "turn_id": "top-level-turn", + "client_metadata": {"x-codex-turn-metadata": "not-json"} + }); + assert_eq!( + codex_request_signals_from_request(&headers, Some(&top_level_body)) + .turn_id + .as_deref(), + Some("top-level-turn") + ); + assert_eq!( + codex_request_signals_from_request(&headers, None) + .turn_id + .as_deref(), + Some("header-turn") + ); + } + + #[test] + fn codex_request_signals_ignore_client_request_id() { + let headers = request_headers(&[("x-client-request-id", "request-only-id")]); + let signals = + codex_request_signals_from_request(&headers, Some(&json!({"model": "gpt-5"}))); + + assert_eq!(signals, super::CodexRequestSignals::default()); + } + #[test] fn report_context_round_trips_normalized_session_affinity() { let affinity = ClientSessionAffinity::new( diff --git a/apps/aether-gateway/src/handlers/proxy/mod.rs b/apps/aether-gateway/src/handlers/proxy/mod.rs index 0fe8373dc..00b17feac 100644 --- a/apps/aether-gateway/src/handlers/proxy/mod.rs +++ b/apps/aether-gateway/src/handlers/proxy/mod.rs @@ -1092,6 +1092,7 @@ async fn proxy_request_inner( ), } let (mut parts, body) = request.into_parts(); + crate::ai_serving::codex_context::install_codex_fingerprint_context_slot(&mut parts); let redaction_slot = crate::privacy::RedactionSessionSlot::default(); parts.extensions.insert(redaction_slot.clone()); parts diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/live/planner.rs b/apps/aether-gateway/src/handlers/proxy/websocket/live/planner.rs index 8a8b9fb6a..83c951879 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/live/planner.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/live/planner.rs @@ -49,6 +49,8 @@ pub(super) enum LiveAuthMode { pub(super) struct PlannedLiveCandidate { pub(super) execution: AiExecutionDecision, pub(super) pinned_candidate: ResponsesWebSocketPinnedCandidate, + pub(super) codex_fingerprint_context: + aether_provider_transport::CodexFingerprintConvergenceContext, pub(super) client_model: String, pub(super) provider_model: String, pub(super) auth_mode: LiveAuthMode, @@ -228,8 +230,9 @@ async fn plan_live_candidate_inner( if validate_model(client_model).is_err() || client_model.len() > MAX_LIVE_MODEL_BYTES { return Ok(None); } - let parts = build_live_planning_parts(headers, remote_addr); + let mut parts = build_live_planning_parts(headers, remote_addr); let body = json!({"model": client_model, "input": []}); + crate::ai_serving::codex_context::install_codex_fingerprint_context_slot(&mut parts); let execution = maybe_build_pinned_stream_local_same_format_provider_decision_payload( state, &parts, @@ -338,6 +341,8 @@ async fn plan_live_candidate_inner( Ok(Some(PlannedLiveCandidate { execution, pinned_candidate, + codex_fingerprint_context: + crate::ai_serving::codex_context::resolve_codex_fingerprint_context(&parts, &body), client_model: client_model.to_string(), provider_model, auth_mode, @@ -560,7 +565,11 @@ pub(super) fn build_live_stream_admission_attempt( remote_addr: &SocketAddr, upstream_url: String, ) -> Result, GatewayError> { - let parts = build_live_planning_parts(headers, remote_addr); + let mut parts = build_live_planning_parts(headers, remote_addr); + crate::ai_serving::codex_context::restore_codex_logical_turn_context( + &mut parts, + &candidate.codex_fingerprint_context, + ); let body = json!({"model": candidate.client_model.as_str(), "input": []}); let mut execution = candidate.execution.clone(); execution.upstream_url = Some(upstream_url); @@ -922,6 +931,11 @@ mod tests { "key-1", ) .unwrap(), + codex_fingerprint_context: + aether_provider_transport::CodexFingerprintConvergenceContext::new( + "test-live-turn", + 1, + ), client_model: "global-model".to_string(), provider_model: "provider-model".to_string(), auth_mode, diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/live/registry.rs b/apps/aether-gateway/src/handlers/proxy/websocket/live/registry.rs index 2f738b9e9..b4bb0cc53 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/live/registry.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/live/registry.rs @@ -718,6 +718,11 @@ mod tests { PlannedLiveCandidate { execution, pinned_candidate: binding.pinned_candidate.clone(), + codex_fingerprint_context: + aether_provider_transport::CodexFingerprintConvergenceContext::new( + "test-live-turn", + 1, + ), client_model: binding.client_model.clone(), provider_model: binding.provider_model.clone(), auth_mode: binding.auth_mode, diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/responses/client.rs b/apps/aether-gateway/src/handlers/proxy/websocket/responses/client.rs index a7a978a6a..96d11c344 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/responses/client.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/responses/client.rs @@ -1,5 +1,6 @@ //! Client-side Responses WebSocket event forwarding and follow-up planning. +use aether_provider_transport::CodexFingerprintConvergenceContext; use axum::extract::ws::{Message as AxumWsMessage, WebSocket}; use futures_util::SinkExt; use serde_json::Value; @@ -248,7 +249,14 @@ pub(super) async fn forward_client_message( // derive one strong live control snapshot that every stage below // shares. The connection's Upgrade-time decision is only the // immutable identity seed. - let planning_parts = build_planning_parts(context); + let logical_turn_id = Uuid::now_v7().to_string(); + let mut planning_parts = build_planning_parts(context); + let codex_fingerprint_context = + crate::ai_serving::codex_context::attach_codex_logical_turn_context( + &mut planning_parts, + &client_event, + &logical_turn_id, + ); let turn_control = match resolve_responses_websocket_turn_control( state, context, @@ -453,6 +461,8 @@ pub(super) async fn forward_client_message( context, planning_parts, client_event, + logical_turn_id, + codex_fingerprint_context, turn_control, turn_redaction_session, ) @@ -466,6 +476,8 @@ pub(super) async fn forward_client_message( planning_parts, client_event, requested_model, + logical_turn_id, + codex_fingerprint_context, turn_control, raw_responses_lite_static_config .expect("independent turns always retain their raw static config"), @@ -513,6 +525,8 @@ async fn forward_pinned_continuation( context: &WebSocketRequestContext, planning_parts: http::request::Parts, client_event: Value, + logical_turn_id: String, + codex_fingerprint_context: CodexFingerprintConvergenceContext, turn_control: ResponsesWebSocketTurnControl, turn_redaction_session: Option, ) -> RelayDisposition { @@ -556,7 +570,6 @@ async fn forward_pinned_continuation( }; let turn_request_id = Uuid::new_v4().to_string(); - let logical_turn_id = Uuid::new_v4().to_string(); let planned = match await_owned_responses_websocket_plan(spawn_owned_responses_websocket_plan( state.clone(), planning_parts, @@ -766,6 +779,7 @@ async fn forward_pinned_continuation( bound.body_normalization = normalization; bound.turn_state.begin( LogicalTurn::new(client_event, turn_index, logical_turn_id) + .with_codex_fingerprint_context(codex_fingerprint_context) .with_provider_store(provider_event.get("store") == Some(&Value::Bool(true))) .with_turn_control(turn_control), turn, @@ -800,12 +814,13 @@ async fn forward_replanned_response_create( planning_parts: http::request::Parts, client_event: Value, requested_model: String, + logical_turn_id: String, + codex_fingerprint_context: CodexFingerprintConvergenceContext, turn_control: ResponsesWebSocketTurnControl, raw_responses_lite_static_config: ResponsesLiteStaticConfig, turn_redaction_session: Option, ) -> RelayDisposition { let turn_request_id = Uuid::new_v4().to_string(); - let logical_turn_id = Uuid::new_v4().to_string(); let now_unix_secs = current_unix_secs(); let excluded_key_ids = bound.exhausted_exclusions.key_ids(now_unix_secs); let excluded_codex_account_ids = bound.exhausted_exclusions.codex_account_ids(now_unix_secs); @@ -992,6 +1007,7 @@ async fn forward_replanned_response_create( bound.body_normalization = normalization; bound.turn_state.begin( LogicalTurn::new(client_event.clone(), turn_index, logical_turn_id.clone()) + .with_codex_fingerprint_context(codex_fingerprint_context.clone()) .with_provider_store(provider_event.get("store") == Some(&Value::Bool(true))) .with_turn_control(turn_control), turn, @@ -1076,6 +1092,7 @@ async fn forward_replanned_response_create( bound.binding_identity = replacement.binding_identity; bound.turn_state.begin( LogicalTurn::new(client_event, turn_index, logical_turn_id) + .with_codex_fingerprint_context(codex_fingerprint_context) .with_provider_store(provider_event.get("store") == Some(&Value::Bool(true))) .with_turn_control(turn_control), turn, diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/responses/quota.rs b/apps/aether-gateway/src/handlers/proxy/websocket/responses/quota.rs index 4ad15bcda..116058015 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/responses/quota.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/responses/quota.rs @@ -146,6 +146,7 @@ pub(super) async fn retry_active_turn_after_quota_exhaustion( }; let turn_index = active.turn_index; let logical_turn_id = active.logical_turn_id.clone(); + let codex_fingerprint_context = active.codex_fingerprint_context.clone(); let turn_attempt = active.turn_attempt; let retry_exclusion_until_unix_secs = bound @@ -154,7 +155,13 @@ pub(super) async fn retry_active_turn_after_quota_exhaustion( let exhausted_key = record_exhausted_bound_key(bound, retry_exclusion_until_unix_secs); let exhausted_key_id = exhausted_key.as_ref().map(|(key_id, _)| key_id.clone()); - let planning_parts = build_planning_parts(context); + let mut planning_parts = build_planning_parts(context); + if let Some(codex_fingerprint_context) = codex_fingerprint_context.as_ref() { + crate::ai_serving::codex_context::restore_codex_logical_turn_context( + &mut planning_parts, + codex_fingerprint_context, + ); + } let turn_request_id = Uuid::new_v4().to_string(); let now_unix_secs = current_unix_secs(); let excluded_key_ids = bound.exhausted_exclusions.key_ids(now_unix_secs); diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/responses/session.rs b/apps/aether-gateway/src/handlers/proxy/websocket/responses/session.rs index 266142c48..d8aa948f1 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/responses/session.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/responses/session.rs @@ -444,7 +444,14 @@ async fn bootstrap_responses_websocket( let raw_responses_lite_static_config = ResponsesLiteStaticConfig::from_response_create(&first_event); - let planning_parts = build_planning_parts(context); + let first_logical_turn_id = Uuid::now_v7().to_string(); + let mut planning_parts = build_planning_parts(context); + let first_codex_fingerprint_context = + crate::ai_serving::codex_context::attach_codex_logical_turn_context( + &mut planning_parts, + &first_event, + &first_logical_turn_id, + ); let turn_control = match resolve_responses_websocket_turn_control( &state, context, @@ -860,7 +867,6 @@ async fn bootstrap_responses_websocket( return None; } }; - let first_logical_turn_id = Uuid::new_v4().to_string(); let first_turn_decision = prepare_responses_websocket_turn_decision( &decision, context.trace_id.clone(), @@ -954,6 +960,7 @@ async fn bootstrap_responses_websocket( } bound.turn_state.begin( LogicalTurn::new(first_event, 1, first_logical_turn_id) + .with_codex_fingerprint_context(first_codex_fingerprint_context) .with_provider_store(first_provider_event.get("store") == Some(&Value::Bool(true))) .with_turn_control(turn_control), first_turn, diff --git a/apps/aether-gateway/src/handlers/proxy/websocket/responses/turn_state.rs b/apps/aether-gateway/src/handlers/proxy/websocket/responses/turn_state.rs index 8fbb69e05..ebcd5f2df 100644 --- a/apps/aether-gateway/src/handlers/proxy/websocket/responses/turn_state.rs +++ b/apps/aether-gateway/src/handlers/proxy/websocket/responses/turn_state.rs @@ -5,6 +5,7 @@ //! 非法组合只能靠调用点的 if 和「记得同时改另外两个字段」来避免。这里把它收敛成 //! 一个枚举：合法组合由类型保证，转换只能走受控 API。 +use aether_provider_transport::CodexFingerprintConvergenceContext; use serde_json::Value; use super::control::ResponsesWebSocketTurnControl; @@ -26,6 +27,9 @@ pub(super) struct LogicalTurn { pub(super) provider_store: bool, pub(super) turn_index: u64, pub(super) logical_turn_id: String, + /// Immutable Codex client identity for every provider attempt belonging to + /// this logical turn. A transparent re-plan must never mint a new turn. + pub(super) codex_fingerprint_context: Option, pub(super) turn_attempt: u32, pub(super) retry_attempted: bool, pub(super) retry_unsafe_reason: Option<&'static str>, @@ -42,6 +46,7 @@ impl LogicalTurn { provider_store: false, turn_index, logical_turn_id, + codex_fingerprint_context: None, turn_attempt: 1, retry_attempted: false, retry_unsafe_reason: None, @@ -54,6 +59,14 @@ impl LogicalTurn { self } + pub(super) fn with_codex_fingerprint_context( + mut self, + context: CodexFingerprintConvergenceContext, + ) -> Self { + self.codex_fingerprint_context = Some(context); + self + } + pub(super) fn with_provider_store(mut self, provider_store: bool) -> Self { self.provider_store = provider_store; self 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 14ad409f2..31260643f 100644 --- a/crates/aether-ai/formats/src/formats/openai/responses/codex.rs +++ b/crates/aether-ai/formats/src/formats/openai/responses/codex.rs @@ -851,10 +851,26 @@ fn header_value_contains_media_type(value: &str, media_type: &str) -> bool { #[derive(Debug, Clone, Default, PartialEq, Eq)] pub struct CodexAuthIdentity { pub account_id: Option, + pub account_user_id: Option, + pub user_id: Option, + pub email: Option, + pub codex_identity_fingerprint: Option, pub is_fedramp: bool, pub uses_codex_backend: bool, } +fn first_non_empty_codex_identity_string<'a>( + values: impl IntoIterator>, +) -> Option { + values.into_iter().find_map(|value| { + value + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) + }) +} + pub fn parse_codex_auth_identity(decrypted_auth_config_raw: Option<&str>) -> CodexAuthIdentity { let Some(raw) = decrypted_auth_config_raw .map(str::trim) @@ -868,29 +884,72 @@ pub fn parse_codex_auth_identity(decrypted_auth_config_raw: Option<&str>) -> Cod let namespaced_auth = value .get("https://api.openai.com/auth") .and_then(Value::as_object); + let namespaced_profile = value + .get("https://api.openai.com/profile") + .and_then(Value::as_object); let agent_identity = value .get("agent_identity") .or_else(|| value.get("agentIdentity")) .and_then(Value::as_object); - let account_id = value - .get("account_id") - .or_else(|| value.get("accountId")) - .or_else(|| value.get("chatgpt_account_id")) - .or_else(|| value.get("chatgptAccountId")) - .or_else(|| namespaced_auth.and_then(|auth| auth.get("chatgpt_account_id"))) - .or_else(|| { - agent_identity.and_then(|identity| { - identity - .get("account_id") - .or_else(|| identity.get("accountId")) - .or_else(|| identity.get("chatgpt_account_id")) - .or_else(|| identity.get("chatgptAccountId")) - }) - }) - .and_then(Value::as_str) - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(ToOwned::to_owned); + let account_id = first_non_empty_codex_identity_string([ + value.get("account_id"), + value.get("accountId"), + value.get("chatgpt_account_id"), + value.get("chatgptAccountId"), + namespaced_auth.and_then(|auth| auth.get("chatgpt_account_id")), + agent_identity.and_then(|identity| identity.get("account_id")), + agent_identity.and_then(|identity| identity.get("accountId")), + agent_identity.and_then(|identity| identity.get("chatgpt_account_id")), + agent_identity.and_then(|identity| identity.get("chatgptAccountId")), + ]); + let account_user_id = first_non_empty_codex_identity_string([ + value.get("account_user_id"), + value.get("accountUserId"), + value.get("chatgpt_account_user_id"), + value.get("chatgptAccountUserId"), + namespaced_auth.and_then(|auth| auth.get("chatgpt_account_user_id")), + agent_identity.and_then(|identity| identity.get("account_user_id")), + agent_identity.and_then(|identity| identity.get("accountUserId")), + agent_identity.and_then(|identity| identity.get("chatgpt_account_user_id")), + agent_identity.and_then(|identity| identity.get("chatgptAccountUserId")), + ]); + let user_id = first_non_empty_codex_identity_string([ + value.get("user_id"), + value.get("userId"), + value.get("chatgpt_user_id"), + value.get("chatgptUserId"), + namespaced_auth.and_then(|auth| auth.get("chatgpt_user_id")), + agent_identity.and_then(|identity| identity.get("user_id")), + agent_identity.and_then(|identity| identity.get("userId")), + agent_identity.and_then(|identity| identity.get("chatgpt_user_id")), + agent_identity.and_then(|identity| identity.get("chatgptUserId")), + value.get("sub"), + ]); + let email = first_non_empty_codex_identity_string([ + value.get("email"), + value.get("email_address"), + value.get("emailAddress"), + value.get("outlook_email"), + namespaced_auth.and_then(|auth| auth.get("email")), + namespaced_auth.and_then(|auth| auth.get("email_address")), + namespaced_auth.and_then(|auth| auth.get("emailAddress")), + namespaced_auth.and_then(|auth| auth.get("outlook_email")), + namespaced_profile.and_then(|profile| profile.get("email")), + namespaced_profile.and_then(|profile| profile.get("email_address")), + namespaced_profile.and_then(|profile| profile.get("emailAddress")), + agent_identity.and_then(|identity| identity.get("email")), + agent_identity.and_then(|identity| identity.get("email_address")), + agent_identity.and_then(|identity| identity.get("emailAddress")), + agent_identity.and_then(|identity| identity.get("outlook_email")), + ]); + let codex_identity_fingerprint = first_non_empty_codex_identity_string([ + value.get("codex_identity_fingerprint"), + value.get("codex-identity-fingerprint"), + value.get("codexIdentityFingerprint"), + agent_identity.and_then(|identity| identity.get("codex_identity_fingerprint")), + agent_identity.and_then(|identity| identity.get("codex-identity-fingerprint")), + agent_identity.and_then(|identity| identity.get("codexIdentityFingerprint")), + ]); let is_fedramp = value .get("is_fedramp") .or_else(|| value.get("chatgpt_account_is_fedramp")) @@ -907,6 +966,9 @@ pub fn parse_codex_auth_identity(decrypted_auth_config_raw: Option<&str>) -> Cod .and_then(Value::as_bool) .unwrap_or(false); let uses_codex_backend = account_id.is_some() + || account_user_id.is_some() + || user_id.is_some() + || codex_identity_fingerprint.is_some() || value .get("provider_type") .and_then(Value::as_str) @@ -919,6 +981,10 @@ pub fn parse_codex_auth_identity(decrypted_auth_config_raw: Option<&str>) -> Cod CodexAuthIdentity { account_id, + account_user_id, + user_id, + email, + codex_identity_fingerprint, is_fedramp, uses_codex_backend, } @@ -2125,7 +2191,8 @@ mod tests { apply_codex_openai_responses_websocket_continuation_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, - project_codex_catalog_model_card, resolve_codex_responses_model_capabilities, + parse_codex_auth_identity, 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, @@ -2778,6 +2845,41 @@ mod tests { assert!(capabilities.supported_service_tiers.is_empty()); } + #[test] + fn codex_auth_identity_parses_member_claims_and_persisted_fingerprint() { + let identity = parse_codex_auth_identity(Some( + &json!({ + "provider_type": "codex", + "accountId": "workspace-1", + "codexIdentityFingerprint": "codex-persisted-fingerprint:v1:stable", + "https://api.openai.com/auth": { + "chatgpt_account_user_id": "workspace-member-1", + "chatgpt_user_id": "user-1" + }, + "https://api.openai.com/profile": { + "email": "Alice@Example.com" + } + }) + .to_string(), + )); + + assert_eq!(identity.account_id.as_deref(), Some("workspace-1")); + assert_eq!( + identity.account_user_id.as_deref(), + Some("workspace-member-1") + ); + assert_eq!(identity.user_id.as_deref(), Some("user-1")); + assert_eq!(identity.email.as_deref(), Some("Alice@Example.com")); + assert_eq!( + identity.codex_identity_fingerprint.as_deref(), + Some("codex-persisted-fingerprint:v1:stable") + ); + assert!(identity.uses_codex_backend); + + let sub_fallback = parse_codex_auth_identity(Some(r#"{"sub":"legacy-user-1"}"#)); + assert_eq!(sub_fallback.user_id.as_deref(), Some("legacy-user-1")); + } + #[test] fn codex_identity_headers_are_derived_only_from_auth_config() { let mut headers = std::collections::BTreeMap::from([ diff --git a/crates/aether-oauth/src/provider/providers/generic.rs b/crates/aether-oauth/src/provider/providers/generic.rs index f81c10fe9..c0c563ba6 100644 --- a/crates/aether-oauth/src/provider/providers/generic.rs +++ b/crates/aether-oauth/src/provider/providers/generic.rs @@ -17,6 +17,35 @@ use super::claude_code::{ CLAUDE_CODE_PROVIDER_TYPE, CLAUDE_CODE_REDIRECT_URI, CLAUDE_CODE_TOKEN_URL, }; +const CODEX_IDENTITY_FINGERPRINT_FIELD: &str = "codex_identity_fingerprint"; +const CODEX_IDENTITY_FINGERPRINT_VERSION: &str = "codex-persisted-fingerprint:v1"; + +pub fn derive_codex_identity_fingerprint( + account_id: Option<&str>, + account_user_id: Option<&str>, + user_id: Option<&str>, + email: Option<&str>, +) -> Option { + let account = normalized_codex_identity_value(account_id); + let member = normalized_codex_identity_value(account_user_id) + .or_else(|| normalized_codex_identity_value(user_id)) + .or_else(|| normalized_codex_identity_value(email))?; + + let mut digest = Sha256::new(); + digest.update(CODEX_IDENTITY_FINGERPRINT_VERSION.as_bytes()); + digest.update([0]); + digest.update(account.as_deref().unwrap_or("").as_bytes()); + digest.update([0]); + digest.update(member.as_bytes()); + let digest = digest.finalize(); + let mut encoded = String::with_capacity(digest.len() * 2); + for byte in digest { + use std::fmt::Write as _; + let _ = write!(&mut encoded, "{byte:02x}"); + } + Some(format!("{CODEX_IDENTITY_FINGERPRINT_VERSION}:{encoded}")) +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct GenericProviderOAuthTemplate { pub provider_type: &'static str, @@ -287,6 +316,7 @@ impl GenericProviderOAuthAdapter { auth_config.insert("scope".to_string(), json!(scope)); } enrich_generic_identity(self.template.provider_type, &mut auth_config, &payload); + ensure_codex_identity_fingerprint(self.template.provider_type, &mut auth_config); Ok(ProviderOAuthTokenSet { token_set, auth_config: Value::Object(auth_config), @@ -383,6 +413,16 @@ impl ProviderOAuthAdapter for GenericProviderOAuthAdapter { let mut refreshed = self .exchange_grant(executor, ctx, "refresh_token", refresh_token, None, None) .await?; + let existing_codex_identity_fingerprint = self + .template + .provider_type + .eq_ignore_ascii_case("codex") + .then(|| { + codex_identity_fingerprint_value(&account.auth_config).or_else(|| { + derive_codex_identity_fingerprint_from_auth_config(&account.auth_config) + }) + }) + .flatten(); // Refresh responses often omit stable account metadata, and some providers // do not rotate refresh_token on every refresh. Preserve the stored config @@ -398,6 +438,13 @@ impl ProviderOAuthAdapter for GenericProviderOAuthAdapter { refreshed.token_set.refresh_token = Some(refresh_token.to_string()); merged.insert("refresh_token".to_string(), json!(refresh_token)); } + if let Some(fingerprint) = existing_codex_identity_fingerprint { + merged.insert( + CODEX_IDENTITY_FINGERPRINT_FIELD.to_string(), + Value::String(fingerprint), + ); + } + ensure_codex_identity_fingerprint(self.template.provider_type, &mut merged); refreshed.auth_config = Value::Object(merged); } Ok(refreshed) @@ -411,6 +458,11 @@ impl ProviderOAuthAdapter for GenericProviderOAuthAdapter { } fn account_fingerprint(&self, account: &ProviderOAuthAccount) -> Option { + if self.template.provider_type.eq_ignore_ascii_case("codex") { + return codex_identity_fingerprint_value(&account.auth_config).or_else(|| { + derive_codex_identity_fingerprint_from_auth_config(&account.auth_config) + }); + } let refresh_token = account .auth_config .get("refresh_token") @@ -472,6 +524,89 @@ fn secret_fingerprint(value: &str) -> String { fingerprint } +fn codex_identity_fingerprint_value(auth_config: &Value) -> Option { + [ + CODEX_IDENTITY_FINGERPRINT_FIELD, + "codex-identity-fingerprint", + "codexIdentityFingerprint", + ] + .iter() + .find_map(|field| auth_config.get(*field).and_then(Value::as_str)) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) +} + +fn codex_identity_claim(auth_config: &Value, fields: &[&str]) -> Option { + fields + .iter() + .find_map(|field| auth_config.get(*field).and_then(Value::as_str)) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_ascii_lowercase) +} + +fn normalized_codex_identity_value(value: Option<&str>) -> Option { + value + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_ascii_lowercase) +} + +fn derive_codex_identity_fingerprint_from_auth_config(auth_config: &Value) -> Option { + let account = codex_identity_claim( + auth_config, + &[ + "account_id", + "accountId", + "chatgpt_account_id", + "chatgptAccountId", + ], + ); + let account_user = codex_identity_claim( + auth_config, + &[ + "account_user_id", + "accountUserId", + "chatgpt_account_user_id", + "chatgptAccountUserId", + ], + ); + let user = codex_identity_claim( + auth_config, + &["user_id", "userId", "chatgpt_user_id", "chatgptUserId"], + ); + let email = codex_identity_claim( + auth_config, + &["email", "email_address", "emailAddress", "outlook_email"], + ); + + derive_codex_identity_fingerprint( + account.as_deref(), + account_user.as_deref(), + user.as_deref(), + email.as_deref(), + ) +} + +fn ensure_codex_identity_fingerprint( + provider_type: &str, + auth_config: &mut serde_json::Map, +) { + if !provider_type.eq_ignore_ascii_case("codex") { + return; + } + let auth_config_value = Value::Object(auth_config.clone()); + let fingerprint = codex_identity_fingerprint_value(&auth_config_value) + .or_else(|| derive_codex_identity_fingerprint_from_auth_config(&auth_config_value)); + if let Some(fingerprint) = fingerprint { + auth_config.insert( + CODEX_IDENTITY_FINGERPRINT_FIELD.to_string(), + Value::String(fingerprint), + ); + } +} + fn enrich_generic_identity( provider_type: &str, auth_config: &mut serde_json::Map, @@ -480,12 +615,21 @@ fn enrich_generic_identity( if let Some(object) = token_payload.as_object() { for field in [ "email", + "email_address", + "emailAddress", + "outlook_email", "account_id", + "accountId", "account_user_id", + "accountUserId", "plan_type", "user_id", + "userId", "account_name", "is_fedramp", + CODEX_IDENTITY_FINGERPRINT_FIELD, + "codex-identity-fingerprint", + "codexIdentityFingerprint", ] { if !auth_config.contains_key(field) { if let Some(value) = object.get(field).cloned() { @@ -565,8 +709,11 @@ fn enrich_generic_identity( .get("https://api.openai.com/profile") .and_then(Value::as_object) { - if let Some(value) = profile.get("email").cloned() { - auth_config.entry("email".to_string()).or_insert(value); + for field in ["email", "email_address", "emailAddress", "outlook_email"] { + if let Some(value) = profile.get(field).cloned() { + auth_config.entry("email".to_string()).or_insert(value); + break; + } } } } @@ -653,13 +800,16 @@ fn decode_jwt_claims(token: &str) -> Option> { #[cfg(test)] mod tests { - use super::{enrich_generic_identity, template_for_provider_type, GenericProviderOAuthAdapter}; + use super::{ + derive_codex_identity_fingerprint, enrich_generic_identity, template_for_provider_type, + GenericProviderOAuthAdapter, CODEX_IDENTITY_FINGERPRINT_FIELD, + }; use crate::network::{OAuthHttpExecutor, OAuthHttpRequest, OAuthHttpResponse}; use crate::provider::ProviderOAuthAdapter; use crate::provider::{ProviderOAuthAccount, ProviderOAuthTransportContext}; use async_trait::async_trait; use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _}; - use serde_json::json; + use serde_json::{json, Value}; use std::collections::BTreeMap; use std::sync::{Arc, Mutex}; @@ -678,6 +828,13 @@ mod tests { assert!(adapter.capabilities().supports_refresh_token_import); } + fn encoded_jwt(claims: &Value) -> String { + format!( + "header.{}.signature", + URL_SAFE_NO_PAD.encode(serde_json::to_vec(claims).expect("claims should encode")) + ) + } + #[test] fn codex_identity_extracts_fedramp_workspace_claim() { let claims = json!({ @@ -686,10 +843,7 @@ mod tests { "chatgpt_account_is_fedramp": true } }); - let token = format!( - "header.{}.signature", - URL_SAFE_NO_PAD.encode(serde_json::to_vec(&claims).expect("claims should encode")) - ); + let token = encoded_jwt(&claims); let mut auth_config = serde_json::Map::new(); enrich_generic_identity("codex", &mut auth_config, &json!({"access_token": token})); @@ -698,9 +852,74 @@ mod tests { assert_eq!(auth_config.get("is_fedramp"), Some(&json!(true))); } + #[test] + fn codex_persisted_fingerprint_is_member_scoped_and_token_independent() { + let adapter = GenericProviderOAuthAdapter::for_provider_type("codex") + .expect("codex adapter should exist"); + let claims = json!({ + "sub": "global-user-1", + "https://api.openai.com/auth": { + "chatgpt_account_id": "workspace-1", + "chatgpt_account_user_id": "member-1" + } + }); + let rotated_claims = json!({ + "sub": "global-user-1", + "iat": 12345, + "https://api.openai.com/auth": { + "chatgpt_account_id": "WORKSPACE-1", + "chatgpt_account_user_id": "MEMBER-1" + } + }); + let other_member_claims = json!({ + "sub": "global-user-2", + "https://api.openai.com/auth": { + "chatgpt_account_id": "workspace-1", + "chatgpt_account_user_id": "member-2" + } + }); + + let first = adapter + .token_set_from_payload(json!({"access_token": encoded_jwt(&claims)})) + .expect("first token should parse"); + let rotated = adapter + .token_set_from_payload(json!({"access_token": encoded_jwt(&rotated_claims)})) + .expect("rotated token should parse"); + let other_member = adapter + .token_set_from_payload(json!({"access_token": encoded_jwt(&other_member_claims)})) + .expect("other member token should parse"); + + let first_fingerprint = first.auth_config[CODEX_IDENTITY_FINGERPRINT_FIELD] + .as_str() + .expect("persisted fingerprint") + .to_string(); + assert!(first_fingerprint.starts_with("codex-persisted-fingerprint:v1:")); + assert_eq!( + rotated.auth_config[CODEX_IDENTITY_FINGERPRINT_FIELD].as_str(), + Some(first_fingerprint.as_str()) + ); + assert_ne!( + other_member.auth_config[CODEX_IDENTITY_FINGERPRINT_FIELD].as_str(), + Some(first_fingerprint.as_str()) + ); + + let account = ProviderOAuthAccount { + provider_type: "codex".to_string(), + access_token: "unrelated-rotated-token".to_string(), + auth_config: first.auth_config, + expires_at_unix_secs: None, + identity: BTreeMap::new(), + }; + assert_eq!( + adapter.account_fingerprint(&account).as_deref(), + Some(first_fingerprint.as_str()) + ); + } + #[derive(Debug, Clone)] struct StaticExecutor { seen_request: Arc>>, + response_payload: Value, } #[async_trait] @@ -712,11 +931,7 @@ mod tests { *self.seen_request.lock().expect("mutex should lock") = Some(request); Ok(OAuthHttpResponse { status_code: 200, - body_text: json!({ - "access_token": "new-access-token", - "expires_in": 3600 - }) - .to_string(), + body_text: self.response_payload.to_string(), json_body: None, }) } @@ -725,12 +940,29 @@ mod tests { #[tokio::test] async fn refresh_preserves_existing_metadata_when_refresh_token_is_not_rotated() { let seen_request = Arc::new(Mutex::new(None)); + let refreshed_token = encoded_jwt(&json!({ + "https://api.openai.com/auth": { + "chatgpt_account_id": "acct-123", + "chatgpt_account_user_id": "replacement-member" + } + })); let executor = StaticExecutor { seen_request: Arc::clone(&seen_request), + response_payload: json!({ + "access_token": refreshed_token, + "expires_in": 3600 + }), }; let adapter = GenericProviderOAuthAdapter::for_provider_type("codex") .expect("codex adapter should exist") .with_token_url_override("https://auth.example.test/token"); + let expected_legacy_fingerprint = derive_codex_identity_fingerprint( + Some("acct-123"), + Some("original-member"), + None, + Some("alice@example.com"), + ) + .expect("legacy identity should produce a fingerprint"); let ctx = ProviderOAuthTransportContext { provider_id: "provider-1".to_string(), provider_type: "codex".to_string(), @@ -752,6 +984,7 @@ mod tests { "refresh_token": "old-refresh-token", "email": "alice@example.com", "account_id": "acct-123", + "account_user_id": "original-member", "updated_at": 1 }), expires_at_unix_secs: Some(1), @@ -763,7 +996,7 @@ mod tests { .await .expect("refresh should succeed"); - assert_eq!(refreshed.token_set.access_token, "new-access-token"); + assert_eq!(refreshed.token_set.access_token, refreshed_token); assert_eq!( refreshed.token_set.refresh_token.as_deref(), Some("old-refresh-token") @@ -771,6 +1004,10 @@ mod tests { assert_eq!(refreshed.auth_config["email"], "alice@example.com"); assert_eq!(refreshed.auth_config["account_id"], "acct-123"); assert_eq!(refreshed.auth_config["refresh_token"], "old-refresh-token"); + assert_eq!( + refreshed.auth_config[CODEX_IDENTITY_FINGERPRINT_FIELD].as_str(), + Some(expected_legacy_fingerprint.as_str()) + ); let seen = seen_request .lock() diff --git a/crates/aether-oauth/src/provider/providers/mod.rs b/crates/aether-oauth/src/provider/providers/mod.rs index a9e72c092..4fe916a75 100644 --- a/crates/aether-oauth/src/provider/providers/mod.rs +++ b/crates/aether-oauth/src/provider/providers/mod.rs @@ -13,7 +13,8 @@ pub use claude_code::{ }; pub use codex::CodexProviderOAuthAdapter; pub use generic::{ - GenericProviderOAuthAdapter, GenericProviderOAuthTemplate, GENERIC_PROVIDER_OAUTH_TEMPLATES, + derive_codex_identity_fingerprint, GenericProviderOAuthAdapter, GenericProviderOAuthTemplate, + GENERIC_PROVIDER_OAUTH_TEMPLATES, }; pub use kiro::{ generate_kiro_machine_id, normalize_kiro_machine_id, KiroAuthConfig, KiroProviderOAuthAdapter, diff --git a/crates/aether-provider/transport/src/codex_fingerprint.rs b/crates/aether-provider/transport/src/codex_fingerprint.rs index 32c56ffb3..f406e2eb4 100644 --- a/crates/aether-provider/transport/src/codex_fingerprint.rs +++ b/crates/aether-provider/transport/src/codex_fingerprint.rs @@ -10,6 +10,68 @@ use crate::snapshot::GatewayProviderTransportSnapshot; pub const CODEX_FINGERPRINT_CONFIG_NAMESPACE: &str = "codex"; pub const CODEX_FINGERPRINT_ENABLED_CONFIG_KEY: &str = "fingerprint_convergence_enabled"; +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CodexFingerprintConvergenceContext { + logical_turn_id: String, + original_turn_id: Option, + original_client_session_id: Option, + original_prompt_cache_key: Option, + turn_started_at_unix_ms: u64, +} + +impl CodexFingerprintConvergenceContext { + pub fn new(logical_turn_id: impl Into, turn_started_at_unix_ms: u64) -> Self { + Self { + logical_turn_id: logical_turn_id.into().trim().to_string(), + original_turn_id: None, + original_client_session_id: None, + original_prompt_cache_key: None, + turn_started_at_unix_ms, + } + } + + pub fn with_original_turn_id(mut self, original_turn_id: impl Into) -> Self { + self.original_turn_id = non_empty_owned(original_turn_id.into()); + self + } + + pub fn with_original_client_session_id( + mut self, + original_client_session_id: impl Into, + ) -> Self { + self.original_client_session_id = non_empty_owned(original_client_session_id.into()); + self + } + + pub fn with_original_prompt_cache_key( + mut self, + original_prompt_cache_key: impl Into, + ) -> Self { + self.original_prompt_cache_key = non_empty_owned(original_prompt_cache_key.into()); + self + } + + pub fn logical_turn_id(&self) -> &str { + self.logical_turn_id.as_str() + } + + pub fn original_turn_id(&self) -> Option<&str> { + self.original_turn_id.as_deref() + } + + pub fn original_client_session_id(&self) -> Option<&str> { + self.original_client_session_id.as_deref() + } + + pub fn original_prompt_cache_key(&self) -> Option<&str> { + self.original_prompt_cache_key.as_deref() + } + + pub fn turn_started_at_unix_ms(&self) -> u64 { + self.turn_started_at_unix_ms + } +} + #[derive(Debug, Clone, PartialEq, Eq)] struct CodexConvergedFingerprint { installation_id: String, @@ -18,6 +80,7 @@ struct CodexConvergedFingerprint { turn_id: String, window_id: String, turn_started_at_unix_ms: u64, + prompt_cache_key: Option, } pub fn codex_fingerprint_convergence_enabled( @@ -40,6 +103,27 @@ pub fn apply_codex_oauth_fingerprint_convergence( original_client_session_id: Option<&str>, provider_request_headers: &mut BTreeMap, provider_request_body: &mut Value, +) -> bool { + let mut context = + CodexFingerprintConvergenceContext::new(Uuid::now_v7().to_string(), current_unix_millis()); + if let Some(original_client_session_id) = original_client_session_id { + context = context.with_original_client_session_id(original_client_session_id); + } + apply_codex_oauth_fingerprint_convergence_with_context( + transport, + provider_api_format, + &context, + provider_request_headers, + provider_request_body, + ) +} + +pub fn apply_codex_oauth_fingerprint_convergence_with_context( + transport: &GatewayProviderTransportSnapshot, + provider_api_format: &str, + context: &CodexFingerprintConvergenceContext, + provider_request_headers: &mut BTreeMap, + provider_request_body: &mut Value, ) -> bool { let is_responses = aether_ai_formats::is_openai_responses_format(provider_api_format); let is_live = aether_ai_formats::api_format_alias_matches(provider_api_format, "codex:live"); @@ -68,13 +152,22 @@ pub fn apply_codex_oauth_fingerprint_convergence( let auth_identity = aether_ai_formats::parse_codex_auth_identity( transport.key.decrypted_auth_config.as_deref(), ); - let account_seed = auth_identity - .account_id - .as_deref() + let account_seed = resolve_codex_account_seed(&auth_identity, transport.key.id.as_str()); + // Only namespace a cache key that survived all provider-body conversion and + // routing rules. The client-side value in `context` is a retry signal, not + // permission to resurrect a field that the terminal body deliberately + // removed. + let effective_prompt_cache_key = provider_request_body + .as_object() + .and_then(|body| body.get("prompt_cache_key")) + .and_then(Value::as_str) .map(str::trim) - .filter(|value| !value.is_empty()) - .unwrap_or(transport.key.id.as_str()); - let fingerprint = resolve_converged_fingerprint(account_seed, original_client_session_id); + .filter(|value| !value.is_empty()); + let fingerprint = resolve_converged_fingerprint_with_prompt_cache( + account_seed.as_str(), + context, + effective_prompt_cache_key, + ); apply_converged_headers(provider_request_headers, &fingerprint); // Live uses the converged identity on the WebSocket/call-control headers. @@ -86,14 +179,24 @@ pub fn apply_codex_oauth_fingerprint_convergence( true } +#[cfg(test)] fn resolve_converged_fingerprint( account_seed: &str, - original_client_session_id: Option<&str>, + context: &CodexFingerprintConvergenceContext, +) -> CodexConvergedFingerprint { + resolve_converged_fingerprint_with_prompt_cache(account_seed, context, None) +} + +fn resolve_converged_fingerprint_with_prompt_cache( + account_seed: &str, + context: &CodexFingerprintConvergenceContext, + effective_prompt_cache_key: Option<&str>, ) -> CodexConvergedFingerprint { let installation_id = derive_stable_uuid_v4(&format!("aether:codex-installation-id:v1:{account_seed}")); let session_id = derive_stable_uuid_v4(&format!("aether:codex-session-id:v1:{account_seed}")); - let original_client_session_id = original_client_session_id + let original_client_session_id = context + .original_client_session_id() .map(str::trim) .filter(|value| !value.is_empty()); let thread_id = original_client_session_id @@ -104,18 +207,84 @@ fn resolve_converged_fingerprint( }) .unwrap_or_else(|| session_id.clone()); let window_id = format!("{thread_id}:0"); - let turn_started_at_unix_ms = SystemTime::now() - .duration_since(UNIX_EPOCH) - .map(|duration| duration.as_millis().min(u64::MAX as u128) as u64) - .unwrap_or(0); + let turn_identity = context + .original_turn_id() + .map(|turn_id| ("original", turn_id)) + .unwrap_or_else(|| ("logical", context.logical_turn_id())); + let turn_id = derive_stable_uuid_v7( + context.turn_started_at_unix_ms(), + &format!( + "aether:codex-turn-id:v1\0{account_seed}\0{}\0{}", + turn_identity.0, turn_identity.1 + ), + ); + let prompt_cache_key = context + .original_prompt_cache_key() + .and(effective_prompt_cache_key) + .map(|effective| { + Uuid::new_v5( + &Uuid::NAMESPACE_URL, + format!("aether:codex-prompt-cache-key:v1\0{account_seed}\0{effective}").as_bytes(), + ) + .to_string() + }); CodexConvergedFingerprint { installation_id, session_id, thread_id, - turn_id: Uuid::now_v7().to_string(), + turn_id, window_id, - turn_started_at_unix_ms, + turn_started_at_unix_ms: context.turn_started_at_unix_ms(), + prompt_cache_key, + } +} + +fn non_empty_owned(value: String) -> Option { + let value = value.trim(); + (!value.is_empty()).then(|| value.to_string()) +} + +fn current_unix_millis() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_millis().min(u64::MAX as u128) as u64) + .unwrap_or(0) +} + +fn normalized_identity_part(value: Option<&str>) -> Option { + value + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_ascii_lowercase) +} + +fn resolve_codex_account_seed( + identity: &aether_ai_formats::CodexAuthIdentity, + fallback_key_id: &str, +) -> String { + let fingerprint = normalized_identity_part(identity.codex_identity_fingerprint.as_deref()) + .or_else(|| { + aether_oauth::provider::providers::derive_codex_identity_fingerprint( + identity.account_id.as_deref(), + identity.account_user_id.as_deref(), + identity.user_id.as_deref(), + identity.email.as_deref(), + ) + }); + if let Some(fingerprint) = fingerprint { + return format!("persisted:v1:{fingerprint}"); + } + + let account = normalized_identity_part(identity.account_id.as_deref()); + let member = normalized_identity_part(identity.account_user_id.as_deref()) + .or_else(|| normalized_identity_part(identity.user_id.as_deref())) + .or_else(|| normalized_identity_part(identity.email.as_deref())); + match (account, member) { + (Some(account), Some(member)) => format!("account-member:v1:{account}\0{member}"), + (None, Some(member)) => format!("member:v1:{member}"), + (Some(account), None) => format!("account:v1:{account}"), + (None, None) => format!("key:v1:{}", fallback_key_id.trim()), } } @@ -128,6 +297,17 @@ fn derive_stable_uuid_v4(seed: &str) -> String { Uuid::from_bytes(bytes).to_string() } +fn derive_stable_uuid_v7(timestamp_ms: u64, seed: &str) -> String { + let digest = Sha256::digest(seed.as_bytes()); + let mut bytes = [0_u8; 16]; + let timestamp_bytes = timestamp_ms.min(0x0000_ffff_ffff_ffff).to_be_bytes(); + bytes[..6].copy_from_slice(×tamp_bytes[2..]); + bytes[6..].copy_from_slice(&digest[..10]); + bytes[6] = (bytes[6] & 0x0f) | 0x70; + bytes[8] = (bytes[8] & 0x3f) | 0x80; + Uuid::from_bytes(bytes).to_string() +} + fn apply_converged_headers( headers: &mut BTreeMap, fingerprint: &CodexConvergedFingerprint, @@ -157,6 +337,12 @@ fn apply_converged_client_metadata(body: &mut Value, fingerprint: &CodexConverge let Some(body) = body.as_object_mut() else { return; }; + if let Some(prompt_cache_key) = fingerprint.prompt_cache_key.as_ref() { + body.insert( + "prompt_cache_key".to_string(), + Value::String(prompt_cache_key.clone()), + ); + } let metadata = body .entry("client_metadata".to_string()) .or_insert_with(|| Value::Object(Map::new())); @@ -212,20 +398,23 @@ fn rewrite_embedded_turn_metadata( metadata: &mut Map, fingerprint: &CodexConvergedFingerprint, ) { - let Some(raw) = metadata - .get("x-codex-turn-metadata") - .and_then(Value::as_str) - else { + let Some(turn_metadata) = metadata.get_mut("x-codex-turn-metadata") else { return; }; - let Ok(mut turn_metadata) = serde_json::from_str::>(raw) else { - return; - }; - apply_turn_metadata_fields(&mut turn_metadata, fingerprint); - let Ok(rebuilt) = serde_json::to_string(&turn_metadata) else { - return; - }; - metadata.insert("x-codex-turn-metadata".to_string(), Value::String(rebuilt)); + match turn_metadata { + Value::Object(turn_metadata) => apply_turn_metadata_fields(turn_metadata, fingerprint), + Value::String(raw) => { + let Ok(mut parsed) = serde_json::from_str::>(raw) else { + return; + }; + apply_turn_metadata_fields(&mut parsed, fingerprint); + let Ok(rebuilt) = serde_json::to_string(&parsed) else { + return; + }; + *raw = rebuilt; + } + _ => {} + } } fn apply_turn_metadata_fields( @@ -256,6 +445,12 @@ fn apply_turn_metadata_fields( "turn_started_at_unix_ms".to_string(), Value::from(fingerprint.turn_started_at_unix_ms), ); + if let Some(prompt_cache_key) = fingerprint.prompt_cache_key.as_ref() { + metadata.insert( + "prompt_cache_key".to_string(), + Value::String(prompt_cache_key.clone()), + ); + } } fn set_header(headers: &mut BTreeMap, name: &str, value: String) { @@ -470,19 +665,316 @@ mod tests { } #[test] - fn stable_account_identity_and_per_client_thread_are_deterministic() { - let first = resolve_converged_fingerprint("account-1", Some("client-a")); - let second = resolve_converged_fingerprint("account-1", Some("client-a")); - let other_client = resolve_converged_fingerprint("account-1", Some("client-b")); - let other_account = resolve_converged_fingerprint("account-2", Some("client-a")); + fn object_form_embedded_turn_metadata_is_rewritten_with_the_same_identity() { + let transport = sample_transport(); + let context = CodexFingerprintConvergenceContext::new("logical-turn-1", 1_700_000_000_123) + .with_original_turn_id("client-turn-1") + .with_original_client_session_id("client-session-1") + .with_original_prompt_cache_key("client-cache-1"); + let mut headers = BTreeMap::new(); + let mut body = json!({ + "model": "gpt-5.4", + "client_metadata": { + "x-codex-turn-metadata": { + "installation_id": "old-installation", + "session_id": "old-session", + "thread_id": "old-thread", + "turn_id": "old-turn", + "window_id": "old-window", + "custom": "preserved" + } + } + }); + + assert!(apply_codex_oauth_fingerprint_convergence_with_context( + &transport, + "openai:responses", + &context, + &mut headers, + &mut body, + )); + + let embedded = &body["client_metadata"]["x-codex-turn-metadata"]; + assert!(embedded.is_object()); + assert_eq!(embedded["custom"], "preserved"); + assert_eq!( + embedded["installation_id"], + body["client_metadata"]["x-codex-installation-id"] + ); + assert_eq!( + embedded["session_id"], + body["client_metadata"]["session_id"] + ); + assert_eq!(embedded["thread_id"], body["client_metadata"]["thread_id"]); + assert_eq!(embedded["turn_id"], body["client_metadata"]["turn_id"]); + assert_eq!( + embedded["window_id"], + body["client_metadata"]["x-codex-window-id"] + ); + assert_eq!(embedded["prompt_cache_key"], body["prompt_cache_key"]); + assert_eq!( + embedded["turn_started_at_unix_ms"], + context.turn_started_at_unix_ms() + ); + } + + #[test] + fn stable_account_identity_client_thread_and_logical_turn_are_deterministic() { + let context = CodexFingerprintConvergenceContext::new("logical-turn-1", 1_700_000_000_123) + .with_original_client_session_id("client-a"); + let other_turn_context = + CodexFingerprintConvergenceContext::new("logical-turn-2", 1_700_000_000_124) + .with_original_client_session_id("client-a"); + let other_client_context = context.clone().with_original_client_session_id("client-b"); + let first = resolve_converged_fingerprint("account-1", &context); + let second = resolve_converged_fingerprint("account-1", &context); + let other_turn = resolve_converged_fingerprint("account-1", &other_turn_context); + let other_client = resolve_converged_fingerprint("account-1", &other_client_context); + let other_account = resolve_converged_fingerprint("account-2", &context); assert_eq!(first.installation_id, second.installation_id); assert_eq!(first.session_id, second.session_id); assert_eq!(first.thread_id, second.thread_id); - assert_ne!(first.turn_id, second.turn_id); + assert_eq!(first.turn_id, second.turn_id); + assert_eq!(first.turn_started_at_unix_ms, 1_700_000_000_123); + assert_ne!(first.turn_id, other_turn.turn_id); assert_ne!(first.thread_id, other_client.thread_id); assert_eq!(first.session_id, other_client.session_id); assert_ne!(first.installation_id, other_account.installation_id); + assert_ne!(first.turn_id, other_account.turn_id); + assert_eq!( + Uuid::parse_str(&first.turn_id) + .expect("turn UUID") + .get_version_num(), + 7 + ); + assert_eq!( + first.turn_id.replace('-', "")[..12], + format!("{:012x}", context.turn_started_at_unix_ms()) + ); + } + + #[test] + fn legacy_wrapper_keeps_a_fresh_turn_and_does_not_namespace_prompt_cache() { + let transport = sample_transport(); + let mut first_headers = BTreeMap::new(); + let mut second_headers = BTreeMap::new(); + let mut first_body = json!({ + "model": "gpt-5.4", + "prompt_cache_key": "existing-cache" + }); + let mut second_body = first_body.clone(); + + assert!(apply_codex_oauth_fingerprint_convergence( + &transport, + "openai:responses", + Some("client-session"), + &mut first_headers, + &mut first_body, + )); + assert!(apply_codex_oauth_fingerprint_convergence( + &transport, + "openai:responses", + Some("client-session"), + &mut second_headers, + &mut second_body, + )); + + assert_ne!( + first_body["client_metadata"]["turn_id"], + second_body["client_metadata"]["turn_id"] + ); + assert_eq!(first_body["prompt_cache_key"], "existing-cache"); + assert_eq!(second_body["prompt_cache_key"], "existing-cache"); + } + + #[test] + fn convergence_context_reuses_original_turn_time_and_namespaced_prompt_cache_key() { + let mut transport = sample_transport(); + transport.key.decrypted_auth_config = Some( + json!({ + "account_id": "workspace-1", + "account_user_id": "member-1", + "codex_identity_fingerprint": "codex-persisted-fingerprint:v1:member-1" + }) + .to_string(), + ); + let context = + CodexFingerprintConvergenceContext::new("logical-attempt-1", 1_700_000_000_123) + .with_original_turn_id("client-turn-1") + .with_original_client_session_id("client-session-1") + .with_original_prompt_cache_key("client-cache-1"); + let original_headers = BTreeMap::from([( + "x-codex-turn-metadata".to_string(), + json!({ + "turn_id": "client-turn-1", + "prompt_cache_key": "client-cache-1" + }) + .to_string(), + )]); + let original_body = json!({ + "model": "gpt-5.4", + "prompt_cache_key": "already-adapted-cache", + "client_metadata": { + "x-codex-turn-metadata": json!({ + "turn_id": "client-turn-1", + "prompt_cache_key": "client-cache-1" + }).to_string() + } + }); + + let mut first_headers = original_headers.clone(); + let mut first_body = original_body.clone(); + let mut retried_headers = original_headers; + let mut retried_body = original_body; + assert!(apply_codex_oauth_fingerprint_convergence_with_context( + &transport, + "openai:responses", + &context, + &mut first_headers, + &mut first_body, + )); + assert!(apply_codex_oauth_fingerprint_convergence_with_context( + &transport, + "openai:responses", + &context, + &mut retried_headers, + &mut retried_body, + )); + + assert_eq!(first_headers, retried_headers); + assert_eq!(first_body, retried_body); + assert_eq!(context.logical_turn_id(), "logical-attempt-1"); + assert_eq!(context.original_turn_id(), Some("client-turn-1")); + assert_eq!( + context.original_client_session_id(), + Some("client-session-1") + ); + assert_eq!(context.original_prompt_cache_key(), Some("client-cache-1")); + assert_eq!(context.turn_started_at_unix_ms(), 1_700_000_000_123); + + let prompt_cache_key = first_body["prompt_cache_key"] + .as_str() + .expect("prompt cache key"); + assert_ne!(prompt_cache_key, "client-cache-1"); + assert_ne!(prompt_cache_key, "already-adapted-cache"); + assert_eq!( + Uuid::parse_str(prompt_cache_key) + .expect("prompt cache UUID") + .get_version_num(), + 5 + ); + let header_metadata: Value = + serde_json::from_str(&first_headers["x-codex-turn-metadata"]).expect("header metadata"); + let body_metadata: Value = serde_json::from_str( + first_body["client_metadata"]["x-codex-turn-metadata"] + .as_str() + .expect("body metadata"), + ) + .expect("body metadata json"); + assert_eq!(header_metadata["prompt_cache_key"], prompt_cache_key); + assert_eq!(body_metadata["prompt_cache_key"], prompt_cache_key); + assert_eq!( + header_metadata["turn_started_at_unix_ms"], + 1_700_000_000_123_u64 + ); + assert_eq!( + body_metadata["turn_started_at_unix_ms"], + 1_700_000_000_123_u64 + ); + + let same_original_turn = context.clone().with_original_turn_id("client-turn-1"); + let changed_logical_turn = + CodexFingerprintConvergenceContext::new("logical-attempt-2", 1_700_000_000_123) + .with_original_turn_id("client-turn-1"); + assert_eq!( + resolve_converged_fingerprint("account-1", &same_original_turn).turn_id, + resolve_converged_fingerprint("account-1", &changed_logical_turn).turn_id + ); + } + + #[test] + fn convergence_does_not_resurrect_a_removed_prompt_cache_key() { + let transport = sample_transport(); + let context = CodexFingerprintConvergenceContext::new("logical-turn", 1_700_000_000_123) + .with_original_prompt_cache_key("client-cache"); + let mut body = json!({ + "model": "gpt-5.4", + "client_metadata": {} + }); + let mut headers = BTreeMap::new(); + + assert!(apply_codex_oauth_fingerprint_convergence_with_context( + &transport, + "openai:responses", + &context, + &mut headers, + &mut body, + )); + + assert!(body.get("prompt_cache_key").is_none()); + assert!(body["client_metadata"].get("prompt_cache_key").is_none()); + } + + #[test] + fn persisted_or_canonical_member_identity_drives_the_account_seed() { + let persisted_a = aether_ai_formats::parse_codex_auth_identity(Some( + r#"{"account_id":"workspace-a","email":"old@example.com","codex_identity_fingerprint":"Stable-Member"}"#, + )); + let persisted_b = aether_ai_formats::parse_codex_auth_identity(Some( + r#"{"account_id":"workspace-b","email":"new@example.com","codex_identity_fingerprint":"stable-member"}"#, + )); + assert_eq!( + resolve_codex_account_seed(&persisted_a, "key-a"), + resolve_codex_account_seed(&persisted_b, "key-b") + ); + + let canonical_a = aether_ai_formats::parse_codex_auth_identity(Some( + r#"{"account_id":"Workspace-1","account_user_id":"Member-1","email":"old@example.com"}"#, + )); + let canonical_b = aether_ai_formats::parse_codex_auth_identity(Some( + r#"{"account_id":"workspace-1","account_user_id":"member-1","email":"new@example.com"}"#, + )); + let other_member = aether_ai_formats::parse_codex_auth_identity(Some( + r#"{"account_id":"workspace-1","account_user_id":"member-2"}"#, + )); + assert_eq!( + resolve_codex_account_seed(&canonical_a, "key-a"), + resolve_codex_account_seed(&canonical_b, "key-b") + ); + assert_ne!( + resolve_codex_account_seed(&canonical_a, "key-a"), + resolve_codex_account_seed(&other_member, "key-a") + ); + + let derived_fingerprint = + aether_oauth::provider::providers::derive_codex_identity_fingerprint( + canonical_a.account_id.as_deref(), + canonical_a.account_user_id.as_deref(), + canonical_a.user_id.as_deref(), + canonical_a.email.as_deref(), + ) + .expect("canonical member fingerprint"); + let after_first_refresh = aether_ai_formats::parse_codex_auth_identity(Some( + &json!({ + "account_id": "Workspace-1", + "account_user_id": "Member-1", + "email": "old@example.com", + "codex_identity_fingerprint": derived_fingerprint + }) + .to_string(), + )); + let legacy_seed = resolve_codex_account_seed(&canonical_a, "key-a"); + let refreshed_seed = resolve_codex_account_seed(&after_first_refresh, "key-a"); + assert_eq!(legacy_seed, refreshed_seed); + + let context = CodexFingerprintConvergenceContext::new("logical-turn-1", 1_700_000_000_123) + .with_original_client_session_id("client-session-1") + .with_original_prompt_cache_key("client-cache-1"); + assert_eq!( + resolve_converged_fingerprint(&legacy_seed, &context), + resolve_converged_fingerprint(&refreshed_seed, &context) + ); } #[test] diff --git a/crates/aether-provider/transport/src/lib.rs b/crates/aether-provider/transport/src/lib.rs index 225d52123..287ef31b7 100644 --- a/crates/aether-provider/transport/src/lib.rs +++ b/crates/aether-provider/transport/src/lib.rs @@ -57,8 +57,10 @@ pub use auth::{build_passthrough_headers, ensure_upstream_auth_header}; pub use auth_config::apply_local_auth_config_header_overrides; pub use cache::{provider_transport_snapshot_looks_refreshed, ProviderTransportSnapshotCacheKey}; pub use codex_fingerprint::{ - apply_codex_oauth_fingerprint_convergence, codex_fingerprint_convergence_enabled, - CODEX_FINGERPRINT_CONFIG_NAMESPACE, CODEX_FINGERPRINT_ENABLED_CONFIG_KEY, + apply_codex_oauth_fingerprint_convergence, + apply_codex_oauth_fingerprint_convergence_with_context, codex_fingerprint_convergence_enabled, + CodexFingerprintConvergenceContext, CODEX_FINGERPRINT_CONFIG_NAMESPACE, + CODEX_FINGERPRINT_ENABLED_CONFIG_KEY, }; pub use conversion::{ candidate_common_transport_skip_reason, candidate_transport_pair_skip_reason,