From b6fc0b7788613458376e30a257a23bdd7f7117fb Mon Sep 17 00:00:00 2001 From: Cooper Gamble Date: Wed, 3 Jun 2026 06:25:13 +0000 Subject: [PATCH] [codex] propagate realtime sideband HTTP state [ci changed_files] --- .../endpoint/realtime_websocket/methods.rs | 48 ++++++++- .../codex-api/tests/realtime_websocket_e2e.rs | 99 +++++++++++++++++++ codex-rs/core/src/client.rs | 19 +--- codex-rs/core/src/realtime_conversation.rs | 5 + 4 files changed, 152 insertions(+), 19 deletions(-) diff --git a/codex-rs/codex-api/src/endpoint/realtime_websocket/methods.rs b/codex-rs/codex-api/src/endpoint/realtime_websocket/methods.rs index ad26549a26..5df495786e 100644 --- a/codex-rs/codex-api/src/endpoint/realtime_websocket/methods.rs +++ b/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, } 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(), diff --git a/codex-rs/codex-api/tests/realtime_websocket_e2e.rs b/codex-rs/codex-api/tests/realtime_websocket_e2e.rs index cb9d7122f4..5d29a7bfb7 100644 --- a/codex-rs/codex-api/tests/realtime_websocket_e2e.rs +++ b/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; @@ -75,6 +83,37 @@ fn test_provider(base_url: String) -> Provider { } } +struct HandshakeAuthProvider { + observed_response: Arc, +} + +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"); diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index bef48bbc0a..dc112747ff 100644 --- a/codex-rs/core/src/client.rs +++ b/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, }) } diff --git a/codex-rs/core/src/realtime_conversation.rs b/codex-rs/core/src/realtime_conversation.rs index 7f71142e19..1040d9e2d7 100644 --- a/codex-rs/core/src/realtime_conversation.rs +++ b/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, 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,