[codex] wire native HTTP state transport [ci changed_files]

codex-rs/codex-api/src/auth.rs

@@ -41,6 +41,27 @@ pub trait AuthProvider: Send + Sync {
         headers
     }
 
+    /// Adds auth headers for a concrete HTTP request URL.
+    ///
+    /// URL-sensitive providers can override this to scope request-specific
+    /// headers without exposing them through generic/non-HTTP auth helpers.
+    fn add_auth_headers_for_url(&self, _request_url: &str, headers: &mut HeaderMap) {
+        self.add_auth_headers(headers);
+    }
+
+    /// Observes response headers for auth state that may need to rotate after a request.
+    ///
+    /// Most providers do not need this. Providers with server-minted,
+    /// request-scoped state may validate the URL and selectively persist
+    /// response headers here.
+    fn observe_response_headers(
+        &self,
+        _request_url: &str,
+        _request_headers: &HeaderMap,
+        _response_headers: &HeaderMap,
+    ) {
+    }
+
     /// Applies auth to a complete outbound request and returns the request to send.
     ///
     /// The input `request` is moved into this method. Implementations may mutate
@@ -54,7 +75,7 @@ pub trait AuthProvider: Send + Sync {
     /// If this returns [`AuthError`], the request should not be sent.
     async fn apply_auth(&self, request: Request) -> Result<Request, AuthError> {
         let mut request = request;
-        self.add_auth_headers(&mut request.headers);
+        self.add_auth_headers_for_url(&request.url, &mut request.headers);
         Ok(request)
     }
 }

codex-rs/codex-api/src/endpoint/mod.rs

@@ -6,6 +6,7 @@ pub(crate) mod realtime_call;
 pub(crate) mod realtime_websocket;
 pub(crate) mod responses;
 pub(crate) mod responses_websocket;
+mod responses_websocket_auth;
 pub(crate) mod search;
 mod session;
 

codex-rs/codex-api/src/endpoint/realtime_websocket/methods.rs

@@ -1,3 +1,4 @@
+use crate::auth::SharedAuthProvider;
 use crate::endpoint::realtime_websocket::methods_common::conversation_function_call_output_message;
 use crate::endpoint::realtime_websocket::methods_common::conversation_item_create_message;
 use crate::endpoint::realtime_websocket::methods_common::normalized_session_mode;
@@ -550,11 +551,21 @@ fn contains_transcript_entry(entries: &[RealtimeTranscriptEntry], role: &str, te
 
 pub struct RealtimeWebsocketClient {
     provider: Provider,
+    auth: Option<SharedAuthProvider>,
 }
 
 impl RealtimeWebsocketClient {
     pub fn new(provider: Provider) -> Self {
-        Self { provider }
+        Self {
+            provider,
+            auth: None,
+        }
+    }
+
+    /// Applies URL-aware auth during handshakes and observes handshake responses.
+    pub fn with_auth(mut self, auth: SharedAuthProvider) -> Self {
+        self.auth = Some(auth);
+        self
     }
 
     pub async fn connect(
@@ -653,6 +664,10 @@ impl RealtimeWebsocketClient {
             default_headers,
         );
         request.headers_mut().extend(headers);
+        if let Some(auth) = self.auth.as_ref() {
+            auth.add_auth_headers_for_url(ws_url.as_str(), request.headers_mut());
+        }
+        let request_headers = request.headers().clone();
 
         info!("connecting realtime websocket: {ws_url}");
         // Realtime websocket TLS should honor the same custom-CA env vars as the rest of Codex's
@@ -660,14 +675,39 @@ impl RealtimeWebsocketClient {
         let connector = maybe_build_rustls_client_config_with_custom_ca()
             .map_err(|err| ApiError::Stream(format!("failed to configure websocket TLS: {err}")))?
             .map(tokio_tungstenite::Connector::Rustls);
-        let (stream, response) = tokio_tungstenite::connect_async_tls_with_config(
+        let connect_result = tokio_tungstenite::connect_async_tls_with_config(
             request,
             Some(websocket_config()),
             false,
             connector,
         )
-        .await
-        .map_err(|err| ApiError::Stream(format!("failed to connect realtime websocket: {err}")))?;
+        .await;
+        let (stream, response) = match connect_result {
+            Ok((stream, response)) => {
+                if let Some(auth) = self.auth.as_ref() {
+                    auth.observe_response_headers(
+                        ws_url.as_str(),
+                        &request_headers,
+                        response.headers(),
+                    );
+                }
+                (stream, response)
+            }
+            Err(err) => {
+                if let Some(auth) = self.auth.as_ref()
+                    && let WsError::Http(response) = &err
+                {
+                    auth.observe_response_headers(
+                        ws_url.as_str(),
+                        &request_headers,
+                        response.headers(),
+                    );
+                }
+                return Err(ApiError::Stream(format!(
+                    "failed to connect realtime websocket: {err}"
+                )));
+            }
+        };
         info!(
             ws_url = %ws_url,
             status = %response.status(),

codex-rs/codex-api/src/endpoint/responses_websocket.rs

@@ -1,3 +1,4 @@
+use super::responses_websocket_auth::WebsocketAuthContext;
 use crate::auth::SharedAuthProvider;
 use crate::common::ResponseEvent;
 use crate::common::ResponseProcessedWsRequest;
@@ -167,6 +168,7 @@ pub struct ResponsesWebsocketConnection {
     server_reasoning_included: bool,
     models_etag: Option<String>,
     server_model: Option<String>,
+    auth_context: WebsocketAuthContext,
     telemetry: Option<Arc<dyn WebsocketTelemetry>>,
 }
 
@@ -178,6 +180,7 @@ impl std::fmt::Debug for ResponsesWebsocketConnection {
             .field("server_reasoning_included", &self.server_reasoning_included)
             .field("models_etag", &self.models_etag)
             .field("server_model", &self.server_model)
+            .field("auth_context", &"<auth-context>")
             .field("telemetry", &self.telemetry.as_ref().map(|_| "<telemetry>"))
             .finish()
     }
@@ -190,6 +193,7 @@ impl ResponsesWebsocketConnection {
         server_reasoning_included: bool,
         models_etag: Option<String>,
         server_model: Option<String>,
+        auth_context: WebsocketAuthContext,
         telemetry: Option<Arc<dyn WebsocketTelemetry>>,
     ) -> Self {
         Self {
@@ -198,6 +202,7 @@ impl ResponsesWebsocketConnection {
             server_reasoning_included,
             models_etag,
             server_model,
+            auth_context,
             telemetry,
         }
     }
@@ -258,6 +263,7 @@ impl ResponsesWebsocketConnection {
         let server_reasoning_included = self.server_reasoning_included;
         let models_etag = self.models_etag.clone();
         let server_model = self.server_model.clone();
+        let auth_context = self.auth_context.clone();
         let telemetry = self.telemetry.clone();
         let request_body = serde_json::to_value(&request).map_err(|err| {
             ApiError::Stream(format!("failed to encode websocket request: {err}"))
@@ -297,6 +303,7 @@ impl ResponsesWebsocketConnection {
                         tx_event.clone(),
                         request_body,
                         idle_timeout,
+                        &auth_context,
                         telemetry,
                         connection_reused,
                     )
@@ -377,19 +384,35 @@ impl ResponsesWebsocketClient {
             .provider
             .websocket_url_for_path("responses")
             .map_err(|err| ApiError::Stream(format!("failed to build websocket URL: {err}")))?;
+        let request_url = ws_url.to_string();
 
         let mut headers =
             merge_request_headers(&self.provider.headers, extra_headers, default_headers);
-        self.auth.add_auth_headers(&mut headers);
+        self.auth
+            .add_auth_headers_for_url(&request_url, &mut headers);
+        let auth_context =
+            WebsocketAuthContext::new(self.auth.clone(), request_url, headers.clone());
 
-        let (stream, _status, server_reasoning_included, models_etag, server_model) =
-            connect_websocket(ws_url, headers, turn_state.clone()).await?;
+        let (
+            stream,
+            _status,
+            server_reasoning_included,
+            models_etag,
+            server_model,
+            response_headers,
+        ) = connect_websocket(ws_url, headers, turn_state.clone())
+            .await
+            .inspect_err(|error| {
+                auth_context.observe_error_headers(error);
+            })?;
+        auth_context.observe_response_headers(&response_headers);
         Ok(ResponsesWebsocketConnection::new(
             stream,
             self.provider.stream_idle_timeout,
             server_reasoning_included,
             models_etag,
             server_model,
+            auth_context,
             telemetry,
         ))
     }
@@ -411,13 +434,22 @@ impl ResponsesWebsocketClient {
             .provider
             .websocket_url_for_path("responses")
             .map_err(|err| ApiError::Stream(format!("failed to build websocket URL: {err}")))?;
+        let request_url = ws_url.to_string();
 
         let mut headers =
             merge_request_headers(&self.provider.headers, extra_headers, default_headers);
-        self.auth.add_auth_headers(&mut headers);
+        self.auth
+            .add_auth_headers_for_url(&request_url, &mut headers);
+        let auth_context =
+            WebsocketAuthContext::new(self.auth.clone(), request_url, headers.clone());
 
-        let (mut stream, status, reasoning_included, models_etag, server_model) =
-            connect_websocket(ws_url.clone(), headers, /*turn_state*/ None).await?;
+        let (mut stream, status, reasoning_included, models_etag, server_model, response_headers) =
+            connect_websocket(ws_url.clone(), headers, /*turn_state*/ None)
+                .await
+                .inspect_err(|error| {
+                    auth_context.observe_error_headers(error);
+                })?;
+        auth_context.observe_response_headers(&response_headers);
         let immediate_close = tokio::time::timeout(immediate_close_timeout, stream.next())
             .await
             .ok()
@@ -472,7 +504,17 @@ async fn connect_websocket(
     url: Url,
     headers: HeaderMap,
     turn_state: Option<Arc<OnceLock<String>>>,
-) -> Result<(WsStream, StatusCode, bool, Option<String>, Option<String>), ApiError> {
+) -> Result<
+    (
+        WsStream,
+        StatusCode,
+        bool,
+        Option<String>,
+        Option<String>,
+        HeaderMap,
+    ),
+    ApiError,
+> {
     ensure_rustls_crypto_provider();
     info!("connecting to websocket: {url}");
 
@@ -499,10 +541,7 @@ async fn connect_websocket(
 
     let (stream, response) = match response {
         Ok((stream, response)) => {
-            info!(
-                "successfully connected to websocket: {url}, headers: {:?}",
-                response.headers()
-            );
+            info!("successfully connected to websocket: {url}");
             (stream, response)
         }
         Err(err) => {
@@ -536,6 +575,7 @@ async fn connect_websocket(
         reasoning_included,
         models_etag,
         server_model,
+        response.headers().clone(),
     ))
 }
 
@@ -630,12 +670,12 @@ fn map_wrapped_websocket_error_event(
     Some(ApiError::Transport(TransportError::Http {
         status,
         url: None,
-        headers: headers.map(json_headers_to_http_headers),
+        headers: headers.as_ref().map(json_headers_to_http_headers),
         body: Some(original_payload),
     }))
 }
 
-fn json_headers_to_http_headers(headers: JsonMap<String, Value>) -> HeaderMap {
+fn json_headers_to_http_headers(headers: &JsonMap<String, Value>) -> HeaderMap {
     let mut mapped = HeaderMap::new();
     for (name, value) in headers {
         let Ok(header_name) = HeaderName::from_bytes(name.as_bytes()) else {
@@ -649,9 +689,9 @@ fn json_headers_to_http_headers(headers: JsonMap<String, Value>) -> HeaderMap {
     mapped
 }
 
-fn json_header_value(value: Value) -> Option<HeaderValue> {
+fn json_header_value(value: &Value) -> Option<HeaderValue> {
     let value = match value {
-        Value::String(value) => value,
+        Value::String(value) => value.clone(),
         Value::Number(value) => value.to_string(),
         Value::Bool(value) => value.to_string(),
         _ => return None,
@@ -664,6 +704,7 @@ async fn run_websocket_response_stream(
     tx_event: mpsc::Sender<std::result::Result<ResponseEvent, ApiError>>,
     request_body: Value,
     idle_timeout: Duration,
+    auth_context: &WebsocketAuthContext,
     telemetry: Option<Arc<dyn WebsocketTelemetry>>,
     connection_reused: bool,
 ) -> Result<(), ApiError> {
@@ -702,12 +743,25 @@ async fn run_websocket_response_stream(
 
         match message {
             Message::Text(text) => {
-                trace!("websocket event: {text}");
-                if let Some(wrapped_error) = parse_wrapped_websocket_error_event(&text)
-                    && let Some(error) =
+                let wrapped_error = parse_wrapped_websocket_error_event(&text);
+                if wrapped_error.is_some() {
+                    trace!("websocket error event");
+                } else {
+                    trace!("websocket event: {text}");
+                }
+                if let Some(wrapped_error) = wrapped_error {
+                    if let Some(response_headers) = wrapped_error.headers.as_ref() {
+                        auth_context.observe_response_headers(&json_headers_to_http_headers(
+                            response_headers,
+                        ));
+                    }
+                    if let Some(error) =
                         map_wrapped_websocket_error_event(wrapped_error, text.to_string())
-                {
-                    return Err(error);
+                    {
+                        return Err(error);
+                    }
+                    debug!("ignoring unmapped websocket error event");
+                    continue;
                 }
 
                 let event = match serde_json::from_str::<ResponsesStreamEvent>(&text) {

codex-rs/codex-api/src/endpoint/responses_websocket_auth.rs

@@ -0,0 +1,52 @@
+use crate::auth::SharedAuthProvider;
+use crate::error::ApiError;
+use codex_client::TransportError;
+use http::HeaderMap;
+use std::sync::Arc;
+use std::sync::Mutex;
+
+#[derive(Clone)]
+pub(super) struct WebsocketAuthContext {
+    auth: SharedAuthProvider,
+    request_url: String,
+    request_headers: Arc<Mutex<HeaderMap>>,
+}
+
+impl WebsocketAuthContext {
+    pub(super) fn new(
+        auth: SharedAuthProvider,
+        request_url: String,
+        request_headers: HeaderMap,
+    ) -> Self {
+        Self {
+            auth,
+            request_url,
+            request_headers: Arc::new(Mutex::new(request_headers)),
+        }
+    }
+
+    pub(super) fn observe_response_headers(&self, response_headers: &HeaderMap) {
+        let mut request_headers = self
+            .request_headers
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner);
+        self.auth
+            .observe_response_headers(&self.request_url, &request_headers, response_headers);
+        self.auth
+            .add_auth_headers_for_url(&self.request_url, &mut request_headers);
+    }
+
+    pub(super) fn observe_error_headers(&self, error: &ApiError) {
+        if let ApiError::Transport(TransportError::Http {
+            headers: Some(response_headers),
+            ..
+        }) = error
+        {
+            self.observe_response_headers(response_headers);
+        }
+    }
+}
+
+#[cfg(test)]
+#[path = "responses_websocket_auth_tests.rs"]
+mod tests;

codex-rs/codex-api/src/endpoint/responses_websocket_auth_tests.rs

@@ -0,0 +1,124 @@
+use super::*;
+use crate::auth::AuthProvider;
+use http::HeaderValue;
+use http::StatusCode;
+use pretty_assertions::assert_eq;
+use std::sync::Mutex as StdMutex;
+
+const REQUEST_URL: &str = "wss://chatgpt.com/backend-api/codex/responses";
+const SENT_STATE: &str = "sent-state";
+const HANDSHAKE_STATE: &str = "handshake-state";
+const WRAPPED_ERROR_STATE: &str = "wrapped-error-state";
+
+struct RecordingAuthProvider {
+    state: StdMutex<String>,
+    observed_updates: StdMutex<Vec<(String, String, String)>>,
+}
+
+impl Default for RecordingAuthProvider {
+    fn default() -> Self {
+        Self {
+            state: StdMutex::new(SENT_STATE.to_string()),
+            observed_updates: StdMutex::new(Vec::new()),
+        }
+    }
+}
+
+impl AuthProvider for RecordingAuthProvider {
+    fn add_auth_headers(&self, _headers: &mut HeaderMap) {}
+
+    fn add_auth_headers_for_url(&self, _request_url: &str, headers: &mut HeaderMap) {
+        let state = self
+            .state
+            .lock()
+            .expect("state lock should not be poisoned");
+        headers.insert(
+            "x-test-state",
+            HeaderValue::from_str(&state).expect("state should be a valid header value"),
+        );
+    }
+
+    fn observe_response_headers(
+        &self,
+        request_url: &str,
+        request_headers: &HeaderMap,
+        response_headers: &HeaderMap,
+    ) {
+        let Some(sent_state) = request_headers
+            .get("x-test-state")
+            .and_then(|value| value.to_str().ok())
+        else {
+            return;
+        };
+        let Some(update_state) = response_headers
+            .get("x-test-state-update")
+            .and_then(|value| value.to_str().ok())
+        else {
+            return;
+        };
+        self.observed_updates
+            .lock()
+            .expect("observed updates lock should not be poisoned")
+            .push((
+                request_url.to_string(),
+                sent_state.to_string(),
+                update_state.to_string(),
+            ));
+        *self
+            .state
+            .lock()
+            .expect("state lock should not be poisoned") = update_state.to_string();
+    }
+}
+
+#[test]
+fn observes_handshake_and_wrapped_error_headers() {
+    let auth = Arc::new(RecordingAuthProvider::default());
+    let mut request_headers = HeaderMap::new();
+    auth.add_auth_headers_for_url(REQUEST_URL, &mut request_headers);
+    assert_eq!(
+        request_headers
+            .get("x-test-state")
+            .and_then(|value| value.to_str().ok()),
+        Some(SENT_STATE)
+    );
+    let context = WebsocketAuthContext::new(auth.clone(), REQUEST_URL.to_string(), request_headers);
+
+    let mut handshake_headers = HeaderMap::new();
+    handshake_headers.insert(
+        "x-test-state-update",
+        HeaderValue::from_static(HANDSHAKE_STATE),
+    );
+    context.observe_response_headers(&handshake_headers);
+
+    let mut wrapped_error_headers = HeaderMap::new();
+    wrapped_error_headers.insert(
+        "x-test-state-update",
+        HeaderValue::from_static(WRAPPED_ERROR_STATE),
+    );
+    context.observe_error_headers(&ApiError::Transport(TransportError::Http {
+        status: StatusCode::BAD_REQUEST,
+        url: Some(REQUEST_URL.to_string()),
+        headers: Some(wrapped_error_headers),
+        body: None,
+    }));
+
+    assert_eq!(
+        *auth
+            .observed_updates
+            .lock()
+            .expect("observed updates lock should not be poisoned"),
+        vec![
+            (
+                REQUEST_URL.to_string(),
+                SENT_STATE.to_string(),
+                HANDSHAKE_STATE.to_string(),
+            ),
+            (
+                REQUEST_URL.to_string(),
+                HANDSHAKE_STATE.to_string(),
+                WRAPPED_ERROR_STATE.to_string(),
+            ),
+        ]
+    );
+}

codex-rs/codex-api/src/endpoint/session.rs

@@ -102,7 +102,16 @@ impl<T: HttpTransport> EndpointSession<T> {
                 let transport = &self.transport;
                 async move {
                     let req = auth.apply_auth(req).await.map_err(TransportError::from)?;
-                    transport.execute(req).await
+                    let request_url = req.url.clone();
+                    let request_headers = req.headers.clone();
+                    let response = transport.execute(req).await;
+                    observe_auth_response_headers(
+                        auth.as_ref(),
+                        &request_url,
+                        &request_headers,
+                        &response,
+                    );
+                    response
                 }
             },
         )
@@ -143,7 +152,16 @@ impl<T: HttpTransport> EndpointSession<T> {
                 let transport = &self.transport;
                 async move {
                     let req = auth.apply_auth(req).await.map_err(TransportError::from)?;
-                    transport.stream(req).await
+                    let request_url = req.url.clone();
+                    let request_headers = req.headers.clone();
+                    let response = transport.stream(req).await;
+                    observe_auth_response_headers(
+                        auth.as_ref(),
+                        &request_url,
+                        &request_headers,
+                        &response,
+                    );
+                    response
                 }
             },
         )
@@ -152,3 +170,44 @@ impl<T: HttpTransport> EndpointSession<T> {
         Ok(stream)
     }
 }
+
+fn observe_auth_response_headers<T>(
+    auth: &dyn crate::auth::AuthProvider,
+    request_url: &str,
+    request_headers: &HeaderMap,
+    response: &Result<T, TransportError>,
+) where
+    T: ResponseHeaders,
+{
+    match response {
+        Ok(response) => {
+            auth.observe_response_headers(request_url, request_headers, response.headers())
+        }
+        Err(TransportError::Http {
+            headers: Some(headers),
+            ..
+        }) => auth.observe_response_headers(request_url, request_headers, headers),
+        Err(_) => {}
+    }
+}
+
+/// Exposes transport response headers so auth providers can observe request-scoped updates.
+trait ResponseHeaders {
+    fn headers(&self) -> &HeaderMap;
+}
+
+impl ResponseHeaders for Response {
+    fn headers(&self) -> &HeaderMap {
+        &self.headers
+    }
+}
+
+impl ResponseHeaders for StreamResponse {
+    fn headers(&self) -> &HeaderMap {
+        &self.headers
+    }
+}
+
+#[cfg(test)]
+#[path = "session_tests.rs"]
+mod tests;

codex-rs/codex-api/src/endpoint/session_tests.rs

@@ -0,0 +1,174 @@
+use super::*;
+use crate::provider::RetryConfig;
+use async_trait::async_trait;
+use codex_client::Request;
+use http::HeaderValue;
+use http::StatusCode;
+use pretty_assertions::assert_eq;
+use std::sync::Mutex as StdMutex;
+use std::time::Duration;
+
+const RESPONSES_URL: &str = "https://chatgpt.com/backend-api/codex/responses";
+
+#[derive(Default)]
+struct RecordingAuthProvider {
+    observed_headers: StdMutex<Vec<(String, HeaderMap, HeaderMap)>>,
+}
+
+impl crate::auth::AuthProvider for RecordingAuthProvider {
+    fn add_auth_headers(&self, _headers: &mut HeaderMap) {}
+
+    fn add_auth_headers_for_url(&self, _request_url: &str, headers: &mut HeaderMap) {
+        headers.insert("x-test-state", HeaderValue::from_static("sent-state"));
+    }
+
+    fn observe_response_headers(
+        &self,
+        request_url: &str,
+        request_headers: &HeaderMap,
+        response_headers: &HeaderMap,
+    ) {
+        self.observed_headers
+            .lock()
+            .expect("recording auth lock should not be poisoned")
+            .push((
+                request_url.to_string(),
+                request_headers.clone(),
+                response_headers.clone(),
+            ));
+    }
+}
+
+#[derive(Clone)]
+struct RejectingTransport {
+    requests: Arc<StdMutex<Vec<(String, HeaderMap)>>>,
+    response_headers: HeaderMap,
+}
+
+impl RejectingTransport {
+    fn new(response_headers: HeaderMap) -> Self {
+        Self {
+            requests: Arc::new(StdMutex::new(Vec::new())),
+            response_headers,
+        }
+    }
+
+    fn http_error(&self, request: Request) -> TransportError {
+        self.requests
+            .lock()
+            .expect("recording transport lock should not be poisoned")
+            .push((request.url.clone(), request.headers));
+        TransportError::Http {
+            status: StatusCode::UNAUTHORIZED,
+            url: Some(request.url),
+            headers: Some(self.response_headers.clone()),
+            body: None,
+        }
+    }
+}
+
+#[async_trait]
+impl HttpTransport for RejectingTransport {
+    async fn execute(&self, request: Request) -> Result<Response, TransportError> {
+        Err(self.http_error(request))
+    }
+
+    async fn stream(&self, request: Request) -> Result<StreamResponse, TransportError> {
+        Err(self.http_error(request))
+    }
+}
+
+#[tokio::test]
+async fn execute_attaches_url_scoped_auth_and_observes_http_error_headers() {
+    let (session, transport, auth, request_headers, response_headers) = test_session();
+
+    session
+        .execute(
+            Method::POST,
+            "responses",
+            HeaderMap::new(),
+            /*body*/ None,
+        )
+        .await
+        .expect_err("request should fail");
+
+    assert_request_and_response_headers(&transport, &auth, request_headers, response_headers);
+}
+
+#[tokio::test]
+async fn stream_attaches_url_scoped_auth_and_observes_http_error_headers() {
+    let (session, transport, auth, request_headers, response_headers) = test_session();
+
+    let result = session
+        .stream_with(
+            Method::POST,
+            "responses",
+            HeaderMap::new(),
+            /*body*/ None,
+            |_| {},
+        )
+        .await;
+    assert!(result.is_err(), "request should fail");
+
+    assert_request_and_response_headers(&transport, &auth, request_headers, response_headers);
+}
+
+fn test_session() -> (
+    EndpointSession<RejectingTransport>,
+    RejectingTransport,
+    Arc<RecordingAuthProvider>,
+    HeaderMap,
+    HeaderMap,
+) {
+    let mut request_headers = HeaderMap::new();
+    request_headers.insert("x-test-state", HeaderValue::from_static("sent-state"));
+    let mut response_headers = HeaderMap::new();
+    response_headers.insert(
+        "x-test-state-update",
+        HeaderValue::from_static("rotated-state"),
+    );
+    let transport = RejectingTransport::new(response_headers.clone());
+    let auth = Arc::new(RecordingAuthProvider::default());
+    let session = EndpointSession::new(transport.clone(), provider(), auth.clone());
+    (session, transport, auth, request_headers, response_headers)
+}
+
+fn provider() -> Provider {
+    Provider {
+        name: "test".to_string(),
+        base_url: "https://chatgpt.com/backend-api/codex".to_string(),
+        query_params: None,
+        headers: HeaderMap::new(),
+        retry: RetryConfig {
+            max_attempts: 1,
+            base_delay: Duration::from_millis(1),
+            retry_429: false,
+            retry_5xx: false,
+            retry_transport: false,
+        },
+        stream_idle_timeout: Duration::from_secs(1),
+    }
+}
+
+fn assert_request_and_response_headers(
+    transport: &RejectingTransport,
+    auth: &RecordingAuthProvider,
+    request_headers: HeaderMap,
+    response_headers: HeaderMap,
+) {
+    assert_eq!(
+        transport
+            .requests
+            .lock()
+            .expect("recording transport lock should not be poisoned")
+            .as_slice(),
+        &[(RESPONSES_URL.to_string(), request_headers.clone())]
+    );
+    assert_eq!(
+        auth.observed_headers
+            .lock()
+            .expect("recording auth lock should not be poisoned")
+            .as_slice(),
+        &[(RESPONSES_URL.to_string(), request_headers, response_headers,)]
+    );
+}

codex-rs/codex-api/src/files.rs

@@ -129,7 +129,9 @@ pub async fn upload_local_file(
         .unwrap_or("file")
         .to_string();
     let create_url = format!("{}/files", base_url.trim_end_matches('/'));
-    let create_response = authorized_request(auth, reqwest::Method::POST, &create_url)
+    let (create_request, create_request_headers) =
+        authorized_request(auth, reqwest::Method::POST, &create_url);
+    let create_response = create_request
         .json(&serde_json::json!({
             "file_name": file_name,
             "file_size": metadata.len(),
@@ -141,6 +143,11 @@ pub async fn upload_local_file(
             url: create_url.clone(),
             source,
         })?;
+    auth.observe_response_headers(
+        &create_url,
+        &create_request_headers,
+        create_response.headers(),
+    );
     let create_status = create_response.status();
     let create_body = create_response.text().await.unwrap_or_default();
     if !create_status.is_success() {
@@ -191,7 +198,9 @@ pub async fn upload_local_file(
     );
     let finalize_started_at = Instant::now();
     loop {
-        let finalize_response = authorized_request(auth, reqwest::Method::POST, &finalize_url)
+        let (finalize_request, finalize_request_headers) =
+            authorized_request(auth, reqwest::Method::POST, &finalize_url);
+        let finalize_response = finalize_request
             .json(&serde_json::json!({}))
             .send()
             .await
@@ -199,6 +208,11 @@ pub async fn upload_local_file(
                 url: finalize_url.clone(),
                 source,
             })?;
+        auth.observe_response_headers(
+            &finalize_url,
+            &finalize_request_headers,
+            finalize_response.headers(),
+        );
         let finalize_status = finalize_response.status();
         let finalize_body = finalize_response.text().await.unwrap_or_default();
         if !finalize_status.is_success() {
@@ -255,15 +269,18 @@ fn authorized_request(
     auth: &dyn AuthProvider,
     method: reqwest::Method,
     url: &str,
-) -> reqwest::RequestBuilder {
+) -> (reqwest::RequestBuilder, http::HeaderMap) {
     let mut headers = http::HeaderMap::new();
-    auth.add_auth_headers(&mut headers);
+    auth.add_auth_headers_for_url(url, &mut headers);
 
     let client = build_reqwest_client();
-    client
-        .request(method, url)
-        .timeout(OPENAI_FILE_REQUEST_TIMEOUT)
-        .headers(headers)
+    (
+        client
+            .request(method, url)
+            .timeout(OPENAI_FILE_REQUEST_TIMEOUT)
+            .headers(headers.clone()),
+        headers,
+    )
 }
 
 fn build_reqwest_client() -> reqwest::Client {
@@ -279,6 +296,7 @@ mod tests {
     use pretty_assertions::assert_eq;
     use reqwest::header::HeaderValue;
     use std::sync::Arc;
+    use std::sync::Mutex;
     use std::sync::atomic::AtomicUsize;
     use std::sync::atomic::Ordering;
     use tempfile::TempDir;
@@ -291,8 +309,16 @@ mod tests {
     use wiremock::matchers::method;
     use wiremock::matchers::path;
 
-    #[derive(Clone, Copy)]
-    struct ChatGptTestAuth;
+    const SENT_STATE: &str = "sent-state";
+    const CREATED_STATE: &str = "created-state";
+    const RETRY_STATE: &str = "retry-state";
+    const FINALIZED_STATE: &str = "finalized-state";
+    type ObservedUpdate = (String, Option<String>);
+
+    #[derive(Clone, Default)]
+    struct ChatGptTestAuth {
+        observed_updates: Arc<Mutex<Vec<ObservedUpdate>>>,
+    }
 
     impl AuthProvider for ChatGptTestAuth {
         fn add_auth_headers(&self, headers: &mut reqwest::header::HeaderMap) {
@@ -302,10 +328,33 @@ mod tests {
             );
             headers.insert("ChatGPT-Account-ID", HeaderValue::from_static("account_id"));
         }
+
+        fn add_auth_headers_for_url(&self, _request_url: &str, headers: &mut http::HeaderMap) {
+            self.add_auth_headers(headers);
+            headers.insert("x-test-state", HeaderValue::from_static(SENT_STATE));
+        }
+
+        fn observe_response_headers(
+            &self,
+            request_url: &str,
+            _request_headers: &http::HeaderMap,
+            response_headers: &http::HeaderMap,
+        ) {
+            self.observed_updates
+                .lock()
+                .expect("observed updates lock should not be poisoned")
+                .push((
+                    request_url.to_string(),
+                    response_headers
+                        .get("x-test-state-update")
+                        .and_then(|value| value.to_str().ok())
+                        .map(ToString::to_string),
+                ));
+        }
     }
 
     fn chatgpt_auth() -> ChatGptTestAuth {
-        ChatGptTestAuth
+        ChatGptTestAuth::default()
     }
 
     fn base_url_for(server: &MockServer) -> String {
@@ -318,6 +367,7 @@ mod tests {
         Mock::given(method("POST"))
             .and(path("/backend-api/files"))
             .and(header("chatgpt-account-id", "account_id"))
+            .and(header("x-test-state", SENT_STATE))
             .and(body_json(serde_json::json!({
                 "file_name": "hello.txt",
                 "file_size": 5,
@@ -325,6 +375,7 @@ mod tests {
             })))
             .respond_with(
                 ResponseTemplate::new(200)
+                    .insert_header("x-test-state-update", CREATED_STATE)
                     .set_body_json(serde_json::json!({"file_id": "file_123", "upload_url": format!("{}/upload/file_123", server.uri())})),
             )
             .mount(&server)
@@ -340,20 +391,25 @@ mod tests {
         let download_url = format!("{}/download/file_123", server.uri());
         Mock::given(method("POST"))
             .and(path("/backend-api/files/file_123/uploaded"))
+            .and(header("x-test-state", SENT_STATE))
             .respond_with(move |_request: &Request| {
                 if finalize_attempts_responder.fetch_add(1, Ordering::SeqCst) == 0 {
-                    return ResponseTemplate::new(200).set_body_json(serde_json::json!({
-                        "status": "retry"
-                    }));
+                    return ResponseTemplate::new(200)
+                        .insert_header("x-test-state-update", RETRY_STATE)
+                        .set_body_json(serde_json::json!({
+                            "status": "retry"
+                        }));
                 }
 
-                ResponseTemplate::new(200).set_body_json(serde_json::json!({
-                    "status": "success",
-                    "download_url": download_url,
-                    "file_name": "hello.txt",
-                    "mime_type": "text/plain",
-                    "file_size_bytes": 5
-                }))
+                ResponseTemplate::new(200)
+                    .insert_header("x-test-state-update", FINALIZED_STATE)
+                    .set_body_json(serde_json::json!({
+                        "status": "success",
+                        "download_url": download_url,
+                        "file_name": "hello.txt",
+                        "mime_type": "text/plain",
+                        "file_size_bytes": 5
+                    }))
             })
             .mount(&server)
             .await;
@@ -363,7 +419,8 @@ mod tests {
         let path = dir.path().join("hello.txt");
         tokio::fs::write(&path, b"hello").await.expect("write file");
 
-        let uploaded = upload_local_file(&base_url, &chatgpt_auth(), &path)
+        let auth = chatgpt_auth();
+        let uploaded = upload_local_file(&base_url, &auth, &path)
             .await
             .expect("upload succeeds");
 
@@ -376,5 +433,25 @@ mod tests {
         assert_eq!(uploaded.file_name, "hello.txt");
         assert_eq!(uploaded.mime_type, Some("text/plain".to_string()));
         assert_eq!(finalize_attempts.load(Ordering::SeqCst), 2);
+        assert_eq!(
+            *auth
+                .observed_updates
+                .lock()
+                .expect("observed updates lock should not be poisoned"),
+            vec![
+                (
+                    format!("{}/files", base_url.trim_end_matches('/')),
+                    Some(CREATED_STATE.to_string()),
+                ),
+                (
+                    format!("{}/files/file_123/uploaded", base_url.trim_end_matches('/')),
+                    Some(RETRY_STATE.to_string()),
+                ),
+                (
+                    format!("{}/files/file_123/uploaded", base_url.trim_end_matches('/')),
+                    Some(FINALIZED_STATE.to_string()),
+                ),
+            ]
+        );
     }
 }

codex-rs/codex-api/tests/realtime_websocket_e2e.rs

@@ -1,7 +1,11 @@
 use std::collections::HashMap;
 use std::future::Future;
+use std::sync::Arc;
+use std::sync::atomic::AtomicBool;
+use std::sync::atomic::Ordering;
 use std::time::Duration;
 
+use codex_api::AuthProvider;
 use codex_api::Provider;
 use codex_api::RealtimeAudioFrame;
 use codex_api::RealtimeEvent;
@@ -19,11 +23,15 @@ use codex_protocol::protocol::RealtimeVoice;
 use futures::SinkExt;
 use futures::StreamExt;
 use http::HeaderMap;
+use http::HeaderValue;
 use serde_json::Value;
 use serde_json::json;
 use tokio::net::TcpListener;
 use tokio_tungstenite::accept_async;
+use tokio_tungstenite::accept_hdr_async;
 use tokio_tungstenite::tungstenite::Message;
+use tokio_tungstenite::tungstenite::handshake::server::Request;
+use tokio_tungstenite::tungstenite::handshake::server::Response;
 
 type RealtimeWsStream = tokio_tungstenite::WebSocketStream<tokio::net::TcpStream>;
 
@@ -75,6 +83,37 @@ fn test_provider(base_url: String) -> Provider {
     }
 }
 
+struct HandshakeAuthProvider {
+    observed_response: Arc<AtomicBool>,
+}
+
+impl AuthProvider for HandshakeAuthProvider {
+    fn add_auth_headers(&self, _headers: &mut HeaderMap) {}
+
+    fn add_auth_headers_for_url(&self, request_url: &str, headers: &mut HeaderMap) {
+        assert!(request_url.starts_with("ws://127.0.0.1:"));
+        headers.insert("x-test-auth-state", HeaderValue::from_static("sent-state"));
+    }
+
+    fn observe_response_headers(
+        &self,
+        request_url: &str,
+        request_headers: &HeaderMap,
+        response_headers: &HeaderMap,
+    ) {
+        assert!(request_url.starts_with("ws://127.0.0.1:"));
+        assert_eq!(
+            request_headers.get("x-test-auth-state"),
+            Some(&HeaderValue::from_static("sent-state")),
+        );
+        if response_headers.get("x-test-auth-state-update")
+            == Some(&HeaderValue::from_static("rotated-state"))
+        {
+            self.observed_response.store(true, Ordering::Relaxed);
+        }
+    }
+}
+
 #[tokio::test]
 async fn realtime_ws_e2e_session_create_and_event_flow() {
     let (addr, server) = spawn_realtime_ws_server(|mut ws: RealtimeWsStream| async move {
@@ -202,6 +241,66 @@ async fn realtime_ws_e2e_session_create_and_event_flow() {
     server.await.expect("server task");
 }
 
+#[tokio::test]
+async fn realtime_ws_e2e_applies_and_observes_url_aware_auth() {
+    let listener = TcpListener::bind("127.0.0.1:0")
+        .await
+        .expect("listener should bind");
+    let addr = listener.local_addr().expect("local addr");
+    let server = tokio::spawn(async move {
+        let (stream, _) = listener.accept().await.expect("accept");
+        let mut ws = accept_hdr_async(stream, |request: &Request, mut response: Response| {
+            assert_eq!(
+                request.headers().get("x-test-auth-state"),
+                Some(&HeaderValue::from_static("sent-state")),
+            );
+            response.headers_mut().insert(
+                "x-test-auth-state-update",
+                HeaderValue::from_static("rotated-state"),
+            );
+            Ok(response)
+        })
+        .await
+        .expect("accept ws");
+        let first = ws
+            .next()
+            .await
+            .expect("first msg")
+            .expect("first msg ok")
+            .into_text()
+            .expect("text");
+        let first_json: Value = serde_json::from_str(&first).expect("json");
+        assert_eq!(first_json["type"], "session.update");
+    });
+
+    let observed_response = Arc::new(AtomicBool::new(false));
+    let client = RealtimeWebsocketClient::new(test_provider(format!("http://{addr}"))).with_auth(
+        Arc::new(HandshakeAuthProvider {
+            observed_response: Arc::clone(&observed_response),
+        }),
+    );
+    let connection = client
+        .connect(
+            RealtimeSessionConfig {
+                instructions: "backend prompt".to_string(),
+                model: Some("realtime-test-model".to_string()),
+                session_id: Some("conv_123".to_string()),
+                event_parser: RealtimeEventParser::V1,
+                session_mode: RealtimeSessionMode::Conversational,
+                output_modality: RealtimeOutputModality::Audio,
+                voice: RealtimeVoice::Cove,
+            },
+            HeaderMap::new(),
+            HeaderMap::new(),
+        )
+        .await
+        .expect("connect");
+
+    assert!(observed_response.load(Ordering::Relaxed));
+    connection.close().await.expect("close");
+    server.await.expect("server task");
+}
+
 #[tokio::test]
 async fn realtime_ws_connect_webrtc_sideband_retries_join_until_server_is_available() {
     let reserving_listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");