openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-29 19:08:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

Preserve parent routing affinity for Guardian reviews (#57710)

Browse Source 
## 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
This commit is contained in:
jif
2026-10-28 11:48:00 +00:00
committed by copyberry
parent e6de6a2423
commit e742db9f7b
25 changed files with 512 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/codex-api/src/common.rs
+2
View File
@@ -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<codex_http_client::ChatGptRoutingAffinity>,
pub rx_event: mpsc::Receiver<Result<ResponseEvent, ApiError>>,
/// Server-assigned `x-request-id` response header, when present.
pub upstream_request_id: Option<String>,
codex-rs/codex-api/src/endpoint/responses_websocket.rs
+10 -2
View File
@@ -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<codex_http_client::ChatGptRoutingAffinity>,
tx_command: mpsc::Sender<WsCommand>,
rx_message: mpsc::UnboundedReceiver<Result<Message, WsError>>,
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::<WsCommand>(32);
let (tx_message, rx_message) = mpsc::unbounded_channel::<Result<Message, WsError>>();
@@ -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<codex_http_client::ChatGptRoutingAffinity>,
stream: Arc<Mutex<Option<WsStream>>>,
// 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<String>,
telemetry: Option<Arc<dyn WebsocketTelemetry>>,
) -> 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<ResponsesWebsocketProbe, ResponsesWebsocketProbeError> {
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}")),
codex-rs/codex-api/src/endpoint/responses_websocket/connector.rs
+2 -1
View File
@@ -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<ResponsesWebsocketConnection, ApiError> {
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?;
codex-rs/codex-api/src/lib.rs
+3
View File
@@ -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;
codex-rs/codex-api/src/sse/responses.rs
+1
View File
@@ -98,6 +98,7 @@ pub fn spawn_response_stream(
});
ResponseStream {
routing_affinity: None,
rx_event,
upstream_request_id,
interrupt: None,
codex-rs/core/src/client.rs
+32 -1
View File
@@ -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<ApiWebSocketConnection>,
responses_headers: ApiHeaderMap,
parent_routing_affinity: Option<codex_http_client::ChatGptRoutingAffinity>,
runtime_id: Option<String>,
/// Owner of the cached state, including before a connection is opened.
auth_owner_generation: Option<u64>,
@@ -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<ApiWebSocketConnection, ApiError> {
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::<http::Uri>()
.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<Box<dyn codex_extension_api::ModelResponseInterceptor>>,
last_response_tx: Option<oneshot::Sender<LastResponse>>,
) -> 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,
codex-rs/core/src/client_common.rs
+1
View File
@@ -116,6 +116,7 @@ fn normalize_image_detail(detail: &mut Option<ImageDetail>, model_info: &ModelIn
}
pub struct ResponseStream {
pub(crate) routing_affinity: Option<codex_http_client::ChatGptRoutingAffinity>,
pub(crate) rx_event: mpsc::Receiver<Result<ResponseEvent>>,
pub(crate) interrupt: Option<oneshot::Sender<codex_api::ResponseInterrupt>>,
/// Signals the mapper task that the consumer stopped polling before the
codex-rs/core/src/client_runtime_auth_tests.rs
+1
View File
@@ -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!(
codex-rs/core/src/client_tests.rs
+1
View File
@@ -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,
codex-rs/core/src/compact_remote_v2.rs
+1
View File
@@ -835,6 +835,7 @@ mod tests {
}
drop(tx_event);
ResponseStream {
routing_affinity: None,
rx_event,
interrupt: None,
consumer_dropped: CancellationToken::new(),
codex-rs/core/src/guardian/mod.rs
+8 -15
View File
@@ -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<String>,
/// The latest response ID and its connection affinity when review was requested.
pub(crate) parent_response: Option<Arc<ParentResponse>>,
turn: Arc<TurnContext>,
environments: TurnEnvironmentSnapshot,
project_instructions: Option<Arc<LoadedAgentsMd>>,
@@ -126,10 +129,7 @@ impl GuardianReviewContext {
environments: &TurnEnvironmentSnapshot,
) -> Self {
Self {
parent_response_id: turn
.extension_data
.get::<codex_api::ResponseId>()
.map(|id| id.0.clone()),
parent_response: turn.extension_data.get::<ParentResponse>(),
environments: environments.clone(),
project_instructions: None,
model_info: Arc::clone(&settings.model_info),
@@ -154,11 +154,7 @@ impl GuardianReviewContext {
impl From<&Arc<StepContext>> for GuardianReviewContext {
fn from(step: &Arc<StepContext>) -> Self {
Self {
parent_response_id: step
.turn
.extension_data
.get::<codex_api::ResponseId>()
.map(|id| id.0.clone()),
parent_response: step.turn.extension_data.get::<ParentResponse>(),
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<StepContext>> for GuardianReviewContext {
impl From<Arc<TurnContext>> for GuardianReviewContext {
fn from(turn: Arc<TurnContext>) -> Self {
Self {
parent_response_id: turn
.extension_data
.get::<codex_api::ResponseId>()
.map(|id| id.0.clone()),
parent_response: turn.extension_data.get::<ParentResponse>(),
environments: turn.initial_environments.clone(),
project_instructions: None,
model_info: Arc::clone(turn.model_info()),
codex-rs/core/src/guardian/review_session.rs
+22 -1
View File
@@ -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::<super::PendingReviewRouting>();
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::<super::PendingReviewRouting>();
// Finalizing input consumes this context; failures before admission leave it pending.
let transcript_admitted = review_session
.session
codex-rs/core/src/guardian/review_session_tests.rs
+43
View File
@@ -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(&params.node_repl_policy);
let reviewer_session = Arc::clone(&reviewer.session);
let manager = prewarm_test_session(&params, 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::<super::super::PendingReviewRouting>()
.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::<super::super::PendingReviewRouting>()
.is_none()
);
assert!(!submitted_text.contains("__oailb"));
assert!(matches!(
outcome,
GuardianReviewSessionOutcome::Completed(Ok(_))
codex-rs/core/src/guardian/routing.rs
+15
View File
@@ -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<ChatGptRoutingAffinity>,
}
/// Installed for the exclusive lifetime of one review, including its inference retries.
pub(crate) struct PendingReviewRouting(pub(crate) Option<Arc<ParentResponse>>);
codex-rs/core/src/model_request_tests.rs
+1
View File
@@ -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,
codex-rs/core/src/responses_metadata.rs
+3
View File
@@ -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<String>,
/// Transport-only parent pin; never projected into headers or body metadata.
pub(crate) parent_routing_affinity: Option<codex_http_client::ChatGptRoutingAffinity>,
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,
codex-rs/core/src/session/session.rs
+11
View File
@@ -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::<crate::guardian::PendingReviewRouting>()
&& 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
codex-rs/core/src/session/turn.rs
+6
View File
@@ -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));
codex-rs/core/src/turn_metadata_tests.rs
+66
View File
@@ -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");
codex-rs/http-client/src/chatgpt_cloudflare_cookies.rs
+1 -1
View File
@@ -126,7 +126,7 @@ impl HttpClientFactory {
}
}
fn chatgpt_cookie_url(uri: &Uri) -> Option<reqwest::Url> {
pub(crate) fn chatgpt_cookie_url(uri: &Uri) -> Option<reqwest::Url> {
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" {
codex-rs/http-client/src/chatgpt_routing_affinity.rs
+115
View File
@@ -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<Jar>,
}
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<Self> {
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<HeaderValue> {
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<HeaderValue> {
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::<Vec<_>>();
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;
codex-rs/http-client/src/chatgpt_routing_affinity_tests.rs
+100
View File
@@ -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::<Vec<_>>();
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::<Vec<_>>();
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));
}
codex-rs/http-client/src/lib.rs
+2
View File
@@ -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;
codex-rs/websocket-client/src/cookie_tests.rs
+46 -9
View File
@@ -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<codex_http_client::ChatGptRoutingAffinity> = 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::<Vec<_>>();
parts.sort_unstable();
HeaderValue::from_str(&parts.join("; ")).unwrap()
}))
.collect::<Vec<_>>(),
[
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(),
codex-rs/websocket-client/src/lib.rs
+19 -3
View File
@@ -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<codex_http_client::ChatGptRoutingAffinity>,
inner: Option<ConnectionInner>,
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<'_>,