Preserve WebSocket probe evidence and redact diagnostic details (#57597)

## Why

Handshake probes can lose upgrade evidence when a subsequent read fails, and diagnostic output can expose URLs, server error messages, and close reasons containing sensitive values.

## What changed

- Retain successful upgrade status and allowlisted tracking headers, including on post-upgrade read failures.
- Return structured probe errors that distinguish typed TLS failures from failures with an unknown stage while preserving the underlying API error.
- Recognize empty close frames as immediate closes with code `1005`.
- Redact wire values from probe debug output and doctor error details, and remove URLs and raw errors from connection logs.

## Testing

Add local TLS and WebSocket fixtures covering certificate rejection, ordinary I/O failures, successful and rejected upgrades, immediate closes, malformed frames, header retention, and redaction. Add doctor output snapshots and error-detail redaction assertions.

GitOrigin-RevId: 1ee53720a7a752043943170d8e4da205ec4e55f3

codex-rs/codex-api/src/endpoint/mod.rs

@@ -5,6 +5,7 @@ pub(crate) mod realtime_call;
 pub(crate) mod realtime_websocket;
 pub(crate) mod responses;
 pub(crate) mod responses_websocket;
+mod responses_websocket_probe;
 pub(crate) mod search;
 mod session;
 
@@ -28,7 +29,9 @@ pub use realtime_websocket::session_update_session_json;
 pub use responses::ResponsesClient;
 pub use responses::ResponsesOptions;
 pub use responses_websocket::ResponsesWebsocketClient;
-pub use responses_websocket::ResponsesWebsocketClose;
 pub use responses_websocket::ResponsesWebsocketConnection;
-pub use responses_websocket::ResponsesWebsocketProbe;
+pub use responses_websocket_probe::ResponsesWebsocketClose;
+pub use responses_websocket_probe::ResponsesWebsocketProbe;
+pub use responses_websocket_probe::ResponsesWebsocketProbeError;
+pub use responses_websocket_probe::ResponsesWebsocketProbeFailureStage;
 pub use search::SearchClient;

codex-rs/codex-api/src/endpoint/responses_websocket.rs

@@ -5,6 +5,10 @@ use crate::common::ResponseStream;
 use crate::common::ResponsesWsRequest;
 use crate::common::SafetyBufferingTreatment;
 use crate::common::WS_REQUEST_HEADER_TRACEPARENT_CLIENT_METADATA_KEY;
+use crate::endpoint::responses_websocket_probe::ResponsesWebsocketClose;
+use crate::endpoint::responses_websocket_probe::ResponsesWebsocketProbe;
+use crate::endpoint::responses_websocket_probe::ResponsesWebsocketProbeError;
+use crate::endpoint::responses_websocket_probe::ResponsesWebsocketProbeFailureStage;
 use crate::error::ApiError;
 use crate::provider::Provider;
 use crate::rate_limits::parse_rate_limit_event;
@@ -37,6 +41,7 @@ use tokio::time::Instant;
 use tokio_tungstenite::tungstenite::Error as WsError;
 use tokio_tungstenite::tungstenite::Message;
 use tokio_tungstenite::tungstenite::client::IntoClientRequest;
+use tokio_tungstenite::tungstenite::handshake::client::Response;
 use tokio_tungstenite::tungstenite::protocol::CloseFrame;
 use tracing::Instrument;
 use tracing::Span;
@@ -395,30 +400,6 @@ pub struct ResponsesWebsocketClient {
     auth: SharedAuthProvider,
 }
 
-/// Close frame information captured by a handshake probe.
-#[derive(Debug, Clone, PartialEq, Eq)]
-pub struct ResponsesWebsocketClose {
-    /// WebSocket close code returned by the server.
-    pub code: String,
-    /// Human-readable close reason returned by the server.
-    pub reason: String,
-}
-
-/// Result of a handshake-only Responses WebSocket probe.
-#[derive(Debug, Clone, PartialEq, Eq)]
-pub struct ResponsesWebsocketProbe {
-    /// Redacted by callers before displaying or serializing support reports.
-    pub url: String,
-    /// HTTP status returned by the successful WebSocket upgrade.
-    pub status: StatusCode,
-    /// Whether the server reported reasoning support in the upgrade response.
-    pub reasoning_included: bool,
-    /// Whether the server returned a server-selected model in the upgrade response.
-    pub server_model_present: bool,
-    /// Close frame received immediately after upgrade, when one arrives quickly.
-    pub immediate_close: Option<ResponsesWebsocketClose>,
-}
-
 impl ResponsesWebsocketClient {
     /// Creates a Responses WebSocket client for an already-resolved provider and auth source.
     pub fn new(provider: Provider, auth: SharedAuthProvider) -> Self {
@@ -458,35 +439,67 @@ impl ResponsesWebsocketClient {
         extra_headers: HeaderMap,
         default_headers: HeaderMap,
         immediate_close_timeout: Duration,
-    ) -> Result<ResponsesWebsocketProbe, ApiError> {
+    ) -> Result<ResponsesWebsocketProbe, ResponsesWebsocketProbeError> {
         let ws_url = self
             .provider
             .websocket_url_for_path("/responses")
-            .map_err(|err| ApiError::Stream(format!("failed to build websocket URL: {err}")))?;
+            .map_err(|err| {
+                ResponsesWebsocketProbeError::before_upgrade(
+                    ApiError::Stream(format!("failed to build websocket URL: {err}")),
+                    ResponsesWebsocketProbeFailureStage::Unknown,
+                )
+            })?;
 
         let headers = merge_request_headers(&self.provider.headers, extra_headers, default_headers);
-        let headers = self.auth.apply_auth_headers(headers).await?;
+        let headers = self.auth.apply_auth_headers(headers).await.map_err(|err| {
+            ResponsesWebsocketProbeError::before_upgrade(
+                err.into(),
+                ResponsesWebsocketProbeFailureStage::Unknown,
+            )
+        })?;
 
-        let connector = WebSocketConnector::new(http_client_factory)
-            .map_err(|err| ApiError::Stream(format!("failed to configure websocket TLS: {err}")))?;
-        let (mut stream, status, reasoning_included, server_model) = connect_websocket(
+        let connector = WebSocketConnector::new(http_client_factory).map_err(|err| {
+            ResponsesWebsocketProbeError::before_upgrade(
+                ApiError::Stream(format!("failed to configure websocket TLS: {err}")),
+                ResponsesWebsocketProbeFailureStage::Tls,
+            )
+        })?;
+        let (mut stream, response, reasoning_included, server_model) = connect_websocket(
             ws_url.clone(),
             headers,
             &connector,
             /*turn_state*/ None,
         )
         .await?;
+        let status = response.status();
+        let mut headers = HeaderMap::new();
+        for name in [
+            "x-request-id",
+            "x-oai-request-id",
+            "cf-ray",
+            "x-ms-request-id",
+        ] {
+            for value in response.headers().get_all(name) {
+                headers.append(name, value.clone());
+            }
+        }
         let immediate_close = tokio::time::timeout(immediate_close_timeout, stream.next())
             .await
             .ok()
             .flatten()
             .transpose()
-            .map_err(map_ws_stream_error)?
+            .map_err(|error| ResponsesWebsocketProbeError {
+                error: Box::new(map_ws_stream_error(error)),
+                stage: ResponsesWebsocketProbeFailureStage::Unknown,
+                status: Some(status),
+                headers: Some(headers.clone()),
+            })?
             .and_then(immediate_close_from_message);
 
         Ok(ResponsesWebsocketProbe {
             url: ws_url.to_string(),
             status,
+            headers,
             reasoning_included,
             server_model_present: server_model.is_some(),
             immediate_close,
@@ -498,7 +511,13 @@ fn immediate_close_from_message(message: Message) -> Option<ResponsesWebsocketCl
     let Message::Close(frame) = message else {
         return None;
     };
-    frame.map(close_frame_to_probe)
+    Some(match frame {
+        Some(frame) => close_frame_to_probe(frame),
+        None => ResponsesWebsocketClose {
+            code: "1005".to_string(),
+            reason: String::new(),
+        },
+    })
 }
 
 fn close_frame_to_probe(frame: CloseFrame) -> ResponsesWebsocketClose {
@@ -528,25 +547,46 @@ async fn connect_websocket(
     headers: HeaderMap,
     connector: &WebSocketConnector,
     turn_state: Option<Arc<OnceLock<String>>>,
-) -> Result<(WsStream, StatusCode, bool, Option<String>), ApiError> {
-    info!("connecting to websocket: {url}");
+) -> Result<(WsStream, Response, bool, Option<String>), ResponsesWebsocketProbeError> {
+    info!("connecting to Responses WebSocket");
 
-    let mut request = url
-        .as_str()
-        .into_client_request()
-        .map_err(|err| ApiError::Stream(format!("failed to build websocket request: {err}")))?;
+    let mut request = url.as_str().into_client_request().map_err(|err| {
+        ResponsesWebsocketProbeError::before_upgrade(
+            ApiError::Stream(format!("failed to build websocket request: {err}")),
+            ResponsesWebsocketProbeFailureStage::Unknown,
+        )
+    })?;
     request.headers_mut().extend(headers);
 
     let response = connector.connect(request, websocket_config()).await;
 
     let (stream, response) = match response {
         Ok((stream, response)) => {
-            info!("successfully connected to websocket: {url}");
+            info!(
+                status = response.status().as_u16(),
+                "successfully connected to Responses WebSocket"
+            );
             (stream, response)
         }
         Err(err) => {
-            error!("failed to connect to websocket: {err}, url: {url}");
-            return Err(map_ws_error(err, &url));
+            // Tokio-rustls wraps certificate failures in I/O errors.
+            let is_tls = match &err {
+                WsError::Tls(_) => true,
+                WsError::Io(error) => {
+                    matches!(error.get_ref(), Some(source) if source.is::<rustls::Error>())
+                }
+                _ => false,
+            };
+            let stage = if is_tls {
+                ResponsesWebsocketProbeFailureStage::Tls
+            } else {
+                ResponsesWebsocketProbeFailureStage::Unknown
+            };
+            error!(?stage, "failed to connect to Responses WebSocket");
+            return Err(ResponsesWebsocketProbeError::before_upgrade(
+                map_ws_error(err, &url),
+                stage,
+            ));
         }
     };
 
@@ -566,7 +606,7 @@ async fn connect_websocket(
     }
     Ok((
         WsStream::new(stream),
-        response.status(),
+        response,
         reasoning_included,
         server_model,
     ))

codex-rs/codex-api/src/endpoint/responses_websocket/connector.rs

@@ -35,8 +35,10 @@ impl ResponsesWebsocketClient {
             .map_err(|err| ApiError::Stream(format!("failed to build websocket URL: {err}")))?;
         let headers = merge_request_headers(&self.provider.headers, extra_headers, default_headers);
         let headers = self.auth.apply_auth_headers(headers).await?;
-        let (stream, _status, server_reasoning_included, server_model) =
-            connect_websocket(ws_url, headers, connector, turn_state).await?;
+        let (stream, _response, server_reasoning_included, server_model) =
+            connect_websocket(ws_url, headers, connector, turn_state)
+                .await
+                .map_err(|failure| *failure.error)?;
         Ok(ResponsesWebsocketConnection::new(
             stream,
             self.provider.stream_idle_timeout,

codex-rs/codex-api/src/endpoint/responses_websocket_probe.rs

@@ -0,0 +1,101 @@
+//! Handshake probe evidence. Failure stages require typed evidence, and Debug output hides wire
+//! values that may contain credentials; diagnostic callers must redact displayed values themselves.
+
+use crate::error::ApiError;
+use http::HeaderMap;
+use http::StatusCode;
+
+/// Close frame information captured by a handshake probe.
+#[derive(Clone, PartialEq, Eq)]
+pub struct ResponsesWebsocketClose {
+    /// WebSocket close code returned by the server.
+    pub code: String,
+    /// Human-readable close reason returned by the server.
+    pub reason: String,
+}
+
+impl std::fmt::Debug for ResponsesWebsocketClose {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("ResponsesWebsocketClose")
+            .field("code", &self.code)
+            .finish_non_exhaustive()
+    }
+}
+
+/// Result of a handshake-only Responses WebSocket probe.
+#[derive(Clone, PartialEq, Eq)]
+pub struct ResponsesWebsocketProbe {
+    /// Redacted by callers before displaying or serializing support reports.
+    pub url: String,
+    /// HTTP status returned by the successful WebSocket upgrade.
+    pub status: StatusCode,
+    /// Allowlisted tracking headers from the successful upgrade.
+    pub headers: HeaderMap,
+    /// Whether the server reported reasoning support in the upgrade response.
+    pub reasoning_included: bool,
+    /// Whether the server returned a server-selected model in the upgrade response.
+    pub server_model_present: bool,
+    /// Close frame received immediately after upgrade, when one arrives quickly.
+    pub immediate_close: Option<ResponsesWebsocketClose>,
+}
+
+impl std::fmt::Debug for ResponsesWebsocketProbe {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("ResponsesWebsocketProbe")
+            .field("status", &self.status)
+            .field("reasoning_included", &self.reasoning_included)
+            .field("server_model_present", &self.server_model_present)
+            .field("immediate_close", &self.immediate_close)
+            .finish_non_exhaustive()
+    }
+}
+
+/// A failure stage established by typed evidence at the handshake boundary.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum ResponsesWebsocketProbeFailureStage {
+    /// TLS configuration or a typed TLS handshake failure.
+    Tls,
+    /// The transport did not expose enough evidence to identify a failure stage.
+    Unknown,
+}
+
+/// Probe failure with any HTTP upgrade evidence received before the failure.
+///
+/// Rejected upgrades retain status and headers in the original [`ApiError`]. An error after a
+/// successful upgrade also retains that upgrade's status and allowlisted tracking headers here.
+#[derive(thiserror::Error)]
+#[error("Responses WebSocket probe failed")]
+pub struct ResponsesWebsocketProbeError {
+    /// Original API error, preserving the production connection's error behavior.
+    #[source]
+    pub error: Box<ApiError>,
+    /// Failure stage, when typed evidence establishes one.
+    pub stage: ResponsesWebsocketProbeFailureStage,
+    /// Status received before a subsequent read failure.
+    pub status: Option<StatusCode>,
+    /// Allowlisted tracking headers received before a subsequent read failure.
+    pub headers: Option<HeaderMap>,
+}
+
+impl ResponsesWebsocketProbeError {
+    pub(crate) fn before_upgrade(
+        error: ApiError,
+        stage: ResponsesWebsocketProbeFailureStage,
+    ) -> Self {
+        Self {
+            error: Box::new(error),
+            stage,
+            status: None,
+            headers: None,
+        }
+    }
+}
+
+impl std::fmt::Debug for ResponsesWebsocketProbeError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("ResponsesWebsocketProbeError")
+            .field("stage", &self.stage)
+            .field("status", &self.status)
+            .finish_non_exhaustive()
+    }
+}

codex-rs/codex-api/src/lib.rs

@@ -74,6 +74,8 @@ pub use crate::endpoint::ResponsesWebsocketClient;
 pub use crate::endpoint::ResponsesWebsocketClose;
 pub use crate::endpoint::ResponsesWebsocketConnection;
 pub use crate::endpoint::ResponsesWebsocketProbe;
+pub use crate::endpoint::ResponsesWebsocketProbeError;
+pub use crate::endpoint::ResponsesWebsocketProbeFailureStage;
 pub use crate::endpoint::SearchClient;
 pub use crate::endpoint::session_update_session_json;
 pub use crate::error::ApiError;