From 35bb7ab558b80c07bd5cd99cc5e901cad87f1930 Mon Sep 17 00:00:00 2001 From: Cooper Gamble Date: Wed, 3 Jun 2026 05:03:05 +0000 Subject: [PATCH] [codex-api] isolate websocket auth tests [ci changed_files] --- .../src/endpoint/responses_websocket.rs | 203 ------------------ .../src/endpoint/responses_websocket_auth.rs | 129 +++++++++++ 2 files changed, 129 insertions(+), 203 deletions(-) diff --git a/codex-rs/codex-api/src/endpoint/responses_websocket.rs b/codex-rs/codex-api/src/endpoint/responses_websocket.rs index de69ac73d8..70a0365b77 100644 --- a/codex-rs/codex-api/src/endpoint/responses_websocket.rs +++ b/codex-rs/codex-api/src/endpoint/responses_websocket.rs @@ -870,211 +870,8 @@ 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 = "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), - } - } #[test] fn websocket_config_enables_permessage_deflate() { 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 4586be382b..7cad78f5b5 100644 --- a/codex-rs/codex-api/src/endpoint/responses_websocket_auth.rs +++ b/codex-rs/codex-api/src/endpoint/responses_websocket_auth.rs @@ -46,3 +46,132 @@ 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(), + ), + ] + ); + } +}