[codex-api] cover native state transport bypasses [ci changed_files]

2026-10-21 17:48:44 +00:00 · 2026-06-03 09:54:29 +00:00
parent f98c7c473b
commit 6b2c4e8475
3 changed files with 375 additions and 56 deletions
@@ -160,6 +160,41 @@ const OPENAI_MODEL_HEADER: &str = "openai-model";
 const WEBSOCKET_CONNECTION_LIMIT_REACHED_CODE: &str = "websocket_connection_limit_reached";
 const WEBSOCKET_CONNECTION_LIMIT_REACHED_MESSAGE: &str = "Responses websocket connection limit reached (60 minutes). Create a new websocket connection to continue.";

+#[derive(Clone)]
+struct WebsocketAuthContext {
+    auth: SharedAuthProvider,
+    request_url: String,
+    request_headers: HeaderMap,
+}
+
+impl WebsocketAuthContext {
+    fn new(auth: SharedAuthProvider, request_url: String, request_headers: HeaderMap) -> Self {
+        Self {
+            auth,
+            request_url,
+            request_headers,
+        }
+    }
+
+    fn observe_response_headers(&self, response_headers: &HeaderMap) {
+        self.auth.observe_response_headers(
+            &self.request_url,
+            &self.request_headers,
+            response_headers,
+        );
+    }
+
+    fn observe_error_headers(&self, error: &ApiError) {
+        if let ApiError::Transport(TransportError::Http {
+            headers: Some(response_headers),
+            ..
+        }) = error
+        {
+            self.observe_response_headers(response_headers);
+        }
+    }
+}
+
 pub struct ResponsesWebsocketConnection {
    stream: Arc<Mutex<Option<WsStream>>>,
    // TODO (pakrym): is this the right place for timeout?
@@ -167,6 +202,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 +214,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 +227,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 +236,7 @@ impl ResponsesWebsocketConnection {
            server_reasoning_included,
            models_etag,
            server_model,
+            auth_context,
            telemetry,
        }
    }
@@ -258,6 +297,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 +337,7 @@ impl ResponsesWebsocketConnection {
                        tx_event.clone(),
                        request_body,
                        idle_timeout,
+                        &auth_context,
                        telemetry,
                        connection_reused,
                    )
@@ -383,7 +424,8 @@ impl ResponsesWebsocketClient {
            merge_request_headers(&self.provider.headers, extra_headers, default_headers);
        self.auth
            .add_auth_headers_for_url(&request_url, &mut headers);
-        let request_headers = headers.clone();
+        let auth_context =
+            WebsocketAuthContext::new(self.auth.clone(), request_url, headers.clone());

        let (
            stream,
@@ -395,21 +437,16 @@ impl ResponsesWebsocketClient {
        ) = connect_websocket(ws_url, headers, turn_state.clone())
            .await
            .inspect_err(|error| {
-                observe_websocket_auth_error_headers(
-                    self.auth.as_ref(),
-                    &request_url,
-                    &request_headers,
-                    error,
-                );
+                auth_context.observe_error_headers(error);
            })?;
-        self.auth
-            .observe_response_headers(&request_url, &request_headers, &response_headers);
+        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,
        ))
    }
@@ -437,21 +474,16 @@ impl ResponsesWebsocketClient {
            merge_request_headers(&self.provider.headers, extra_headers, default_headers);
        self.auth
            .add_auth_headers_for_url(&request_url, &mut headers);
-        let request_headers = headers.clone();
+        let auth_context =
+            WebsocketAuthContext::new(self.auth.clone(), request_url, headers.clone());

        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| {
-                    observe_websocket_auth_error_headers(
-                        self.auth.as_ref(),
-                        &request_url,
-                        &request_headers,
-                        error,
-                    );
+                    auth_context.observe_error_headers(error);
                })?;
-        self.auth
-            .observe_response_headers(&request_url, &request_headers, &response_headers);
+        auth_context.observe_response_headers(&response_headers);
        let immediate_close = tokio::time::timeout(immediate_close_timeout, stream.next())
            .await
            .ok()
@@ -473,21 +505,6 @@ impl ResponsesWebsocketClient {
    }
 }

-fn observe_websocket_auth_error_headers(
-    auth: &dyn crate::auth::AuthProvider,
-    request_url: &str,
-    request_headers: &HeaderMap,
-    error: &ApiError,
-) {
-    if let ApiError::Transport(TransportError::Http {
-        headers: Some(response_headers),
-        ..
-    }) = error
-    {
-        auth.observe_response_headers(request_url, request_headers, response_headers);
-    }
-}
-
 fn immediate_close_from_message(message: Message) -> Option<ResponsesWebsocketClose> {
    let Message::Close(frame) = message else {
        return None;
@@ -721,6 +738,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> {
@@ -764,6 +782,7 @@ async fn run_websocket_response_stream(
                    && let Some(error) =
                        map_wrapped_websocket_error_event(wrapped_error, text.to_string())
                {
+                    auth_context.observe_error_headers(&error);
                    return Err(error);
                }

@@ -873,8 +892,186 @@ async fn send_websocket_request(
 #[cfg(test)]
 mod tests {
    use super::*;
+    use crate::auth::AuthProvider;
+    use crate::provider::RetryConfig;
    use pretty_assertions::assert_eq;
    use serde_json::json;
+    use std::sync::Mutex as StdMutex;
+    use tokio::net::TcpListener;
+    use tokio_tungstenite::accept_hdr_async_with_config;
+
+    const SENT_STATE: &str = "ois1.sent.nonce.ciphertext";
+    const HANDSHAKE_STATE: &str = "ois1.handshake.nonce.ciphertext";
+    const WRAPPED_ERROR_STATE: &str = "ois1.error.nonce.ciphertext";
+
+    #[derive(Default)]
+    struct RecordingAuthProvider {
+        observed_updates: StdMutex<Vec<(String, String, String)>>,
+    }
+
+    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) {
+            headers.insert("x-oai-is", HeaderValue::from_static(SENT_STATE));
+        }
+
+        fn observe_response_headers(
+            &self,
+            request_url: &str,
+            request_headers: &HeaderMap,
+            response_headers: &HeaderMap,
+        ) {
+            let Some(sent_state) = request_headers
+                .get("x-oai-is")
+                .and_then(|value| value.to_str().ok())
+            else {
+                return;
+            };
+            let Some(update_state) = response_headers
+                .get("x-oai-is-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(),
+                ));
+        }
+    }
+
+    #[tokio::test]
+    async fn connect_attaches_auth_and_observes_handshake_and_wrapped_error_headers() {
+        let listener = TcpListener::bind("127.0.0.1:0")
+            .await
+            .expect("bind websocket server");
+        let address = listener.local_addr().expect("websocket server address");
+        let observed_request_state = Arc::new(StdMutex::new(None));
+        let server_request_state = Arc::clone(&observed_request_state);
+        let server_task = tokio::spawn(async move {
+            let (stream, _) = listener.accept().await.expect("accept websocket client");
+            let mut stream = accept_hdr_async_with_config(
+                stream,
+                move |request: &tokio_tungstenite::tungstenite::handshake::server::Request,
+                      mut response: tokio_tungstenite::tungstenite::handshake::server::Response| {
+                    *server_request_state
+                        .lock()
+                        .expect("request state lock should not be poisoned") = request
+                        .headers()
+                        .get("x-oai-is")
+                        .and_then(|value| value.to_str().ok())
+                        .map(ToString::to_string);
+                    response.headers_mut().insert(
+                        "x-oai-is-update",
+                        HeaderValue::from_static(HANDSHAKE_STATE),
+                    );
+                    Ok(response)
+                },
+                Some(websocket_config()),
+            )
+            .await
+            .expect("complete websocket handshake");
+            let message = stream
+                .next()
+                .await
+                .expect("receive websocket request")
+                .expect("read websocket request");
+            assert!(matches!(message, Message::Text(_)));
+            stream
+                .send(Message::Text(
+                    json!({
+                        "type": "error",
+                        "status": 401,
+                        "headers": {
+                            "x-oai-is-update": WRAPPED_ERROR_STATE,
+                        },
+                    })
+                    .to_string()
+                    .into(),
+                ))
+                .await
+                .expect("send wrapped websocket error");
+        });
+
+        let auth = Arc::new(RecordingAuthProvider::default());
+        let request_url = format!("ws://{address}/backend-api/codex/responses");
+        let client = ResponsesWebsocketClient::new(websocket_provider(address), auth.clone());
+        let connection = client
+            .connect(
+                HeaderMap::new(),
+                HeaderMap::new(),
+                /*turn_state*/ None,
+                /*telemetry*/ None,
+            )
+            .await
+            .expect("connect websocket client");
+        let mut response_stream = connection
+            .stream_request(
+                ResponsesWsRequest::ResponseProcessed(ResponseProcessedWsRequest {
+                    response_id: "response-id".to_string(),
+                }),
+                /*connection_reused*/ false,
+            )
+            .await
+            .expect("start websocket response stream");
+        let error = response_stream
+            .next()
+            .await
+            .expect("receive websocket response")
+            .expect_err("wrapped error should fail the response stream");
+
+        server_task.await.expect("websocket server task");
+        let ApiError::Transport(TransportError::Http { status, .. }) = error else {
+            panic!("expected transport HTTP error");
+        };
+        assert_eq!(status, StatusCode::UNAUTHORIZED);
+        assert_eq!(
+            *observed_request_state
+                .lock()
+                .expect("request state lock should not be poisoned"),
+            Some(SENT_STATE.to_string())
+        );
+        assert_eq!(
+            *auth
+                .observed_updates
+                .lock()
+                .expect("observed updates lock should not be poisoned"),
+            vec![
+                (
+                    request_url.clone(),
+                    SENT_STATE.to_string(),
+                    HANDSHAKE_STATE.to_string(),
+                ),
+                (
+                    request_url,
+                    SENT_STATE.to_string(),
+                    WRAPPED_ERROR_STATE.to_string(),
+                ),
+            ]
+        );
+    }
+
+    fn websocket_provider(address: std::net::SocketAddr) -> Provider {
+        Provider {
+            name: "test".to_string(),
+            base_url: format!("http://{address}/backend-api/codex"),
+            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),
+        }
+    }

    #[test]
    fn websocket_config_enables_permessage_deflate() {
@@ -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 = "ois1.sent.nonce.ciphertext";
+    const CREATED_STATE: &str = "ois1.created.nonce.ciphertext";
+    const RETRY_STATE: &str = "ois1.retry.nonce.ciphertext";
+    const FINALIZED_STATE: &str = "ois1.finalized.nonce.ciphertext";
+    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-oai-is", 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-oai-is-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-oai-is", 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-oai-is-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-oai-is", 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-oai-is-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-oai-is-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()),
+                ),
+            ]
+        );
    }
 }