mirror of
https://github.com/openai/codex.git
synced 2026-10-29 19:08:43 +00:00
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
This commit is contained in:
@@ -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>,
|
||||
|
||||
@@ -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}")),
|
||||
|
||||
@@ -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?;
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -98,6 +98,7 @@ pub fn spawn_response_stream(
|
||||
});
|
||||
|
||||
ResponseStream {
|
||||
routing_affinity: None,
|
||||
rx_event,
|
||||
upstream_request_id,
|
||||
interrupt: None,
|
||||
|
||||
Reference in New Issue
Block a user