[codex] propagate realtime sideband HTTP state [ci changed_files]

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/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");

codex-rs/core/src/client.rs

@@ -302,17 +302,7 @@ pub(crate) struct RealtimeWebrtcCallStart {
     pub(crate) sdp: String,
     pub(crate) call_id: String,
     pub(crate) sideband_headers: ApiHeaderMap,
-}
-
-/// Reuses the API-auth material that created the WebRTC call for the sideband WebSocket join.
-///
-/// API-key sessions send that API bearer. ChatGPT-auth sessions send their bearer plus account id;
-/// transceiver is responsible for accepting that same call-create identity on the direct
-/// `api.openai.com` sideband path.
-fn sideband_websocket_auth_headers(api_auth: &dyn AuthProvider) -> ApiHeaderMap {
-    let mut headers = ApiHeaderMap::new();
-    api_auth.add_auth_headers(&mut headers);
-    headers
+    pub(crate) sideband_auth: SharedAuthProvider,
 }
 
 impl ModelClient {
@@ -624,10 +614,8 @@ impl ModelClient {
         if let Some(header_value) = self.generate_attestation_header_for().await {
             extra_headers.insert(X_OAI_ATTESTATION_HEADER, header_value);
         }
-        let mut sideband_headers = extra_headers.clone();
-        sideband_headers.extend(sideband_websocket_auth_headers(
-            client_setup.api_auth.as_ref(),
-        ));
+        let sideband_headers = extra_headers.clone();
+        let sideband_auth = Arc::clone(&client_setup.api_auth);
         let transport = ReqwestTransport::new(build_reqwest_client());
         let response =
             ApiRealtimeCallClient::new(transport, client_setup.api_provider, client_setup.api_auth)
@@ -638,6 +626,7 @@ impl ModelClient {
             sdp: response.sdp,
             call_id: response.call_id,
             sideband_headers,
+            sideband_auth,
         })
     }
 

codex-rs/core/src/realtime_conversation.rs

@@ -19,6 +19,7 @@ use codex_api::RealtimeSessionMode;
 use codex_api::RealtimeWebsocketClient;
 use codex_api::RealtimeWebsocketEvents;
 use codex_api::RealtimeWebsocketWriter;
+use codex_api::SharedAuthProvider;
 use codex_api::map_api_error;
 use codex_app_server_protocol::AuthMode;
 use codex_config::config_toml::RealtimeWsMode;
@@ -320,6 +321,7 @@ impl RealtimeConversationManager {
                 session_config,
                 call_id: call.call_id,
                 sideband_headers: call.sideband_headers,
+                sideband_auth: call.sideband_auth,
                 input_channels,
                 events_tx,
                 handoff_state: handoff.clone(),
@@ -1029,6 +1031,7 @@ struct RealtimeWebrtcSidebandInputTask {
     session_config: RealtimeSessionConfig,
     call_id: String,
     sideband_headers: HeaderMap,
+    sideband_auth: SharedAuthProvider,
     input_channels: RealtimeInputChannels,
     events_tx: Sender<RealtimeEvent>,
     handoff_state: RealtimeHandoffState,
@@ -1043,6 +1046,7 @@ fn spawn_webrtc_sideband_input_task(input: RealtimeWebrtcSidebandInputTask) -> J
         session_config,
         call_id,
         sideband_headers,
+        sideband_auth,
         input_channels,
         events_tx,
         handoff_state,
@@ -1057,6 +1061,7 @@ fn spawn_webrtc_sideband_input_task(input: RealtimeWebrtcSidebandInputTask) -> J
         }
 
         let connection = match client
+            .with_auth(sideband_auth)
             .connect_webrtc_sideband(
                 session_config,
                 &call_id,