[codex-api] isolate websocket auth state handling [ci changed_files]

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/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;
@@ -22,7 +23,6 @@ use serde::Deserialize;
 use serde_json::Value;
 use serde_json::map::Map as JsonMap;
 use std::sync::Arc;
-use std::sync::Mutex as SyncMutex;
 use std::sync::OnceLock;
 use std::time::Duration;
 use tokio::net::TcpStream;
@@ -161,50 +161,6 @@ 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: Arc<SyncMutex<HeaderMap>>,
-}
-
-impl WebsocketAuthContext {
-    fn new(auth: SharedAuthProvider, request_url: String, request_headers: HeaderMap) -> Self {
-        Self {
-            auth,
-            request_url,
-            request_headers: Arc::new(SyncMutex::new(request_headers)),
-        }
-    }
-
-    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);
-    }
-
-    fn observe_error_headers(&self, error: &ApiError) {
-        if let ApiError::Transport(TransportError::Http {
-            headers: Some(response_headers),
-            ..
-        }) = error
-        {
-            self.observe_response_headers(response_headers);
-        }
-    }
-
-    fn observe_wrapped_error_headers(&self, error: &WrappedWebsocketErrorEvent) {
-        if let Some(response_headers) = error.headers.as_ref() {
-            self.observe_response_headers(&json_headers_to_http_headers(response_headers));
-        }
-    }
-}
-
 pub struct ResponsesWebsocketConnection {
     stream: Arc<Mutex<Option<WsStream>>>,
     // TODO (pakrym): is this the right place for timeout?
@@ -794,12 +750,18 @@ async fn run_websocket_response_stream(
                     trace!("websocket event: {text}");
                 }
                 if let Some(wrapped_error) = wrapped_error {
-                    auth_context.observe_wrapped_error_headers(&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);
                     }
+                    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,48 @@
+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);
+        }
+    }
+}