Share ChatGPT cookies between HTTP and WebSocket transports (#46506)

## Why

WebSocket handshakes did not reuse the HTTP cookie store or retain response cookies, so routing cookies such as `__oailb` were unavailable to subsequent connections.

## What changed

- Reuse the HTTP factory's ChatGPT cookies for secure WebSocket handshakes, preserving explicit `Cookie` headers and marking generated headers sensitive.
- Retain allowlisted infrastructure cookies from both successful and rejected upgrades. Keep configured cookies scoped to their factory and exclude account and session cookies from the shared store.
- Apply HTTPS cookie scope to `wss` requests, preserving host and path restrictions and excluding insecure `ws` requests and non-ChatGPT hosts.

## Testing

Add HTTP/WebSocket cookie-sharing coverage and local TLS handshake tests for routing-cookie reuse across connectors, rejected-upgrade refreshes, explicit header precedence, cookie scope, session-cookie exclusion, and deletion.

GitOrigin-RevId: 6bd6afe16e97cf9758ca7ba207a4e88c969497a4

codex-rs/http-client/src/chatgpt_cloudflare_cookies.rs

@@ -1,10 +1,17 @@
+//! ChatGPT cookies shared by HTTP and WebSocket transports. Only infrastructure cookies may be
+//! stored globally; configured cookies remain scoped to their factory.
+
 use std::sync::Arc;
 use std::sync::LazyLock;
 
+use http::HeaderMap;
+use http::HeaderValue;
+use http::Uri;
+use http::header::SET_COOKIE;
 use reqwest::cookie::CookieStore;
 use reqwest::cookie::Jar;
-use reqwest::header::HeaderValue;
 
+use crate::HttpClientFactory;
 use crate::chatgpt_hosts::is_allowed_chatgpt_host;
 
 // WARNING: this HTTP cookie store is process-global and may be shared across auth contexts.
@@ -96,6 +103,38 @@ impl CookieStore for ChatGptCookieStore {
     }
 }
 
+impl HttpClientFactory {
+    /// Returns cookies for a ChatGPT HTTPS or WSS request, using the same store as HTTP clients.
+    /// Explicit request Cookie headers should take precedence over this value.
+    pub fn chatgpt_cookie_header(&self, uri: &Uri) -> Option<HeaderValue> {
+        let url = chatgpt_cookie_url(uri)?;
+        let mut cookies = match self.chatgpt_cookie_store() {
+            Some(store) => store.cookies(&url),
+            None => SHARED_CHATGPT_CLOUDFLARE_COOKIE_STORE.cookies(&url),
+        }?;
+        cookies.set_sensitive(true);
+        Some(cookies)
+    }
+
+    /// Retains only allowlisted infrastructure cookies from a ChatGPT HTTPS or WSS response.
+    /// Account and session cookies are never added to the shared store.
+    pub fn store_chatgpt_response_cookies(&self, uri: &Uri, headers: &HeaderMap) {
+        if let Some(url) = chatgpt_cookie_url(uri) {
+            SHARED_CHATGPT_CLOUDFLARE_COOKIE_STORE
+                .set_cookies(&mut headers.get_all(SET_COOKIE).iter(), &url);
+        }
+    }
+}
+
+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" {
+        url.set_scheme("https").ok()?;
+    }
+    is_chatgpt_cookie_url(&url).then_some(url)
+}
+
 /// Adds the process-local ChatGPT infrastructure cookie jar used by Codex HTTP clients.
 ///
 /// WARNING: this jar is global within the process. It is only acceptable because it hardcodes a
@@ -177,9 +216,52 @@ fn is_allowed_cloudflare_cookie_name(name: &str) -> bool {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::OutboundProxyPolicy;
     use pretty_assertions::assert_eq;
     use reqwest::cookie::CookieStore;
 
+    #[test]
+    fn http_and_websocket_cookies_share_the_factory_store() {
+        let factory = HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault)
+            .with_chatgpt_cookies([HeaderValue::from_static("configured=owner")]);
+        let store = factory.chatgpt_cookie_store().unwrap();
+        let https = reqwest::Url::parse("https://http-websocket-cookies.chatgpt.com/api").unwrap();
+        let wss: Uri = "wss://http-websocket-cookies.chatgpt.com/api"
+            .parse()
+            .unwrap();
+        let cookie = HeaderValue::from_static("__oailb=from-http; Path=/; Secure");
+        store.set_cookies(&mut std::iter::once(&cookie), &https);
+
+        let header = factory.chatgpt_cookie_header(&wss).unwrap();
+        assert_eq!(header, "__oailb=from-http; configured=owner");
+        assert!(header.is_sensitive());
+
+        let mut response_headers = HeaderMap::new();
+        response_headers.insert(
+            SET_COOKIE,
+            HeaderValue::from_static("__oailb=from-wss; Path=/; Secure"),
+        );
+        factory.store_chatgpt_response_cookies(&wss, &response_headers);
+        assert_eq!(
+            store.cookies(&https),
+            Some(HeaderValue::from_static(
+                "__oailb=from-wss; configured=owner"
+            ))
+        );
+        assert_eq!(
+            factory.chatgpt_cookie_header(
+                &"ws://http-websocket-cookies.chatgpt.com/api"
+                    .parse()
+                    .unwrap()
+            ),
+            None
+        );
+        assert_eq!(
+            factory.chatgpt_cookie_header(&"wss://api.openai.com/api".parse().unwrap()),
+            None
+        );
+    }
+
     #[test]
     fn additional_cookies_use_current_path_scoped_cloudflare_cookies() {
         let cloudflare = Arc::new(ChatGptCloudflareCookieStore::default());

codex-rs/websocket-client/src/cookie_tests.rs

@@ -0,0 +1,221 @@
+//! Real WebSocket handshakes exercising shared routing cookies and their transport boundaries.
+
+use std::sync::Arc;
+
+use codex_http_client::HttpClientFactory;
+use codex_http_client::OutboundProxyPolicy;
+use codex_http_client::OutboundProxyRoute;
+use codex_utils_rustls_provider::ensure_rustls_crypto_provider;
+use pretty_assertions::assert_eq;
+use rcgen::CertifiedKey;
+use rcgen::generate_simple_self_signed;
+use rustls::ClientConfig;
+use rustls::RootCertStore;
+use rustls::ServerConfig;
+use rustls::pki_types::PrivateKeyDer;
+use rustls::pki_types::PrivatePkcs8KeyDer;
+use tokio::io::AsyncReadExt;
+use tokio::io::AsyncWriteExt;
+use tokio::net::TcpListener;
+use tokio_rustls::TlsAcceptor;
+use tokio_tungstenite::accept_hdr_async;
+use tokio_tungstenite::tungstenite::Error as WebSocketError;
+use tokio_tungstenite::tungstenite::client::IntoClientRequest;
+use tokio_tungstenite::tungstenite::handshake::server::Request;
+use tokio_tungstenite::tungstenite::handshake::server::Response;
+use tokio_tungstenite::tungstenite::http::HeaderValue;
+use tokio_tungstenite::tungstenite::http::StatusCode;
+use tokio_tungstenite::tungstenite::http::header::COOKIE;
+use tokio_tungstenite::tungstenite::http::header::SET_COOKIE;
+use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
+
+use crate::AsyncIo;
+use crate::TcpNodelay;
+use crate::WebSocketConnector;
+
+#[tokio::test]
+async fn websocket_handshakes_share_routing_cookies_and_respect_cookie_scope() {
+    ensure_rustls_crypto_provider();
+    let host = "websocket-cookie-test.chatgpt.com";
+    let other_host = "other-websocket-cookie-test.chatgpt.com";
+    let CertifiedKey { cert, signing_key } = generate_simple_self_signed(vec![
+        host.to_string(),
+        other_host.to_string(),
+        "api.openai.com".to_string(),
+    ])
+    .unwrap();
+    let server_config = ServerConfig::builder()
+        .with_no_client_auth()
+        .with_single_cert(
+            vec![cert.der().clone()],
+            PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(signing_key.serialize_der())),
+        )
+        .unwrap();
+    let acceptor = TlsAcceptor::from(Arc::new(server_config));
+    let mut roots = RootCertStore::empty();
+    roots.add(cert.der().clone()).unwrap();
+    let tls_config = Arc::new(
+        ClientConfig::builder()
+            .with_root_certificates(roots)
+            .with_no_client_auth(),
+    );
+
+    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),
+        (review_url.clone(), None),
+        (review_url.clone(), Some("explicit=keep")),
+        (format!("wss://{host}/outside"), None),
+        (
+            format!("wss://{other_host}/backend-api/codex/responses"),
+            None,
+        ),
+        ("wss://api.openai.com/v1/responses".to_string(), None),
+        ("wss://api.openai.com/v1/responses".to_string(), None),
+        (format!("ws://{host}/backend-api/codex/responses"), None),
+        (review_url.clone(), None),
+        (review_url.clone(), None),
+        (review_url, None),
+    ];
+    let responses: Vec<(StatusCode, Vec<&str>)> = vec![
+        (
+            StatusCode::SWITCHING_PROTOCOLS,
+            vec![
+                "__oailb=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"],
+        ),
+        (StatusCode::SWITCHING_PROTOCOLS, vec![]),
+        (StatusCode::SWITCHING_PROTOCOLS, vec![]),
+        (StatusCode::SWITCHING_PROTOCOLS, vec![]),
+        (
+            StatusCode::SWITCHING_PROTOCOLS,
+            vec!["__oailb=untrusted; Path=/; Secure"],
+        ),
+        (StatusCode::SWITCHING_PROTOCOLS, vec![]),
+        (
+            StatusCode::SWITCHING_PROTOCOLS,
+            vec!["__oailb=insecure; Path=/backend-api"],
+        ),
+        (StatusCode::SWITCHING_PROTOCOLS, vec![]),
+        (
+            StatusCode::SWITCHING_PROTOCOLS,
+            vec!["__oailb=; Path=/backend-api; Max-Age=0; Secure"],
+        ),
+        (StatusCode::SWITCHING_PROTOCOLS, vec![]),
+    ];
+    let expected_statuses = responses
+        .iter()
+        .map(|(status, _)| *status)
+        .collect::<Vec<_>>();
+    let secure = requests
+        .iter()
+        .map(|(url, _)| url.starts_with("wss:"))
+        .collect::<Vec<_>>();
+
+    // Terminate CONNECT locally so the real TLS/upgrade path can use ChatGPT hostnames without DNS
+    // overrides, environment changes, or external traffic.
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let proxy_url = format!("http://{}", listener.local_addr().unwrap());
+    let server = tokio::spawn(async move {
+        let mut received_cookies = Vec::new();
+        for ((status, cookies), secure) in responses.into_iter().zip(secure) {
+            let (mut stream, _) = listener.accept().await.unwrap();
+            let mut connect = Vec::new();
+            while !connect.ends_with(b"\r\n\r\n") {
+                connect.push(stream.read_u8().await.unwrap());
+            }
+            stream
+                .write_all(b"HTTP/1.1 200 Connection Established\r\n\r\n")
+                .await
+                .unwrap();
+            let stream: Box<dyn AsyncIo> = if secure {
+                Box::new(acceptor.accept(stream).await.unwrap())
+            } else {
+                Box::new(stream)
+            };
+            let result = accept_hdr_async(stream, |request: &Request, mut response: Response| {
+                received_cookies.push(request.headers().get(COOKIE).cloned());
+                for cookie in cookies {
+                    response
+                        .headers_mut()
+                        .append(SET_COOKIE, HeaderValue::from_static(cookie));
+                }
+                if status == StatusCode::SWITCHING_PROTOCOLS {
+                    Ok(response)
+                } else {
+                    *response.status_mut() = status;
+                    Err(response.map(|()| None))
+                }
+            })
+            .await;
+            if status == StatusCode::SWITCHING_PROTOCOLS {
+                drop(result.unwrap());
+            } else {
+                assert!(matches!(result, Err(WebSocketError::Http(_))));
+            }
+        }
+        received_cookies
+    });
+
+    let mut actual_statuses = Vec::new();
+    for (url, explicit_cookie) in requests {
+        // Separate connector/factory instances must still see the shared routing cookie.
+        let connector = WebSocketConnector {
+            http_client_factory: HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
+            tls_config: Some(Arc::clone(&tls_config)),
+            tcp_nodelay: TcpNodelay::Default,
+        };
+        let mut request = url.into_client_request().unwrap();
+        if let Some(cookie) = explicit_cookie {
+            request
+                .headers_mut()
+                .insert(COOKIE, HeaderValue::from_static(cookie));
+        }
+        let result = connector
+            .connect_with_route(
+                request,
+                WebSocketConfig::default(),
+                OutboundProxyRoute::Proxy {
+                    url: proxy_url.clone(),
+                    no_proxy: None,
+                },
+                /*loopback_direct*/ false,
+            )
+            .await;
+        let status = match result {
+            Ok((connection, response)) => {
+                drop(connection);
+                response.status()
+            }
+            Err(WebSocketError::Http(response)) => response.status(),
+            Err(error) => panic!("unexpected handshake error: {error}"),
+        };
+        actual_statuses.push(status);
+    }
+    assert_eq!(actual_statuses, expected_statuses);
+    assert_eq!(
+        server.await.unwrap(),
+        [
+            None,
+            Some("__oailb=west"),
+            Some("explicit=keep"),
+            None,
+            None,
+            None,
+            None,
+            None,
+            Some("__oailb=east"),
+            Some("__oailb=east"),
+            None
+        ]
+        .map(|cookie| cookie.map(HeaderValue::from_static))
+        .to_vec(),
+    );
+}

codex-rs/websocket-client/src/lib.rs

@@ -1,4 +1,5 @@
-//! Proxy-aware WebSocket connection setup shared by Codex API clients.
+//! Proxy-aware WebSocket connection setup shared by Codex API clients, reusing the HTTP factory's
+//! ChatGPT cookie store for secure handshakes.
 
 mod dialer;
 
@@ -27,6 +28,7 @@ use tokio_tungstenite::tungstenite::Message;
 use tokio_tungstenite::tungstenite::handshake::client::Request;
 use tokio_tungstenite::tungstenite::handshake::client::Response;
 use tokio_tungstenite::tungstenite::http::Uri;
+use tokio_tungstenite::tungstenite::http::header::COOKIE;
 use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
 
 /// Connects WebSockets using the outbound proxy policy resolved by application configuration.
@@ -133,12 +135,18 @@ impl WebSocketConnector {
 
     async fn connect_with_route(
         &self,
-        request: Request,
+        mut request: Request,
         config: WebSocketConfig,
         proxy_route: OutboundProxyRoute,
         loopback_direct: bool,
     ) -> Result<(WebSocketConnection, Response), WebSocketError> {
-        let (inner, response) = dialer::connect(
+        let uri = request.uri().clone();
+        if !request.headers().contains_key(COOKIE)
+            && let Some(cookies) = self.http_client_factory.chatgpt_cookie_header(&uri)
+        {
+            request.headers_mut().insert(COOKIE, cookies);
+        }
+        let result = dialer::connect(
             request,
             config,
             self.tls_config.clone(),
@@ -147,7 +155,18 @@ impl WebSocketConnector {
             loopback_direct,
         )
         .boxed()
-        .await?;
+        .await;
+        // Like HTTP responses, rejected upgrades can also refresh infrastructure cookies.
+        match &result {
+            Ok((_, response)) => self
+                .http_client_factory
+                .store_chatgpt_response_cookies(&uri, response.headers()),
+            Err(WebSocketError::Http(response)) => self
+                .http_client_factory
+                .store_chatgpt_response_cookies(&uri, response.headers()),
+            Err(_) => {}
+        }
+        let (inner, response) = result?;
         Ok((WebSocketConnection { inner }, response))
     }
 }
@@ -239,3 +258,7 @@ impl<T> AsyncIo for T where T: AsyncRead + AsyncWrite + Send + Unpin {}
 #[cfg(test)]
 #[path = "lib_tests.rs"]
 mod tests;
+
+#[cfg(test)]
+#[path = "cookie_tests.rs"]
+mod cookie_tests;