From e742db9f7ba98e63e502b310857d49d047888518 Mon Sep 17 00:00:00 2001 From: jif Date: Wed, 28 Oct 2026 11:38:56 +0000 Subject: [PATCH] Preserve parent routing affinity for Guardian reviews (#57710) ## Why Shared routing cookies can change after a parent connection is established. Guardian reviews need the routing affinity captured with their parent response so unrelated traffic does not replace that hint. ## What changed - Capture `__oailb` and `__cflb` cookies from successful ChatGPT WebSocket handshakes and pair them with the parent response ID. - Retain the snapshot throughout a review and its retries, and apply it only to Guardian requests with the matching parent response. Reconnect cached reviewer sockets when the parent affinity changes. - Preserve cookie expiry and restrict reuse to the original secure origin and endpoint path. Replace routing cookies while retaining unrelated cookies, and keep the captured affinity out of prompts, history, and serialized metadata. ## Testing Add coverage for cookie scope, expiry, duplicate rejection, and merging; replaying a captured pin after shared cookies change or are deleted; and retaining the matching parent snapshot through compaction without exposing cookies in review prompts or metadata. GitOrigin-RevId: d180679e6265b612a3617f20ca1f5dec89e73aa8 --- codex-rs/codex-api/src/common.rs | 2 + .../src/endpoint/responses_websocket.rs | 12 +- .../endpoint/responses_websocket/connector.rs | 3 +- codex-rs/codex-api/src/lib.rs | 3 + codex-rs/codex-api/src/sse/responses.rs | 1 + codex-rs/core/src/client.rs | 33 ++++- codex-rs/core/src/client_common.rs | 1 + .../core/src/client_runtime_auth_tests.rs | 1 + codex-rs/core/src/client_tests.rs | 1 + codex-rs/core/src/compact_remote_v2.rs | 1 + codex-rs/core/src/guardian/mod.rs | 23 ++-- codex-rs/core/src/guardian/review_session.rs | 23 +++- .../core/src/guardian/review_session_tests.rs | 43 +++++++ codex-rs/core/src/guardian/routing.rs | 15 +++ codex-rs/core/src/model_request_tests.rs | 1 + codex-rs/core/src/responses_metadata.rs | 3 + codex-rs/core/src/session/session.rs | 11 ++ codex-rs/core/src/session/turn.rs | 6 + codex-rs/core/src/turn_metadata_tests.rs | 66 ++++++++++ .../src/chatgpt_cloudflare_cookies.rs | 2 +- .../src/chatgpt_routing_affinity.rs | 115 ++++++++++++++++++ .../src/chatgpt_routing_affinity_tests.rs | 100 +++++++++++++++ codex-rs/http-client/src/lib.rs | 2 + codex-rs/websocket-client/src/cookie_tests.rs | 55 +++++++-- codex-rs/websocket-client/src/lib.rs | 22 +++- 25 files changed, 512 insertions(+), 33 deletions(-) create mode 100644 codex-rs/core/src/guardian/routing.rs create mode 100644 codex-rs/http-client/src/chatgpt_routing_affinity.rs create mode 100644 codex-rs/http-client/src/chatgpt_routing_affinity_tests.rs diff --git a/codex-rs/codex-api/src/common.rs b/codex-rs/codex-api/src/common.rs index 1abc89c1b7..8696df4813 100644 --- a/codex-rs/codex-api/src/common.rs +++ b/codex-rs/codex-api/src/common.rs @@ -420,6 +420,8 @@ pub enum ResponseInterrupt { } pub struct ResponseStream { + /// Routing pin from the connection that produced these responses. Never model metadata. + pub routing_affinity: Option, pub rx_event: mpsc::Receiver>, /// Server-assigned `x-request-id` response header, when present. pub upstream_request_id: Option, diff --git a/codex-rs/codex-api/src/endpoint/responses_websocket.rs b/codex-rs/codex-api/src/endpoint/responses_websocket.rs index 27361aef0c..0c31969017 100644 --- a/codex-rs/codex-api/src/endpoint/responses_websocket.rs +++ b/codex-rs/codex-api/src/endpoint/responses_websocket.rs @@ -1,3 +1,4 @@ +use crate::RESPONSES_PATH; use crate::auth::SharedAuthProvider; use crate::common::ResponseEvent; use crate::common::ResponseInterrupt; @@ -55,6 +56,7 @@ use tungstenite::protocol::WebSocketConfig; use url::Url; struct WsStream { + routing_affinity: Option, tx_command: mpsc::Sender, rx_message: mpsc::UnboundedReceiver>, pump_task: tokio::task::JoinHandle<()>, @@ -69,6 +71,7 @@ enum WsCommand { impl WsStream { fn new(inner: WebSocketConnection) -> Self { + let routing_affinity = inner.routing_affinity().cloned(); let (tx_command, mut rx_command) = mpsc::channel::(32); let (tx_message, rx_message) = mpsc::unbounded_channel::>(); @@ -126,6 +129,7 @@ impl WsStream { }); Self { + routing_affinity, tx_command, rx_message, pump_task, @@ -193,6 +197,7 @@ struct ResponsesWebsocketTimingLogContext { } pub struct ResponsesWebsocketConnection { + routing_affinity: Option, stream: Arc>>, // TODO (pakrym): is this the right place for timeout? idle_timeout: Duration, @@ -215,18 +220,20 @@ impl std::fmt::Debug for ResponsesWebsocketConnection { impl ResponsesWebsocketConnection { fn new( - stream: WsStream, + mut stream: WsStream, idle_timeout: Duration, server_reasoning_included: bool, server_model: Option, telemetry: Option>, ) -> Self { + let routing_affinity = stream.routing_affinity.take(); Self { stream: Arc::new(Mutex::new(Some(stream))), idle_timeout, server_reasoning_included, server_model, telemetry, + routing_affinity, } } @@ -387,6 +394,7 @@ impl ResponsesWebsocketConnection { Ok(ResponseStream { rx_event, upstream_request_id: None, + routing_affinity: self.routing_affinity.clone(), interrupt: Some(tx_interrupt), }) } @@ -442,7 +450,7 @@ impl ResponsesWebsocketClient { ) -> Result { let ws_url = self .provider - .websocket_url_for_path("/responses") + .websocket_url_for_path(RESPONSES_PATH) .map_err(|err| { ResponsesWebsocketProbeError::before_upgrade( ApiError::Stream(format!("failed to build websocket URL: {err}")), diff --git a/codex-rs/codex-api/src/endpoint/responses_websocket/connector.rs b/codex-rs/codex-api/src/endpoint/responses_websocket/connector.rs index 4dffd362f1..cd4fc604c1 100644 --- a/codex-rs/codex-api/src/endpoint/responses_websocket/connector.rs +++ b/codex-rs/codex-api/src/endpoint/responses_websocket/connector.rs @@ -6,6 +6,7 @@ use super::ResponsesWebsocketConnection; use super::connect_websocket; use super::merge_request_headers; use crate::ApiError; +use crate::RESPONSES_PATH; use crate::WebsocketTelemetry; use codex_websocket_client::WebSocketConnector; use http::HeaderMap; @@ -31,7 +32,7 @@ impl ResponsesWebsocketClient { ) -> Result { let ws_url = self .provider - .websocket_url_for_path("/responses") + .websocket_url_for_path(RESPONSES_PATH) .map_err(|err| ApiError::Stream(format!("failed to build websocket URL: {err}")))?; let headers = merge_request_headers(&self.provider.headers, extra_headers, default_headers); let headers = self.auth.apply_auth_headers(headers).await?; diff --git a/codex-rs/codex-api/src/lib.rs b/codex-rs/codex-api/src/lib.rs index ad1a361a79..221010a44f 100644 --- a/codex-rs/codex-api/src/lib.rs +++ b/codex-rs/codex-api/src/lib.rs @@ -14,6 +14,9 @@ pub(crate) mod search; pub(crate) mod sse; pub(crate) mod telemetry; +/// Path of the Responses endpoint, relative to a provider's base URL. +pub const RESPONSES_PATH: &str = "/responses"; + pub use crate::requests::headers::build_session_headers; pub use codex_client::RequestTelemetry; pub use codex_client::ReqwestTransport; diff --git a/codex-rs/codex-api/src/sse/responses.rs b/codex-rs/codex-api/src/sse/responses.rs index 7f076a139c..40478f4088 100644 --- a/codex-rs/codex-api/src/sse/responses.rs +++ b/codex-rs/codex-api/src/sse/responses.rs @@ -98,6 +98,7 @@ pub fn spawn_response_stream( }); ResponseStream { + routing_affinity: None, rx_event, upstream_request_id, interrupt: None, diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 3308777b5a..2ac9b931ff 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -47,6 +47,7 @@ use codex_api::MemoriesClient as ApiMemoriesClient; use codex_api::MemorySummarizeInput as ApiMemorySummarizeInput; use codex_api::MemorySummarizeOutput as ApiMemorySummarizeOutput; use codex_api::Provider as ApiProvider; +use codex_api::RESPONSES_PATH; use codex_api::RawMemory as ApiRawMemory; use codex_api::RealtimeCallClient as ApiRealtimeCallClient; use codex_api::RealtimeSessionConfig as ApiRealtimeSessionConfig; @@ -360,6 +361,7 @@ struct WebsocketContinuation { struct WebsocketSession { connection: Option, responses_headers: ApiHeaderMap, + parent_routing_affinity: Option, runtime_id: Option, /// Owner of the cached state, including before a connection is opened. auth_owner_generation: Option, @@ -1425,9 +1427,26 @@ impl ModelClient { auth_context: AuthRequestTelemetryContext, request_route_telemetry: RequestRouteTelemetry, responses_headers: &ApiHeaderMap, + parent_routing_affinity: Option<&codex_http_client::ChatGptRoutingAffinity>, ) -> std::result::Result { let mut headers = self.build_websocket_headers(responses_metadata).await; headers.extend(responses_headers.clone()); + if let Some(affinity) = parent_routing_affinity { + let uri = api_provider + .websocket_url_for_path(RESPONSES_PATH) + .map_err(|err| ApiError::Stream(format!("failed to build websocket URL: {err}")))? + .as_str() + .parse::() + .map_err(|err| ApiError::Stream(format!("invalid websocket URI: {err}")))?; + let cookies = headers + .get(http::header::COOKIE) + .cloned() + .or_else(|| api_provider.headers.get(http::header::COOKIE).cloned()) + .or_else(|| http_client_factory.chatgpt_cookie_header(&uri)); + if let Some(cookies) = affinity.merge_cookie_header(&uri, cookies.as_ref()) { + headers.insert(http::header::COOKIE, cookies); + } + } let websocket_telemetry = ModelClientSession::build_websocket_telemetry( session_telemetry, auth_context.clone(), @@ -1798,12 +1817,19 @@ impl ModelClientSession { responses_headers, runtime_id, } = params; + let parent_routing_affinity = if is_guardian_reviewer(responses_headers) { + responses_metadata.parent_routing_affinity.clone() + } else { + None + }; let connection_key = ResponsesConnectionKey::new(&api_provider, auth_revision); let reset_reason = match self.websocket_session.connection.as_ref() { Some(_) if self.websocket_session.responses_headers != *responses_headers || self.websocket_session.connection_key.as_ref() != Some(&connection_key) - || self.websocket_session.runtime_id != runtime_id => + || self.websocket_session.runtime_id != runtime_id + || self.websocket_session.parent_routing_affinity + != parent_routing_affinity => { Some("other") } @@ -1835,9 +1861,11 @@ impl ModelClientSession { auth_context, request_route_telemetry, responses_headers, + parent_routing_affinity.as_ref(), ) .await?; self.websocket_session.connection = Some(new_conn); + self.websocket_session.parent_routing_affinity = parent_routing_affinity; self.websocket_session.responses_headers = responses_headers.clone(); self.websocket_session.runtime_id = runtime_id; self.websocket_session.auth_owner_generation = auth_owner_generation; @@ -2624,6 +2652,7 @@ fn map_response_stream( interceptors: Vec>, last_response_tx: Option>, ) -> ResponseStream { + let routing_affinity = api_stream.routing_affinity.take(); let upstream_request_id = api_stream.upstream_request_id.take(); let interrupt = api_stream.interrupt.take(); let mut stream = map_response_events( @@ -2635,6 +2664,7 @@ fn map_response_stream( last_response_tx, ); stream.interrupt = interrupt; + stream.routing_affinity = routing_affinity; stream } @@ -2801,6 +2831,7 @@ where }); ResponseStream { + routing_affinity: None, rx_event, interrupt: None, consumer_dropped: consumer_dropped_for_stream, diff --git a/codex-rs/core/src/client_common.rs b/codex-rs/core/src/client_common.rs index a19431c5de..81bf9422ea 100644 --- a/codex-rs/core/src/client_common.rs +++ b/codex-rs/core/src/client_common.rs @@ -116,6 +116,7 @@ fn normalize_image_detail(detail: &mut Option, model_info: &ModelIn } pub struct ResponseStream { + pub(crate) routing_affinity: Option, pub(crate) rx_event: mpsc::Receiver>, pub(crate) interrupt: Option>, /// Signals the mapper task that the consumer stopped polling before the diff --git a/codex-rs/core/src/client_runtime_auth_tests.rs b/codex-rs/core/src/client_runtime_auth_tests.rs index a1906c9504..35a41b68cd 100644 --- a/codex-rs/core/src/client_runtime_auth_tests.rs +++ b/codex-rs/core/src/client_runtime_auth_tests.rs @@ -117,6 +117,7 @@ async fn selected_runtime_transport_stays_revoked(mode: RuntimeAuthMode) { auth_context, RequestRouteTelemetry::for_endpoint("/responses"), &http::HeaderMap::new(), + /*parent_routing_affinity*/ None, ) .await; assert!(matches!( diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index e5dd102416..4f60892b53 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -2278,6 +2278,7 @@ async fn intercepted_output_reaches_trace_and_websocket_bookkeeping() -> anyhow: let (last_response_tx, last_response_rx) = tokio::sync::oneshot::channel(); let mut stream = super::map_response_stream( codex_api::ResponseStream { + routing_affinity: None, rx_event, upstream_request_id: None, interrupt: None, diff --git a/codex-rs/core/src/compact_remote_v2.rs b/codex-rs/core/src/compact_remote_v2.rs index bacf8c3a5c..85a3e6c91d 100644 --- a/codex-rs/core/src/compact_remote_v2.rs +++ b/codex-rs/core/src/compact_remote_v2.rs @@ -835,6 +835,7 @@ mod tests { } drop(tx_event); ResponseStream { + routing_affinity: None, rx_event, interrupt: None, consumer_dropped: CancellationToken::new(), diff --git a/codex-rs/core/src/guardian/mod.rs b/codex-rs/core/src/guardian/mod.rs index 9737236bc1..a714e4dd40 100644 --- a/codex-rs/core/src/guardian/mod.rs +++ b/codex-rs/core/src/guardian/mod.rs @@ -21,6 +21,9 @@ mod review; mod review_session; mod reviewer_config; pub(crate) use reviewer_config::resolve_review_model; +mod routing; +pub(crate) use routing::ParentResponse; +pub(crate) use routing::PendingReviewRouting; mod runtime; #[cfg(test)] pub(crate) mod test_host; @@ -83,8 +86,8 @@ pub(crate) use codex_guardian_reviewer::GUARDIAN_MAX_NODE_REPL_TOOL_RESULT_TOKEN /// MCP elicitations use their live invocation's issuing step when available, otherwise turn inputs. #[derive(Clone)] pub(crate) struct GuardianReviewContext { - /// The latest response ID received in this turn when review was requested. - pub(crate) parent_response_id: Option, + /// The latest response ID and its connection affinity when review was requested. + pub(crate) parent_response: Option>, turn: Arc, environments: TurnEnvironmentSnapshot, project_instructions: Option>, @@ -126,10 +129,7 @@ impl GuardianReviewContext { environments: &TurnEnvironmentSnapshot, ) -> Self { Self { - parent_response_id: turn - .extension_data - .get::() - .map(|id| id.0.clone()), + parent_response: turn.extension_data.get::(), environments: environments.clone(), project_instructions: None, model_info: Arc::clone(&settings.model_info), @@ -154,11 +154,7 @@ impl GuardianReviewContext { impl From<&Arc> for GuardianReviewContext { fn from(step: &Arc) -> Self { Self { - parent_response_id: step - .turn - .extension_data - .get::() - .map(|id| id.0.clone()), + parent_response: step.turn.extension_data.get::(), turn: Arc::clone(&step.turn), environments: step.environments.clone(), project_instructions: Some(step.loaded_agents_md.clone().unwrap_or_default()), @@ -175,10 +171,7 @@ impl From<&Arc> for GuardianReviewContext { impl From> for GuardianReviewContext { fn from(turn: Arc) -> Self { Self { - parent_response_id: turn - .extension_data - .get::() - .map(|id| id.0.clone()), + parent_response: turn.extension_data.get::(), environments: turn.initial_environments.clone(), project_instructions: None, model_info: Arc::clone(turn.model_info()), diff --git a/codex-rs/core/src/guardian/review_session.rs b/codex-rs/core/src/guardian/review_session.rs index 9cdc141584..a0aa858219 100644 --- a/codex-rs/core/src/guardian/review_session.rs +++ b/codex-rs/core/src/guardian/review_session.rs @@ -632,6 +632,13 @@ async fn run_review_on_session( .insert(super::input_budget::PendingReviewContext( prompt_items.context, )); + review_session + .session + .services + .thread_extension_data + .insert(super::PendingReviewRouting( + params.parent_context.parent_response.clone(), + )); let request = codex_guardian_reviewer::ReviewerTurn { items, environments: codex_protocol::protocol::TurnEnvironmentSelections::new( @@ -644,7 +651,11 @@ async fn run_review_on_session( model: review_model.model.clone(), reasoning_effort: review_model.reasoning_effort.clone(), cyber_access_program: parent_turn.cyber_access_program, - parent_response_id: params.parent_context.parent_response_id.clone(), + parent_response_id: params + .parent_context + .parent_response + .as_ref() + .map(|parent| parent.id.clone()), schema: params.schema.clone(), parent_turn_id: parent_turn.sub_id.clone(), root_turn_id: parent_turn.turn_metadata_state.root_turn_id(), @@ -660,6 +671,11 @@ async fn run_review_on_session( { Ok(turn_id) => turn_id, Err(outcome) => { + review_session + .session + .services + .thread_extension_data + .remove::(); review_session .session .services @@ -688,6 +704,11 @@ async fn run_review_on_session( &mut analytics_result, ) .await; + review_session + .session + .services + .thread_extension_data + .remove::(); // Finalizing input consumes this context; failures before admission leave it pending. let transcript_admitted = review_session .session diff --git a/codex-rs/core/src/guardian/review_session_tests.rs b/codex-rs/core/src/guardian/review_session_tests.rs index 252cd43235..db4218b685 100644 --- a/codex-rs/core/src/guardian/review_session_tests.rs +++ b/codex-rs/core/src/guardian/review_session_tests.rs @@ -46,6 +46,21 @@ async fn run_review_preserves_evidence_during_parent_compaction() { ) .unwrap(); params.parent_session = Arc::clone(&parent); + let parent_uri = "wss://review-snapshot.chatgpt.com/backend-api/codex/responses" + .parse() + .unwrap(); + let mut headers = http::HeaderMap::new(); + headers.insert( + http::header::SET_COOKIE, + http::HeaderValue::from_static("__oailb=original-parent; Path=/; Secure; Max-Age=3600"), + ); + turn.extension_data.insert(super::super::ParentResponse { + id: "original-response".to_owned(), + routing_affinity: codex_http_client::ChatGptRoutingAffinity::from_response( + &parent_uri, + &headers, + ), + }); params.parent_context = GuardianReviewContext::from(Arc::clone(&turn)); let evidence: ResponseItem = serde_json::from_value(serde_json::json!({ "type": "function_call_output", "call_id": "prior-inspection", "output": EVIDENCE @@ -73,9 +88,14 @@ async fn run_review_preserves_evidence_during_parent_compaction() { .node_repl_auto_review_required, ) .with_node_repl_policy(¶ms.node_repl_policy); + let reviewer_session = Arc::clone(&reviewer.session); let manager = prewarm_test_session(¶ms, reviewer).await; // Capture the review context, then compact the parent before the reviewer builds its prompt. let prepared = setup::prepare_review(params).await.unwrap(); + turn.extension_data.insert(super::super::ParentResponse { + id: "later-response".to_owned(), + routing_affinity: None, + }); let checkpoint: ResponseItem = serde_json::from_value(serde_json::json!({ "type": "compaction", "id": "cmp_new", "encrypted_content": "new-checkpoint" })) @@ -99,6 +119,21 @@ async fn run_review_preserves_evidence_during_parent_compaction() { .await; let ((outcome, _), submitted_text) = tokio::join!(manager.review(prepared), async { let submission = rx_sub.recv().await.unwrap(); + let pending = reviewer_session + .services + .thread_extension_data + .get::() + .unwrap(); + let parent = pending.0.as_ref().unwrap(); + assert_eq!(parent.id, "original-response"); + assert_eq!( + parent + .routing_affinity + .as_ref() + .unwrap() + .cookie_header(&parent_uri), + Some(http::HeaderValue::from_static("__oailb=original-parent")) + ); let id = submission.id; let Op::TurnInput { request, reply, .. } = submission.op else { panic!("expected reviewer prompt"); @@ -123,6 +158,14 @@ async fn run_review_preserves_evidence_during_parent_compaction() { .unwrap(); text }); + assert!( + reviewer_session + .services + .thread_extension_data + .get::() + .is_none() + ); + assert!(!submitted_text.contains("__oailb")); assert!(matches!( outcome, GuardianReviewSessionOutcome::Completed(Ok(_)) diff --git a/codex-rs/core/src/guardian/routing.rs b/codex-rs/core/src/guardian/routing.rs new file mode 100644 index 0000000000..3d8f9e0432 --- /dev/null +++ b/codex-rs/core/src/guardian/routing.rs @@ -0,0 +1,15 @@ +//! Pairs a parent response with its connection's routing pin. Reviewers retain this snapshot +//! through queueing/retries; the cookie never enters serialized request metadata or history. + +use std::sync::Arc; + +use codex_http_client::ChatGptRoutingAffinity; + +#[derive(Clone, Debug)] +pub(crate) struct ParentResponse { + pub(crate) id: String, + pub(crate) routing_affinity: Option, +} + +/// Installed for the exclusive lifetime of one review, including its inference retries. +pub(crate) struct PendingReviewRouting(pub(crate) Option>); diff --git a/codex-rs/core/src/model_request_tests.rs b/codex-rs/core/src/model_request_tests.rs index 84ab5014d1..319e97bd00 100644 --- a/codex-rs/core/src/model_request_tests.rs +++ b/codex-rs/core/src/model_request_tests.rs @@ -30,6 +30,7 @@ async fn interceptors_forward_before_completion_in_order_and_cancel_upstream() { let mut stream = intercept_stream( Box::pin( ResponseStream { + routing_affinity: None, rx_event, interrupt: None, consumer_dropped, diff --git a/codex-rs/core/src/responses_metadata.rs b/codex-rs/core/src/responses_metadata.rs index 2d82bb719e..27bd6e4e0c 100644 --- a/codex-rs/core/src/responses_metadata.rs +++ b/codex-rs/core/src/responses_metadata.rs @@ -238,6 +238,8 @@ pub struct CodexResponsesMetadata { pub(crate) guardian_recorded_items: bool, /// Guardian parent reference; projected only onto a Guardian request. pub(crate) parent_response_id: Option, + /// Transport-only parent pin; never projected into headers or body metadata. + pub(crate) parent_routing_affinity: Option, pub(crate) installation_id: String, pub(crate) session_id: String, pub(crate) thread_id: String, @@ -284,6 +286,7 @@ impl CodexResponsesMetadata { ) -> Self { Self { parent_response_id: None, + parent_routing_affinity: None, guardian_recorded_items: false, installation_id, session_id, diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 8e2d7b31c8..a41f1471e4 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -725,6 +725,17 @@ impl Session { ); ExecutionMetadata::from_settings(&step_context.settings).apply_to(&mut responses_metadata); responses_metadata.guardian_recorded_items = self.parent_observer_event_tx.get().is_some(); + if responses_metadata.subagent_header.as_deref() + == Some(crate::guardian::GUARDIAN_REVIEWER_NAME) + && let Some(pending) = self + .services + .thread_extension_data + .get::() + && let Some(parent) = &pending.0 + && responses_metadata.parent_response_id.as_deref() == Some(parent.id.as_str()) + { + responses_metadata.parent_routing_affinity = parent.routing_affinity.clone(); + } responses_metadata.tool_namespaces_info = if step_context .turn .config diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 5cd4d77e0e..7e40cb4a1a 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -2924,6 +2924,12 @@ async fn try_run_sampling_request( match event { ResponseEvent::Created { response_id } => { if let Some(response_id) = response_id { + turn_context + .extension_data + .insert(crate::guardian::ParentResponse { + id: response_id.clone(), + routing_affinity: stream.routing_affinity.clone(), + }); turn_context .extension_data .insert(codex_api::ResponseId(response_id)); diff --git a/codex-rs/core/src/turn_metadata_tests.rs b/codex-rs/core/src/turn_metadata_tests.rs index 0987dd3468..2d8aad6e15 100644 --- a/codex-rs/core/src/turn_metadata_tests.rs +++ b/codex-rs/core/src/turn_metadata_tests.rs @@ -608,6 +608,72 @@ async fn execution_metadata_applies_captured_settings_to_responses_metadata() { assert_eq!(metadata[NODE_REPL_DISABLED_KEY], true); } +#[tokio::test] +async fn guardian_affinity_requires_the_captured_parent_and_stays_out_of_metadata() { + let uri = "wss://metadata-affinity.chatgpt.com/backend-api/codex/responses" + .parse() + .unwrap(); + let mut headers = http::HeaderMap::new(); + headers.insert( + http::header::SET_COOKIE, + http::HeaderValue::from_static("__oailb=private-parent-pin; Path=/; Secure; Max-Age=3600"), + ); + let affinity = + codex_http_client::ChatGptRoutingAffinity::from_response(&uri, &headers).unwrap(); + for (subagent, request_parent, expected) in [ + (Some("guardian"), "parent", true), + (Some("guardian"), "other-parent", false), + (None, "parent", false), + ] { + let (session, mut turn) = make_session_and_context().await; + Arc::get_mut(&mut turn.turn_metadata_state) + .unwrap() + .subagent_header = subagent.map(str::to_owned); + turn.turn_metadata_state + .update_responsesapi_client_metadata( + codex_protocol::turn_input::ResponsesApiClientMetadataUpdate { + replacement: Some(HashMap::from([( + "parent_response_id".to_owned(), + request_parent.to_owned(), + )])), + patch: None, + }, + ); + session + .services + .thread_extension_data + .insert(crate::guardian::PendingReviewRouting(Some(Arc::new( + crate::guardian::ParentResponse { + id: "parent".to_owned(), + routing_affinity: Some(affinity.clone()), + }, + )))); + let step = StepContext::for_test(Arc::new(turn)); + let metadata = session + .responses_metadata(&step, CodexResponsesRequestKind::Turn) + .await; + assert_eq!( + metadata + .parent_routing_affinity + .as_ref() + .and_then(|pin| pin.cookie_header(&uri)), + expected.then(|| http::HeaderValue::from_static("__oailb=private-parent-pin")) + ); + assert!( + !metadata + .turn_metadata_json() + .unwrap() + .contains("private-parent-pin") + ); + assert!( + !serde_json::to_string(&metadata.client_metadata(/*include_internal*/ true)) + .unwrap() + .contains("private-parent-pin") + ); + assert!(!format!("{:?}", metadata.compatibility_headers()).contains("private-parent-pin")); + } +} + #[test] fn turn_metadata_state_marks_user_input_requested_during_turn_only_for_mcp_request_meta() { let temp_dir = TempDir::new().expect("temp dir"); diff --git a/codex-rs/http-client/src/chatgpt_cloudflare_cookies.rs b/codex-rs/http-client/src/chatgpt_cloudflare_cookies.rs index 847bad77c5..54db221bbf 100644 --- a/codex-rs/http-client/src/chatgpt_cloudflare_cookies.rs +++ b/codex-rs/http-client/src/chatgpt_cloudflare_cookies.rs @@ -126,7 +126,7 @@ impl HttpClientFactory { } } -fn chatgpt_cookie_url(uri: &Uri) -> Option { +pub(crate) fn chatgpt_cookie_url(uri: &Uri) -> Option { let mut url = reqwest::Url::parse(&uri.to_string()).ok()?; // A secure WebSocket handshake has the same cookie scope as HTTPS. if url.scheme() == "wss" { diff --git a/codex-rs/http-client/src/chatgpt_routing_affinity.rs b/codex-rs/http-client/src/chatgpt_routing_affinity.rs new file mode 100644 index 0000000000..21522aa858 --- /dev/null +++ b/codex-rs/http-client/src/chatgpt_routing_affinity.rs @@ -0,0 +1,115 @@ +//! Connection-owned ChatGPT routing pins. Keep only the handshake's routing cookies, +//! with their original expiry, and never forward them outside the original endpoint. + +use std::sync::Arc; + +use http::HeaderMap; +use http::HeaderValue; +use http::Uri; +use http::header::SET_COOKIE; +use reqwest::cookie::CookieStore; +use reqwest::cookie::Jar; +use url::Url; + +use crate::chatgpt_cloudflare_cookies::chatgpt_cookie_url; + +const ROUTING_COOKIE_NAMES: [&str; 2] = ["__oailb", "__cflb"]; + +/// A routing pin returned by one connection, independent of the process-wide cookie jar. +/// This is a best-effort hint: the edge still validates the cookie and may fail over. +/// Equality identifies the captured connection, independently of cookie expiry, so a live +/// colocated reviewer socket need not be disconnected when its handshake cookie expires. +#[derive(Clone)] +pub struct ChatGptRoutingAffinity { + url: Url, + jar: Arc, +} + +impl std::fmt::Debug for ChatGptRoutingAffinity { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ChatGptRoutingAffinity") + .finish_non_exhaustive() + } +} + +impl PartialEq for ChatGptRoutingAffinity { + fn eq(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.jar, &other.jar) + } +} + +impl Eq for ChatGptRoutingAffinity {} + +impl ChatGptRoutingAffinity { + /// Captures valid `__oailb` and Cloudflare's native `__cflb` response cookies. + /// Ambiguous or oversized cookies reject the pin; request cookies alone cannot establish + /// where failover landed. Cookie scope and expiry remain owned by the private jar. + pub fn from_response(uri: &Uri, headers: &HeaderMap) -> Option { + let url = chatgpt_cookie_url(uri)?; + let mut cookies = Vec::with_capacity(2); + for name in ROUTING_COOKIE_NAMES { + let mut matches = headers.get_all(SET_COOKIE).iter().filter(|header| { + header + .to_str() + .ok() + .and_then(|value| value.split_once('=')) + .is_some_and(|(cookie_name, _)| cookie_name.trim() == name) + }); + if let Some(cookie) = matches.next() { + if cookie.as_bytes().len() > 8192 || matches.next().is_some() { + return None; + } + cookies.push(cookie); + } + } + let jar = Jar::default(); + jar.set_cookies(&mut cookies.into_iter(), &url); + jar.cookies(&url)?; + Some(Self { + url, + jar: Arc::new(jar), + }) + } + + /// Returns the still-valid pin only for the same secure ChatGPT origin and endpoint path. + /// The private jar retains Max-Age/Expires and is never refreshed by unrelated traffic. + pub fn cookie_header(&self, uri: &Uri) -> Option { + let url = chatgpt_cookie_url(uri)?; + if url.origin() != self.url.origin() || url.path() != self.url.path() { + return None; + } + let mut cookie = self.jar.cookies(&url)?; + cookie.set_sensitive(true); + Some(cookie) + } + + /// Replaces routing cookies with the scoped pin while preserving unrelated cookies. + /// Returns `None` when the pin is no longer valid for this endpoint. + pub fn merge_cookie_header( + &self, + uri: &Uri, + existing: Option<&HeaderValue>, + ) -> Option { + let pin = self.cookie_header(uri)?; + let mut parts = existing + .and_then(|value| value.to_str().ok()) + .into_iter() + .flat_map(|value| value.split(';')) + .map(str::trim) + .filter(|part| { + !part.is_empty() + && !part + .split_once('=') + .is_some_and(|(name, _)| ROUTING_COOKIE_NAMES.contains(&name.trim())) + }) + .collect::>(); + parts.push(pin.to_str().ok()?); + let mut cookies = HeaderValue::from_str(&parts.join("; ")).ok()?; + cookies.set_sensitive(true); + Some(cookies) + } +} + +#[cfg(test)] +#[path = "chatgpt_routing_affinity_tests.rs"] +mod tests; diff --git a/codex-rs/http-client/src/chatgpt_routing_affinity_tests.rs b/codex-rs/http-client/src/chatgpt_routing_affinity_tests.rs new file mode 100644 index 0000000000..5111fc3a42 --- /dev/null +++ b/codex-rs/http-client/src/chatgpt_routing_affinity_tests.rs @@ -0,0 +1,100 @@ +use super::*; +use pretty_assertions::assert_eq; + +#[test] +fn affinity_retains_only_the_response_pin_with_its_cookie_scope_and_expiry() { + let uri: Uri = "wss://affinity.chatgpt.com/backend-api/codex/responses" + .parse() + .unwrap(); + let mut headers = HeaderMap::new(); + headers.append( + SET_COOKIE, + HeaderValue::from_static("__oailb=parent; Path=/backend-api; Max-Age=3600; Secure"), + ); + headers.append( + SET_COOKIE, + HeaderValue::from_static("__cflb=parent-native; Path=/backend-api; Max-Age=3600; Secure"), + ); + headers.append( + SET_COOKIE, + HeaderValue::from_static("chatgpt_session=secret; Path=/; Secure"), + ); + let affinity = ChatGptRoutingAffinity::from_response(&uri, &headers).unwrap(); + let pin = affinity.cookie_header(&uri).unwrap(); + let mut cookies = pin.to_str().unwrap().split("; ").collect::>(); + cookies.sort_unstable(); + assert_eq!(cookies, ["__cflb=parent-native", "__oailb=parent"]); + assert!(pin.is_sensitive()); + let existing = HeaderValue::from_static( + "__oailb=old; cf_clearance=keep; __cflb=old-native; configured=keep", + ); + let merged = affinity.merge_cookie_header(&uri, Some(&existing)).unwrap(); + let mut cookies = merged.to_str().unwrap().split("; ").collect::>(); + cookies.sort_unstable(); + assert_eq!( + cookies, + [ + "__cflb=parent-native", + "__oailb=parent", + "cf_clearance=keep", + "configured=keep", + ] + ); + assert!(merged.is_sensitive()); + assert_eq!( + affinity.merge_cookie_header(&uri, /*existing*/ None), + Some(pin.clone()) + ); + assert_eq!(format!("{affinity:?}"), "ChatGptRoutingAffinity { .. }"); + for outside_scope in [ + "ws://affinity.chatgpt.com/backend-api/codex/responses", + "wss://other.chatgpt.com/backend-api/codex/responses", + "wss://affinity.chatgpt.com:444/backend-api/codex/responses", + "wss://affinity.chatgpt.com/backend-api/codex/other", + "wss://api.openai.com/backend-api/codex/responses", + ] { + assert_eq!( + affinity.cookie_header(&outside_scope.parse().unwrap()), + None + ); + assert_eq!( + affinity.merge_cookie_header(&outside_scope.parse().unwrap(), Some(&existing)), + None + ); + } + for name in ["__oailb", "__cflb"] { + headers.insert( + SET_COOKIE, + HeaderValue::from_str(&format!("{name}=only; Path=/; Secure")).unwrap(), + ); + assert_eq!( + ChatGptRoutingAffinity::from_response(&uri, &headers) + .unwrap() + .cookie_header(&uri), + Some(HeaderValue::from_str(&format!("{name}=only")).unwrap()) + ); + for invalid in [ + "__oailb=deleted; Path=/; Max-Age=0; Secure", + "__oailb=expired; Path=/; Expires=Thu, 01 Jan 1970 00:00:00 GMT; Secure", + "__oailb=wrong-path; Path=/elsewhere; Secure", + "__oailb=wrong-domain; Domain=other.chatgpt.com; Path=/; Secure", + "chatgpt_session=secret; Path=/; Secure", + ] { + headers.insert( + SET_COOKIE, + HeaderValue::from_str(&invalid.replace("__oailb", name)).unwrap(), + ); + assert!(ChatGptRoutingAffinity::from_response(&uri, &headers).is_none()); + } + headers.insert( + SET_COOKIE, + HeaderValue::from_str(&format!("{name}=one; Path=/; Secure")).unwrap(), + ); + headers.append( + SET_COOKIE, + HeaderValue::from_str(&format!("{name}=two; Path=/; Secure")).unwrap(), + ); + assert!(ChatGptRoutingAffinity::from_response(&uri, &headers).is_none()); + } + assert_eq!(affinity.cookie_header(&uri), Some(pin)); +} diff --git a/codex-rs/http-client/src/lib.rs b/codex-rs/http-client/src/lib.rs index db66e6b18b..9eec07dd39 100644 --- a/codex-rs/http-client/src/lib.rs +++ b/codex-rs/http-client/src/lib.rs @@ -3,6 +3,7 @@ use url::Url; mod chatgpt_cloudflare_cookies; mod chatgpt_hosts; +mod chatgpt_routing_affinity; mod client; mod client_builder; mod client_tls; @@ -25,6 +26,7 @@ mod transport; pub use crate::chatgpt_cloudflare_cookies::with_chatgpt_cloudflare_cookie_store; pub use crate::chatgpt_hosts::is_allowed_chatgpt_host; +pub use crate::chatgpt_routing_affinity::ChatGptRoutingAffinity; pub use crate::client::HttpClient; pub use crate::client::HttpError; pub use crate::client_builder::HttpClientBuilder; diff --git a/codex-rs/websocket-client/src/cookie_tests.rs b/codex-rs/websocket-client/src/cookie_tests.rs index 1d48eab20c..2db2f8275f 100644 --- a/codex-rs/websocket-client/src/cookie_tests.rs +++ b/codex-rs/websocket-client/src/cookie_tests.rs @@ -63,7 +63,7 @@ async fn websocket_handshakes_share_routing_cookies_and_respect_cookie_scope() { let parent_url = format!("wss://{host}/backend-api/codex/responses"); let review_url = format!("wss://{host}/backend-api/codex/guardian"); let requests = [ - (parent_url, None), + (parent_url.clone(), None), (review_url.clone(), None), (review_url.clone(), Some("explicit=keep")), (format!("wss://{host}/outside"), None), @@ -77,19 +77,24 @@ async fn websocket_handshakes_share_routing_cookies_and_respect_cookie_scope() { (review_url.clone(), None), (review_url.clone(), None), (review_url, None), + (parent_url, None), ]; let responses: Vec<(StatusCode, Vec<&str>)> = vec![ ( StatusCode::SWITCHING_PROTOCOLS, vec![ "__oailb=west; Path=/backend-api; Secure; HttpOnly", + "__cflb=west; Path=/backend-api; Secure; HttpOnly", "chatgpt_session=never-store; Path=/; Secure", ], ), // Even a failed upgrade can refresh a cookie used by the next connection. ( StatusCode::FORBIDDEN, - vec!["__oailb=east; Path=/backend-api; Secure"], + vec![ + "__oailb=east; Path=/backend-api; Secure", + "__cflb=east; Path=/backend-api; Secure", + ], ), (StatusCode::SWITCHING_PROTOCOLS, vec![]), (StatusCode::SWITCHING_PROTOCOLS, vec![]), @@ -106,10 +111,15 @@ async fn websocket_handshakes_share_routing_cookies_and_respect_cookie_scope() { (StatusCode::SWITCHING_PROTOCOLS, vec![]), ( StatusCode::SWITCHING_PROTOCOLS, - vec!["__oailb=; Path=/backend-api; Max-Age=0; Secure"], + vec![ + "__oailb=; Path=/backend-api; Max-Age=0; Secure", + "__cflb=; Path=/backend-api; Max-Age=0; Secure", + ], ), (StatusCode::SWITCHING_PROTOCOLS, vec![]), ]; + let mut responses = responses; + responses.push((StatusCode::SWITCHING_PROTOCOLS, vec![])); let expected_statuses = responses .iter() .map(|(status, _)| *status) @@ -165,7 +175,9 @@ async fn websocket_handshakes_share_routing_cookies_and_respect_cookie_scope() { }); let mut actual_statuses = Vec::new(); - for (url, explicit_cookie) in requests { + let mut parent_affinity: Option = None; + let last_request = requests.len() - 1; + for (index, (url, explicit_cookie)) in requests.into_iter().enumerate() { // Separate connector/factory instances must still see the shared routing cookie. let connector = WebSocketConnector { http_client_factory: HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), @@ -178,6 +190,15 @@ async fn websocket_handshakes_share_routing_cookies_and_respect_cookie_scope() { .headers_mut() .insert(COOKIE, HeaderValue::from_static(cookie)); } + // Replay the parent's captured pin after other traffic changed and deleted the shared pin. + if index == last_request { + let pin = parent_affinity + .as_ref() + .unwrap() + .cookie_header(request.uri()) + .unwrap(); + request.headers_mut().insert(COOKIE, pin); + } let result = connector .connect_with_route( request, @@ -191,6 +212,12 @@ async fn websocket_handshakes_share_routing_cookies_and_respect_cookie_scope() { .await; let status = match result { Ok((connection, response)) => { + if index == 0 { + parent_affinity = connection.routing_affinity().cloned(); + } else { + // All later successful handshakes return no usable routing cookie. + assert!(connection.routing_affinity().is_none()); + } drop(connection); response.status() } @@ -201,19 +228,29 @@ async fn websocket_handshakes_share_routing_cookies_and_respect_cookie_scope() { } assert_eq!(actual_statuses, expected_statuses); assert_eq!( - server.await.unwrap(), + server + .await + .unwrap() + .into_iter() + .map(|cookie| cookie.map(|value| { + let mut parts = value.to_str().unwrap().split("; ").collect::>(); + parts.sort_unstable(); + HeaderValue::from_str(&parts.join("; ")).unwrap() + })) + .collect::>(), [ None, - Some("__oailb=west"), + Some("__cflb=west; __oailb=west"), Some("explicit=keep"), None, None, None, None, None, - Some("__oailb=east"), - Some("__oailb=east"), - None + Some("__cflb=east; __oailb=east"), + Some("__cflb=east; __oailb=east"), + None, + Some("__cflb=west; __oailb=west"), ] .map(|cookie| cookie.map(HeaderValue::from_static)) .to_vec(), diff --git a/codex-rs/websocket-client/src/lib.rs b/codex-rs/websocket-client/src/lib.rs index 98c92114ef..4bd8513fe7 100644 --- a/codex-rs/websocket-client/src/lib.rs +++ b/codex-rs/websocket-client/src/lib.rs @@ -194,10 +194,19 @@ impl WebSocketConnector { .store_chatgpt_response_cookies(&uri, response.headers()), Err(_) => {} } - result + result.map(|(inner, response)| { + let routing_affinity = codex_http_client::ChatGptRoutingAffinity::from_response( + &uri, + response.headers(), + ); + (inner, response, routing_affinity) + }) }; - let (inner, response) = permit.run(connect).await.map_err(policy_error)??; - Ok((WebSocketConnection::new(inner, permit), response)) + let (inner, response, routing_affinity) = + permit.run(connect).await.map_err(policy_error)??; + let mut connection = WebSocketConnection::new(inner, permit); + connection.routing_affinity = routing_affinity; + Ok((connection, response)) } } @@ -221,6 +230,7 @@ fn is_loopback_destination(uri: &Uri) -> bool { /// without knowing which concrete network stream route selection produced. Independent revocation /// waits preserve both task wakers when the stream and sink are split. pub struct WebSocketConnection { + routing_affinity: Option, inner: Option, permit: NetworkPermit, read_revoked: BoxFuture<'static, ()>, @@ -245,6 +255,7 @@ impl WebSocketConnection { let read_permit = permit.clone(); let write_permit = permit.clone(); Self { + routing_affinity: None, inner: Some(inner), permit, read_revoked: async move { read_permit.revoked().await }.boxed(), @@ -253,6 +264,11 @@ impl WebSocketConnection { } } + /// The scoped routing cookie from this connection's handshake, never account cookies. + pub fn routing_affinity(&self) -> Option<&codex_http_client::ChatGptRoutingAffinity> { + self.routing_affinity.as_ref() + } + fn poll_policy( &mut self, context: &mut Context<'_>,