diff --git a/codex-rs/codex-api/src/endpoint/responses_websocket.rs b/codex-rs/codex-api/src/endpoint/responses_websocket.rs index 70a0365b77..23134f40b1 100644 --- a/codex-rs/codex-api/src/endpoint/responses_websocket.rs +++ b/codex-rs/codex-api/src/endpoint/responses_websocket.rs @@ -22,6 +22,7 @@ use http::StatusCode; use serde::Deserialize; use serde_json::Value; use serde_json::map::Map as JsonMap; +use std::net::IpAddr; use std::sync::Arc; use std::sync::OnceLock; use std::time::Duration; @@ -32,6 +33,7 @@ use tokio::sync::oneshot; use tokio::time::Instant; use tokio_tungstenite::MaybeTlsStream; use tokio_tungstenite::WebSocketStream; +use tokio_tungstenite::client_async_tls_with_config; use tokio_tungstenite::connect_async_tls_with_config; use tokio_tungstenite::tungstenite::Error as WsError; use tokio_tungstenite::tungstenite::Message; @@ -531,13 +533,30 @@ async fn connect_websocket( .map_err(|err| ApiError::Stream(format!("failed to configure websocket TLS: {err}")))? .map(tokio_tungstenite::Connector::Rustls); - let response = connect_async_tls_with_config( - request, - Some(websocket_config()), - false, // `false` means "do not disable Nagle", which is tungstenite's recommended default. - connector, - ) - .await; + let loopback_address = url.host_str().and_then(|host| { + host.parse::() + .ok() + .filter(IpAddr::is_loopback) + .and_then(|_| url.port_or_known_default().map(|port| (host, port))) + }); + let response = if let Some(address) = loopback_address { + // Explicit loopback providers must not be routed through an ambient proxy. + match TcpStream::connect(address).await { + Ok(stream) => { + client_async_tls_with_config(request, stream, Some(websocket_config()), connector) + .await + } + Err(error) => Err(WsError::Io(error)), + } + } else { + connect_async_tls_with_config( + request, + Some(websocket_config()), + false, // `false` means "do not disable Nagle", which is tungstenite's recommended default. + connector, + ) + .await + }; let (stream, response) = match response { Ok((stream, response)) => { @@ -867,6 +886,10 @@ async fn send_websocket_request( Ok(()) } +#[cfg(test)] +#[path = "responses_websocket_auth_integration_tests.rs"] +mod auth_integration_tests; + #[cfg(test)] mod tests { use super::*; diff --git a/codex-rs/codex-api/src/endpoint/responses_websocket_auth.rs b/codex-rs/codex-api/src/endpoint/responses_websocket_auth.rs index 7cad78f5b5..83d17963ba 100644 --- a/codex-rs/codex-api/src/endpoint/responses_websocket_auth.rs +++ b/codex-rs/codex-api/src/endpoint/responses_websocket_auth.rs @@ -48,130 +48,5 @@ impl WebsocketAuthContext { } #[cfg(test)] -mod tests { - 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, - observed_updates: StdMutex>, - } - - 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(), - ), - ] - ); - } -} +#[path = "responses_websocket_auth_tests.rs"] +mod tests; diff --git a/codex-rs/codex-api/src/endpoint/responses_websocket_auth_integration_tests.rs b/codex-rs/codex-api/src/endpoint/responses_websocket_auth_integration_tests.rs new file mode 100644 index 0000000000..df6d2e143f --- /dev/null +++ b/codex-rs/codex-api/src/endpoint/responses_websocket_auth_integration_tests.rs @@ -0,0 +1,206 @@ +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 = "sent-state"; +const HANDSHAKE_STATE: &str = "handshake-state"; +const WRAPPED_ERROR_STATE: &str = "wrapped-error-state"; + +struct RecordingAuthProvider { + state: StdMutex, + observed_updates: StdMutex>, +} + +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(); + } +} + +#[tokio::test] +async fn connect_attaches_auth_and_observes_handshake_and_retryable_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-test-state") + .and_then(|value| value.to_str().ok()) + .map(ToString::to_string); + response.headers_mut().insert( + "x-test-state-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": 400, + "error": { + "code": WEBSOCKET_CONNECTION_LIMIT_REACHED_CODE, + "message": WEBSOCKET_CONNECTION_LIMIT_REACHED_MESSAGE, + }, + "headers": { + "x-test-state-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::Retryable { message, delay } = error else { + panic!("expected retryable error"); + }; + assert_eq!(message, WEBSOCKET_CONNECTION_LIMIT_REACHED_MESSAGE); + assert_eq!(delay, None); + 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, + HANDSHAKE_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), + } +} diff --git a/codex-rs/codex-api/src/endpoint/responses_websocket_auth_tests.rs b/codex-rs/codex-api/src/endpoint/responses_websocket_auth_tests.rs new file mode 100644 index 0000000000..5c38e2dbc8 --- /dev/null +++ b/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, + observed_updates: StdMutex>, +} + +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(), + ), + ] + ); +}