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

use crate::auth::SharedAuthProvider;
use crate::common::ResponseEvent;
use crate::common::ResponseStream;
use crate::common::ResponsesWsRequest;
use crate::common::SafetyBufferingTreatment;
use crate::common::WS_REQUEST_HEADER_TRACEPARENT_CLIENT_METADATA_KEY;
use crate::error::ApiError;
use crate::provider::Provider;
use crate::rate_limits::parse_rate_limit_event;
use crate::responses_headers::json_headers_to_http_headers;
use crate::safety_buffering::treatment_from_headers;
use crate::sse::ResponsesStreamEvent;
use crate::sse::process_responses_event;
use crate::telemetry::WebsocketTelemetry;
use codex_client::TransportError;
use codex_http_client::HttpClientFactory;
use codex_http_client::RetryAfter;
use codex_websocket_client::WebSocketConnection;
use codex_websocket_client::WebSocketConnector;
use futures::FutureExt;
use futures::SinkExt;
use futures::StreamExt;
use http::HeaderMap;
use http::StatusCode;
use serde::Deserialize;
use serde_json::Value;
use serde_json::map::Map as JsonMap;
use std::sync::Arc;
use std::sync::OnceLock;
use std::time::Duration;
use tokio::sync::Mutex;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
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::protocol::CloseFrame;
use tracing::Instrument;
use tracing::Span;
use tracing::debug;
use tracing::error;
use tracing::info;
use tracing::instrument;
use tungstenite::extensions::ExtensionsConfig;
use tungstenite::extensions::compression::deflate::DeflateConfig;
use tungstenite::protocol::WebSocketConfig;
use url::Url;

struct WsStream {
    tx_command: mpsc::Sender<WsCommand>,
    rx_message: mpsc::UnboundedReceiver<Result<Message, WsError>>,
    pump_task: tokio::task::JoinHandle<()>,
}

enum WsCommand {
    Send {
        message: Message,
        tx_result: oneshot::Sender<Result<(), WsError>>,
    },
}

impl WsStream {
    fn new(inner: WebSocketConnection) -> Self {
        let (tx_command, mut rx_command) = mpsc::channel::<WsCommand>(32);
        let (tx_message, rx_message) = mpsc::unbounded_channel::<Result<Message, WsError>>();

        let pump_task = tokio::spawn(async move {
            let mut inner = inner;
            loop {
                tokio::select! {
                    command = rx_command.recv() => {
                        let Some(command) = command else {
                            break;
                        };
                        match command {
                            WsCommand::Send { message, tx_result } => {
                                let result = inner.send(message).await;
                                let should_break = result.is_err();
                                let _ = tx_result.send(result);
                                if should_break {
                                    break;
                                }
                            }
                        }
                    }
                    message = inner.next() => {
                        let Some(message) = message else {
                            break;
                        };
                        match message {
                            Ok(Message::Ping(payload)) => {
                                if let Err(err) = inner.send(Message::Pong(payload)).await {
                                    let _ = tx_message.send(Err(err));
                                    break;
                                }
                            }
                            Ok(Message::Pong(_)) => {}
                            Ok(message @ (Message::Text(_)
                            | Message::Binary(_)
                            | Message::Close(_)
                            | Message::Frame(_))) => {
                                let is_close = matches!(message, Message::Close(_));
                                if tx_message.send(Ok(message)).is_err() {
                                    break;
                                }
                                if is_close {
                                    break;
                                }
                            }
                            Err(err) => {
                                let _ = tx_message.send(Err(err));
                                break;
                            }
                        }
                    }
                }
            }
        });

        Self {
            tx_command,
            rx_message,
            pump_task,
        }
    }

    async fn request(
        &mut self,
        make_command: impl FnOnce(oneshot::Sender<Result<(), WsError>>) -> WsCommand,
    ) -> Result<(), WsError> {
        let (tx_result, rx_result) = oneshot::channel();
        if self.tx_command.send(make_command(tx_result)).await.is_ok()
            && let Ok(result) = rx_result.await
        {
            return result;
        }
        while let Some(message) = self.rx_message.recv().await {
            message?;
        }
        Err(WsError::ConnectionClosed)
    }

    async fn send(&mut self, message: Message) -> Result<(), WsError> {
        self.request(|tx_result| WsCommand::Send { message, tx_result })
            .await
    }

    async fn next(&mut self) -> Option<Result<Message, WsError>> {
        self.rx_message.recv().await
    }
}

impl Drop for WsStream {
    fn drop(&mut self) {
        self.pump_task.abort();
    }
}

const X_CODEX_TURN_STATE_HEADER: &str = "x-codex-turn-state";
const X_MODELS_ETAG_HEADER: &str = "x-models-etag";
const X_REASONING_INCLUDED_HEADER: &str = "x-reasoning-included";
const OPENAI_MODEL_HEADER: &str = "openai-model";
const WEBSOCKET_CONNECTION_LIMIT_REACHED_CODE: &str = "websocket_connection_limit_reached";
const WEBSOCKET_CONNECTION_LIMIT_REACHED_MESSAGE: &str = "Responses websocket connection limit reached (60 minutes). Create a new websocket connection to continue.";
const PREVIOUS_RESPONSE_NOT_FOUND_CODE: &str = "previous_response_not_found";
const PREVIOUS_RESPONSE_NOT_FOUND_MESSAGE: &str =
    "Previous response was not found. Retrying the full request.";
const RESPONSES_WEBSOCKET_TIMING_KIND: &str = "responsesapi.websocket_timing";
const RESPONSES_WEBSOCKET_TIMING_EVENT_TARGET: &str = "codex_api::responses_websocket_timing";
const SESSION_ID_CLIENT_METADATA_KEY: &str = "session_id";
const THREAD_ID_CLIENT_METADATA_KEY: &str = "thread_id";
const TURN_ID_CLIENT_METADATA_KEY: &str = "turn_id";
const WS_STREAM_REQUEST_START_MS_CLIENT_METADATA_KEY: &str = "x-codex-ws-stream-request-start-ms";

struct ResponsesWebsocketTimingLogContext {
    model: String,
    session_id: Option<String>,
    thread_id: Option<String>,
    turn_id: Option<String>,
    traceparent: Option<String>,
    previous_response_id: Option<String>,
    request_start_ms: Option<String>,
    warmup: bool,
    connection_reused: bool,
}

pub struct ResponsesWebsocketConnection {
    stream: Arc<Mutex<Option<WsStream>>>,
    // TODO (pakrym): is this the right place for timeout?
    idle_timeout: Duration,
    server_reasoning_included: bool,
    server_model: Option<String>,
    telemetry: Option<Arc<dyn WebsocketTelemetry>>,
}

impl std::fmt::Debug for ResponsesWebsocketConnection {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ResponsesWebsocketConnection")
            .field("stream", &"<ws-stream>")
            .field("idle_timeout", &self.idle_timeout)
            .field("server_reasoning_included", &self.server_reasoning_included)
            .field("server_model", &self.server_model)
            .field("telemetry", &self.telemetry.as_ref().map(|_| "<telemetry>"))
            .finish()
    }
}

impl ResponsesWebsocketConnection {
    fn new(
        stream: WsStream,
        idle_timeout: Duration,
        server_reasoning_included: bool,
        server_model: Option<String>,
        telemetry: Option<Arc<dyn WebsocketTelemetry>>,
    ) -> Self {
        Self {
            stream: Arc::new(Mutex::new(Some(stream))),
            idle_timeout,
            server_reasoning_included,
            server_model,
            telemetry,
        }
    }

    pub async fn is_closed(&self) -> bool {
        self.stream
            .lock()
            .await
            .as_ref()
            .is_none_or(|stream| stream.pump_task.is_finished())
    }

    #[instrument(
        name = "responses_websocket.stream_request",
        level = "info",
        skip_all,
        fields(transport = "responses_websocket", api.path = "/responses")
    )]
    pub async fn stream_request(
        &self,
        request: ResponsesWsRequest<'_>,
        connection_reused: bool,
        turn_state: Option<Arc<OnceLock<String>>>,
    ) -> Result<ResponseStream, ApiError> {
        let (tx_event, rx_event) =
            mpsc::channel::<std::result::Result<ResponseEvent, ApiError>>(1600);
        let (tx_interrupt, rx_interrupt) = oneshot::channel();
        let stream = Arc::clone(&self.stream);
        let idle_timeout = self.idle_timeout;
        let server_reasoning_included = self.server_reasoning_included;
        let server_model = self.server_model.clone();
        let telemetry = self.telemetry.clone();
        let ResponsesWsRequest::ResponseCreate(ws_request) = &request;
        let client_metadata = ws_request.client_metadata.as_ref();
        let timing_log_context = ResponsesWebsocketTimingLogContext {
            model: ws_request.model.to_string(),
            session_id: client_metadata
                .and_then(|metadata| metadata.get(SESSION_ID_CLIENT_METADATA_KEY))
                .cloned(),
            thread_id: client_metadata
                .and_then(|metadata| metadata.get(THREAD_ID_CLIENT_METADATA_KEY))
                .cloned(),
            turn_id: client_metadata
                .and_then(|metadata| metadata.get(TURN_ID_CLIENT_METADATA_KEY))
                .cloned(),
            traceparent: client_metadata
                .and_then(|metadata| {
                    metadata.get(WS_REQUEST_HEADER_TRACEPARENT_CLIENT_METADATA_KEY)
                })
                .cloned(),
            previous_response_id: ws_request.previous_response_id.clone(),
            request_start_ms: client_metadata
                .and_then(|metadata| metadata.get(WS_STREAM_REQUEST_START_MS_CLIENT_METADATA_KEY))
                .cloned(),
            warmup: ws_request.generate == Some(false),
            connection_reused,
        };
        let request_text = serialize_websocket_request(&request)?;

        let current_span = Span::current();
        tokio::spawn(
            #[expect(
                clippy::await_holding_invalid_type,
                reason = "the guard serializes exclusive use of the websocket stream for the lifetime of the response stream"
            )]
            async move {
                if let Some(model) = server_model {
                    let _ = tx_event.send(Ok(ResponseEvent::ServerModel(model))).await;
                }
                if server_reasoning_included {
                    let _ = tx_event
                        .send(Ok(ResponseEvent::ServerReasoningIncluded(true)))
                        .await;
                }
                let mut guard = tokio::select! {
                    biased;
                    _ = tx_event.closed() => return,
                    guard = stream.lock() => guard,
                };
                if tx_event.is_closed() {
                    return;
                }
                let result = {
                    let Some(ws_stream) = guard.as_mut() else {
                        let _ = tx_event
                            .send(Err(ApiError::Stream(
                                "websocket connection is closed".to_string(),
                            )))
                            .await;
                        return;
                    };

                    tokio::select! {
                        biased;
                        result = run_websocket_response_stream(
                            ws_stream,
                            tx_event.clone(),
                            request_text,
                            idle_timeout,
                            telemetry,
                            turn_state.as_deref(),
                            &timing_log_context,
                            rx_interrupt,
                        ) => result,
                        _ = tx_event.closed() => Err(ApiError::Stream(
                            "response event consumer dropped".to_string(),
                        )),
                    }
                };

                if let Err(err) = result {
                    // A terminal stream error should reach the caller immediately. Waiting for a
                    // graceful close handshake here can stall indefinitely and mask the error.
                    let failed_stream = guard.take();
                    drop(guard);
                    drop(failed_stream);
                    let _ = tx_event.send(Err(err)).await;
                }
            }
            .instrument(current_span),
        );

        Ok(ResponseStream {
            rx_event,
            upstream_request_id: None,
            interrupt: Some(tx_interrupt),
        })
    }
}

mod connector;

/// Client for connecting to the Responses WebSocket endpoint for one provider.
pub struct ResponsesWebsocketClient {
    provider: Provider,
    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 {
        Self { provider, auth }
    }

    pub async fn connect(
        &self,
        http_client_factory: &HttpClientFactory,
        extra_headers: HeaderMap,
        default_headers: HeaderMap,
        turn_state: Option<Arc<OnceLock<String>>>,
        telemetry: Option<Arc<dyn WebsocketTelemetry>>,
    ) -> Result<ResponsesWebsocketConnection, ApiError> {
        let connector = WebSocketConnector::new(http_client_factory)
            .map_err(|err| ApiError::Stream(format!("failed to configure websocket TLS: {err}")))?;
        self.connect_with_connector(
            &connector,
            extra_headers,
            default_headers,
            turn_state,
            telemetry,
        )
        .await
    }

    /// Opens a WebSocket connection long enough to validate the upgrade response.
    ///
    /// The probe uses the same URL construction, headers, authentication, TLS,
    /// and custom-CA path as a real Responses WebSocket connection, but it does
    /// not send a request frame. After the HTTP 101 upgrade succeeds, it waits
    /// briefly for an immediate server close frame so diagnostics can distinguish
    /// a usable connection from a policy rejection that closes right away.
    pub async fn probe_handshake(
        &self,
        http_client_factory: &HttpClientFactory,
        extra_headers: HeaderMap,
        default_headers: HeaderMap,
        immediate_close_timeout: Duration,
    ) -> Result<ResponsesWebsocketProbe, ApiError> {
        let ws_url = self
            .provider
            .websocket_url_for_path("/responses")
            .map_err(|err| ApiError::Stream(format!("failed to build websocket URL: {err}")))?;

        let mut headers =
            merge_request_headers(&self.provider.headers, extra_headers, default_headers);
        self.auth.add_auth_headers(&mut headers);

        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(
            ws_url.clone(),
            headers,
            &connector,
            /*turn_state*/ None,
        )
        .await?;
        let immediate_close = tokio::time::timeout(immediate_close_timeout, stream.next())
            .await
            .ok()
            .flatten()
            .transpose()
            .map_err(map_ws_stream_error)?
            .and_then(immediate_close_from_message);

        Ok(ResponsesWebsocketProbe {
            url: ws_url.to_string(),
            status,
            reasoning_included,
            server_model_present: server_model.is_some(),
            immediate_close,
        })
    }
}

fn immediate_close_from_message(message: Message) -> Option<ResponsesWebsocketClose> {
    let Message::Close(frame) = message else {
        return None;
    };
    frame.map(close_frame_to_probe)
}

fn close_frame_to_probe(frame: CloseFrame) -> ResponsesWebsocketClose {
    ResponsesWebsocketClose {
        code: frame.code.to_string(),
        reason: frame.reason.to_string(),
    }
}

fn merge_request_headers(
    provider_headers: &HeaderMap,
    extra_headers: HeaderMap,
    default_headers: HeaderMap,
) -> HeaderMap {
    let mut headers = provider_headers.clone();
    headers.extend(extra_headers);
    for (name, value) in &default_headers {
        if let http::header::Entry::Vacant(entry) = headers.entry(name) {
            entry.insert(value.clone());
        }
    }
    headers
}

async fn connect_websocket(
    url: Url,
    headers: HeaderMap,
    connector: &WebSocketConnector,
    turn_state: Option<Arc<OnceLock<String>>>,
) -> Result<(WsStream, StatusCode, bool, Option<String>), ApiError> {
    info!("connecting to websocket: {url}");

    let mut request = url
        .as_str()
        .into_client_request()
        .map_err(|err| ApiError::Stream(format!("failed to build websocket request: {err}")))?;
    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}");
            (stream, response)
        }
        Err(err) => {
            error!("failed to connect to websocket: {err}, url: {url}");
            return Err(map_ws_error(err, &url));
        }
    };

    let reasoning_included = response.headers().contains_key(X_REASONING_INCLUDED_HEADER);
    let server_model = response
        .headers()
        .get(OPENAI_MODEL_HEADER)
        .and_then(|value| value.to_str().ok())
        .map(ToString::to_string);
    if let Some(turn_state) = turn_state
        && let Some(header_value) = response
            .headers()
            .get(X_CODEX_TURN_STATE_HEADER)
            .and_then(|value| value.to_str().ok())
    {
        let _ = turn_state.set(header_value.to_string());
    }
    Ok((
        WsStream::new(stream),
        response.status(),
        reasoning_included,
        server_model,
    ))
}

fn websocket_config() -> WebSocketConfig {
    let mut extensions = ExtensionsConfig::default();
    extensions.permessage_deflate = Some(DeflateConfig::default());

    let mut config = WebSocketConfig::default();
    config.extensions = extensions;
    config
}

fn map_ws_stream_error(error: WsError) -> ApiError {
    match codex_websocket_client::network_policy_denial(&error) {
        Some(denied) => ApiError::Transport(TransportError::Policy(denied)),
        None => ApiError::Stream(error.to_string()),
    }
}

fn map_ws_error(err: WsError, url: &Url) -> ApiError {
    if let Some(denied) = codex_websocket_client::network_policy_denial(&err) {
        return ApiError::Transport(TransportError::Policy(denied));
    }
    match err {
        WsError::Http(response) => {
            let status = response.status();
            let headers = response.headers().clone();
            let retry_after = RetryAfter::from_headers(&headers);
            let body = response
                .body()
                .as_ref()
                .and_then(|bytes| String::from_utf8(bytes.clone()).ok());
            ApiError::Transport(TransportError::Http {
                status,
                url: Some(url.to_string()),
                headers: Some(headers),
                body,
                retry_after,
            })
        }
        WsError::ConnectionClosed | WsError::AlreadyClosed => {
            ApiError::Stream("websocket closed".to_string())
        }
        WsError::Io(err) => ApiError::Transport(TransportError::Network(err.to_string())),
        other => ApiError::Transport(TransportError::Network(other.to_string())),
    }
}

#[derive(Debug, Deserialize)]
struct WrappedWebsocketError {
    code: Option<String>,
    message: Option<String>,
    #[serde(default)]
    headers: Option<Value>,
}

#[derive(Debug, Deserialize)]
struct WrappedWebsocketErrorEvent {
    #[serde(rename = "type")]
    kind: String,
    #[serde(alias = "status_code")]
    status: Option<u16>,
    #[serde(default)]
    error: Option<WrappedWebsocketError>,
    #[serde(default)]
    headers: Option<JsonMap<String, Value>>,
}

fn parse_wrapped_websocket_error_event(payload: &str) -> Option<WrappedWebsocketErrorEvent> {
    let event: WrappedWebsocketErrorEvent = serde_json::from_str(payload).ok()?;
    if event.kind != "error" {
        return None;
    }
    Some(event)
}

fn map_wrapped_websocket_error_event(
    event: WrappedWebsocketErrorEvent,
    original_payload: String,
) -> Option<ApiError> {
    let WrappedWebsocketErrorEvent {
        status,
        error,
        headers,
        ..
    } = event;
    let retry_after = [
        error
            .as_ref()
            .and_then(|error| error.headers.as_ref())
            .and_then(Value::as_object),
        headers.as_ref(),
    ]
    .into_iter()
    .flatten()
    .find_map(|headers| RetryAfter::from_headers(&json_headers_to_http_headers(headers)));

    if let Some(error) = error.as_ref()
        && let Some(code) = error.code.as_deref()
        && let Some(fallback_message) = match code {
            WEBSOCKET_CONNECTION_LIMIT_REACHED_CODE => {
                Some(WEBSOCKET_CONNECTION_LIMIT_REACHED_MESSAGE)
            }
            PREVIOUS_RESPONSE_NOT_FOUND_CODE => Some(PREVIOUS_RESPONSE_NOT_FOUND_MESSAGE),
            _ => None,
        }
    {
        return Some(ApiError::Retryable {
            message: error
                .message
                .clone()
                .unwrap_or_else(|| fallback_message.to_string()),
            retry_after,
        });
    }

    let status = StatusCode::from_u16(status?).ok()?;
    if status.is_success() {
        return None;
    }

    Some(ApiError::Transport(TransportError::Http {
        status,
        url: None,
        headers: headers.as_ref().map(json_headers_to_http_headers),
        body: Some(original_payload),
        retry_after,
    }))
}

#[allow(clippy::too_many_arguments)]
async fn run_websocket_response_stream(
    ws_stream: &mut WsStream,
    tx_event: mpsc::Sender<std::result::Result<ResponseEvent, ApiError>>,
    request_text: String,
    idle_timeout: Duration,
    telemetry: Option<Arc<dyn WebsocketTelemetry>>,
    turn_state: Option<&OnceLock<String>>,
    timing_log_context: &ResponsesWebsocketTimingLogContext,
    interrupt: oneshot::Receiver<()>,
) -> Result<(), ApiError> {
    let mut last_server_model: Option<String> = None;
    let mut safety_buffering_treatment = SafetyBufferingTreatment::default();
    send_websocket_request(
        ws_stream,
        request_text,
        idle_timeout,
        telemetry.as_ref(),
        timing_log_context.connection_reused,
    )
    .await?;

    // A response owns its interrupt, and create must be sent before interrupt.
    let mut interrupt = interrupt.fuse();
    let mut response_id = None;
    loop {
        let poll_start = Instant::now();
        let response = tokio::select! {
            response = tokio::time::timeout(idle_timeout, ws_stream.next()) => {
                response.map_err(|_| ApiError::Stream("idle timeout waiting for websocket".into()))
            }
            Ok(()) = &mut interrupt, if response_id.is_some() => {
                send_websocket_request(
                    ws_stream,
                    serde_json::json!({
                        "type": "response.interrupt",
                        "response_id": response_id,
                        "mode": "discard_partial_items",
                    }).to_string(),
                    idle_timeout,
                    /*telemetry*/ None,
                    timing_log_context.connection_reused,
                )
                .await?;
                continue;
            }
        };
        if let Some(t) = telemetry.as_ref() {
            t.on_ws_event(&response, poll_start.elapsed());
        }
        let message = match response {
            Ok(Some(Ok(msg))) => msg,
            Ok(Some(Err(err))) => {
                return Err(map_ws_stream_error(err));
            }
            Ok(None) => {
                return Err(ApiError::Stream(
                    "stream closed before response.completed".into(),
                ));
            }
            Err(err) => {
                return Err(err);
            }
        };

        match message {
            Message::Text(text) => {
                if let Some(wrapped_error) = parse_wrapped_websocket_error_event(&text)
                    && let Some(error) =
                        map_wrapped_websocket_error_event(wrapped_error, text.to_string())
                {
                    return Err(error);
                }

                let event = match serde_json::from_str::<ResponsesStreamEvent>(&text) {
                    Ok(event) => event,
                    Err(err) => {
                        debug!("failed to parse websocket event: {err}, data: {text}");
                        continue;
                    }
                };
                emit_responses_websocket_timing_event(
                    event.kind(),
                    text.as_str(),
                    timing_log_context,
                );
                if event.kind() == "codex.response.metadata"
                    && let Some(etag) =
                        event
                            .headers
                            .as_ref()
                            .and_then(Value::as_object)
                            .and_then(|headers| {
                                json_headers_to_http_headers(headers)
                                    .get(X_MODELS_ETAG_HEADER)
                                    .and_then(|value| value.to_str().ok())
                                    .map(str::to_string)
                            })
                {
                    let _ = tx_event.send(Ok(ResponseEvent::ModelsEtag(etag))).await;
                }
                if let Some(response_turn_state) = event.turn_state()
                    && let Some(turn_state) = turn_state
                {
                    let _ = turn_state.set(response_turn_state);
                }
                let model_verifications = event.model_verifications();
                let turn_moderation_metadata = event.turn_moderation_metadata();
                let safety_buffering =
                    safety_buffering_for_event(&event, &mut safety_buffering_treatment);
                if event.kind() == "codex.rate_limits" {
                    if let Some(snapshot) = parse_rate_limit_event(&text) {
                        let _ = tx_event.send(Ok(ResponseEvent::RateLimits(snapshot))).await;
                    }
                    continue;
                }
                if let Some(model) = event.response_model()
                    && last_server_model.as_deref() != Some(model.as_str())
                {
                    let _ = tx_event
                        .send(Ok(ResponseEvent::ServerModel(model.clone())))
                        .await;
                    last_server_model = Some(model);
                }
                if let Some(verifications) = model_verifications
                    && tx_event
                        .send(Ok(ResponseEvent::ModelVerifications(verifications)))
                        .await
                        .is_err()
                {
                    return Err(ApiError::Stream(
                        "response event consumer dropped".to_string(),
                    ));
                }
                if let Some(metadata) = turn_moderation_metadata
                    && tx_event
                        .send(Ok(ResponseEvent::TurnModerationMetadata(metadata)))
                        .await
                        .is_err()
                {
                    return Err(ApiError::Stream(
                        "response event consumer dropped".to_string(),
                    ));
                }
                if let Some(buffering) = safety_buffering
                    && tx_event
                        .send(Ok(ResponseEvent::SafetyBuffering(buffering)))
                        .await
                        .is_err()
                {
                    return Err(ApiError::Stream(
                        "response event consumer dropped".to_string(),
                    ));
                }
                match process_responses_event(event) {
                    Ok(Some(event)) => {
                        if let ResponseEvent::Created { response_id: id } = &event {
                            response_id.clone_from(id);
                        }
                        let is_completed = matches!(event, ResponseEvent::Completed { .. });
                        let _ = tx_event.send(Ok(event)).await;
                        if is_completed {
                            break;
                        }
                    }
                    Ok(None) => {}
                    Err(error) => {
                        return Err(error.into_api_error());
                    }
                }
            }
            Message::Binary(_) => {
                return Err(ApiError::Stream("unexpected binary websocket event".into()));
            }
            Message::Close(_) => {
                return Err(ApiError::Stream(
                    "websocket closed by server before response.completed".into(),
                ));
            }
            Message::Frame(_) => {}
            Message::Ping(_) | Message::Pong(_) => {}
        }
    }

    Ok(())
}

fn emit_responses_websocket_timing_event(
    kind: &str,
    payload: &str,
    context: &ResponsesWebsocketTimingLogContext,
) {
    if kind != RESPONSES_WEBSOCKET_TIMING_KIND {
        return;
    }

    // This full payload is excluded from always-on sinks. Opt in with
    // `RUST_LOG='codex_api::responses_websocket_timing=trace'`.
    tracing::event!(
        name: RESPONSES_WEBSOCKET_TIMING_KIND,
        target: RESPONSES_WEBSOCKET_TIMING_EVENT_TARGET,
        tracing::Level::TRACE,
        model = context.model.as_str(),
        session_id = context.session_id.as_deref().unwrap_or_default(),
        thread_id = context.thread_id.as_deref().unwrap_or_default(),
        turn_id = context.turn_id.as_deref().unwrap_or_default(),
        traceparent = context.traceparent.as_deref().unwrap_or_default(),
        previous_response_id = context.previous_response_id.as_deref().unwrap_or_default(),
        request_start_ms = context.request_start_ms.as_deref().unwrap_or_default(),
        warmup = context.warmup,
        connection_reused = context.connection_reused,
        payload,
        "responses websocket timing"
    );
}

fn safety_buffering_for_event(
    event: &ResponsesStreamEvent,
    treatment: &mut SafetyBufferingTreatment,
) -> Option<crate::common::SafetyBuffering> {
    if let Some(headers) = event.headers.as_ref().and_then(Value::as_object)
        && let Some(updated_treatment) =
            treatment_from_headers(&json_headers_to_http_headers(headers))
    {
        *treatment = updated_treatment;
    }
    event.safety_buffering(treatment)
}

async fn send_websocket_request(
    ws_stream: &mut WsStream,
    request_text: String,
    idle_timeout: Duration,
    telemetry: Option<&Arc<dyn WebsocketTelemetry>>,
    connection_reused: bool,
) -> Result<(), ApiError> {
    let request_start = Instant::now();
    let result = tokio::time::timeout(
        idle_timeout,
        ws_stream.send(Message::Text(request_text.into())),
    )
    .await
    .map_err(|_| ApiError::Stream("idle timeout sending websocket request".into()))
    .and_then(|result| result.map_err(map_ws_stream_error));

    if let Some(t) = telemetry.as_ref() {
        t.on_ws_request(
            request_start.elapsed(),
            result.as_ref().err(),
            connection_reused,
        );
    }

    result?;

    Ok(())
}

fn serialize_websocket_request(request: &ResponsesWsRequest<'_>) -> Result<String, ApiError> {
    serde_json::to_string(request)
        .map_err(|err| ApiError::Stream(format!("failed to encode websocket request: {err}")))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::common::ResponseCreateWsRequest;
    use crate::common::ResponsesApiRequest;
    use codex_protocol::ResponseItemId;
    use codex_protocol::models::ContentItem;
    use codex_protocol::models::ResponseItem;
    use http::HeaderValue;
    use pretty_assertions::assert_eq;
    use serde_json::json;
    use serde_json::value::RawValue;
    use serde_json::value::to_raw_value;
    use std::collections::HashMap;
    use std::sync::Arc;

    #[test]
    fn direct_serialization_preserves_websocket_request_payload() {
        let api_request = ResponsesApiRequest {
            model: "gpt-test".to_string(),
            input: vec![ResponseItem::Message {
                status: None,
                encrypted_content: None,
                id: Some(ResponseItemId::with_suffix("msg", "1")),
                role: "user".to_string(),
                content: vec![ContentItem::InputText {
                    text: "hello".to_string(),
                }],
                phase: None,
                internal_chat_message_metadata_passthrough: None,
            }],
            tools: Some(
                Arc::<RawValue>::from(
                    to_raw_value(&vec![json!({
                        "type": "function",
                        "name": "lookup",
                        "parameters": {"type": "object"}
                    })])
                    .expect("serialize tools"),
                )
                .into(),
            ),
            tool_choice: "auto".to_string(),
            parallel_tool_calls: true,
            reasoning: None,
            store: false,
            stream: true,
            stream_options: None,
            include: vec!["reasoning.encrypted_content".to_string()],
            service_tier: Some("priority".to_string()),
            prompt_cache_key: Some("cache-key".to_string()),
            text: None,
            access_programs: Some(
                codex_protocol::turn_input::CyberAccessProgram::DaybreakBlue.into(),
            ),
            client_metadata: Some(HashMap::from([(
                "traceparent".to_string(),
                "00-0123456789abcdef0123456789abcdef-0123456789abcdef-01".to_string(),
            )])),
        };
        let request = ResponsesWsRequest::ResponseCreate(ResponseCreateWsRequest {
            previous_response_id: Some("resp-1".to_string()),
            generate: Some(false),
            ..ResponseCreateWsRequest::from(&api_request)
        });

        let mut expected_payload =
            serde_json::to_value(&api_request).expect("serialize responses API request");
        expected_payload["type"] = json!("response.create");
        expected_payload["previous_response_id"] = json!("resp-1");
        expected_payload["generate"] = json!(false);
        let request_text =
            serialize_websocket_request(&request).expect("serialize websocket request");
        assert!(request_text.starts_with(
            r#"{"type":"response.create","model":"gpt-test","stream":true,"service_tier":"priority","previous_response_id":"resp-1","input":"#
        ));
        let wire_payload =
            serde_json::from_str::<Value>(&request_text).expect("parse websocket request");

        assert_eq!(wire_payload, expected_payload);
    }

    #[test]
    fn websocket_config_enables_permessage_deflate() {
        let config = websocket_config();
        assert!(config.extensions.permessage_deflate.is_some());
    }

    #[test]
    fn parse_wrapped_websocket_error_event_maps_to_transport_http() {
        let payload = json!({
            "type": "error",
            "status": 429,
            "error": {
                "type": "usage_limit_reached",
                "message": "The usage limit has been reached",
                "plan_type": "pro",
                "resets_at": 1738888888
            },
            "headers": {
                "x-codex-primary-used-percent": "100.0",
                "x-codex-primary-window-minutes": 15
            }
        })
        .to_string();

        let wrapped_error = parse_wrapped_websocket_error_event(&payload)
            .expect("expected websocket error payload to be parsed");
        let api_error = map_wrapped_websocket_error_event(wrapped_error, payload)
            .expect("expected websocket error payload to map to ApiError");

        let ApiError::Transport(TransportError::Http {
            status,
            headers,
            body,
            ..
        }) = api_error
        else {
            panic!("expected ApiError::Transport(Http)");
        };

        assert_eq!(status, StatusCode::TOO_MANY_REQUESTS);
        let headers = headers.expect("expected headers");
        assert_eq!(
            headers
                .get("x-codex-primary-used-percent")
                .and_then(|value| value.to_str().ok()),
            Some("100.0")
        );
        assert_eq!(
            headers
                .get("x-codex-primary-window-minutes")
                .and_then(|value| value.to_str().ok()),
            Some("15")
        );
        let body = body.expect("expected body");
        assert!(body.contains("usage_limit_reached"));
        assert!(body.contains("The usage limit has been reached"));
    }

    /// Websocket error advice uses the shared HTTP value validation and duplicate handling.
    #[tokio::test(start_paused = true)]
    async fn wrapped_websocket_retry_after_uses_http_header_validation() {
        for (error_headers, expected_seconds) in [
            (json!("invalid"), 12),
            (json!({"retry-after": "\n5\n"}), 12),
            (json!({"retry-after": "\t5\t"}), 5),
            (json!({"Retry-After": "5", "retry-after": "30"}), 30),
        ] {
            let payload = json!({
                "type": "error",
                "status": 429,
                "error": {"code": "rate_limit_exceeded", "headers": error_headers},
                "headers": {"retry-after": "12"}
            })
            .to_string();
            let wrapped = parse_wrapped_websocket_error_event(&payload).unwrap();
            let Some(ApiError::Transport(TransportError::Http { retry_after, .. })) =
                map_wrapped_websocket_error_event(wrapped, payload)
            else {
                panic!("expected a websocket HTTP error");
            };
            assert_eq!(
                retry_after,
                RetryAfter::from_delay(Duration::from_secs(expected_seconds))
            );
        }
    }

    #[test]
    fn wrapped_websocket_usage_limit_preserves_optional_window() {
        for (window, expected) in [(Some(10080), Some(10080)), (None, None)] {
            let mut payload = json!({
                "type": "error",
                "status": 429,
                "error": {
                    "type": "usage_limit_reached",
                    "plan_type": "pro"
                }
            });
            if let Some(window) = window {
                payload["error"]["limit_window_minutes"] = json!(window);
            }
            let payload = payload.to_string();
            let wrapped = parse_wrapped_websocket_error_event(&payload).expect("websocket error");
            let api_error = map_wrapped_websocket_error_event(wrapped, payload).expect("API error");
            let err = crate::api_bridge::map_api_error(api_error);
            let codex_protocol::error::CodexErrorDetails::UsageLimitReached(usage_limit) =
                err.details()
            else {
                panic!("expected usage-limit error, got {err:?}");
            };
            assert_eq!(usage_limit.limit_window_minutes, expected);
        }
    }

    #[test]
    fn parse_wrapped_websocket_error_event_ignores_non_error_payloads() {
        let payload = json!({
            "type": "response.created",
            "response": {
                "id": "resp-1"
            }
        })
        .to_string();

        let wrapped_error = parse_wrapped_websocket_error_event(&payload);
        assert!(wrapped_error.is_none());
    }

    #[test]
    fn parse_wrapped_websocket_error_event_with_status_maps_invalid_request() {
        let payload = json!({
            "type": "error",
            "status": 400,
            "error": {
                "type": "invalid_request_error",
                "message": "Model does not support image inputs"
            }
        })
        .to_string();

        let wrapped_error = parse_wrapped_websocket_error_event(&payload)
            .expect("expected websocket error payload to be parsed");
        let api_error = map_wrapped_websocket_error_event(wrapped_error, payload)
            .expect("expected websocket error payload to map to ApiError");
        let ApiError::Transport(TransportError::Http { status, body, .. }) = api_error else {
            panic!("expected ApiError::Transport(Http)");
        };
        assert_eq!(status, StatusCode::BAD_REQUEST);
        let body = body.expect("expected body");
        assert!(body.contains("invalid_request_error"));
        assert!(body.contains("Model does not support image inputs"));
    }

    #[test]
    fn parse_wrapped_websocket_error_event_with_connection_limit_maps_retryable() {
        let payload = json!({
            "type": "error",
            "status": 400,
            "error": {
                "type": "invalid_request_error",
                "code": "websocket_connection_limit_reached",
                "message": "Responses websocket connection limit reached (60 minutes). Create a new websocket connection to continue."
            }
        })
        .to_string();

        let wrapped_error = parse_wrapped_websocket_error_event(&payload)
            .expect("expected websocket error payload to be parsed");
        let api_error = map_wrapped_websocket_error_event(wrapped_error, payload)
            .expect("expected websocket error payload to map to ApiError");
        let ApiError::Retryable {
            message,
            retry_after,
        } = api_error
        else {
            panic!("expected ApiError::Retryable");
        };
        assert_eq!(message, WEBSOCKET_CONNECTION_LIMIT_REACHED_MESSAGE);
        assert_eq!(retry_after, None);
    }

    #[test]
    fn parse_wrapped_websocket_error_event_without_status_is_not_mapped() {
        let payload = json!({
            "type": "error",
            "error": {
                "type": "usage_limit_reached",
                "message": "The usage limit has been reached"
            },
            "headers": {
                "x-codex-primary-used-percent": "100.0",
                "x-codex-primary-window-minutes": 15
            }
        })
        .to_string();

        let wrapped_error = parse_wrapped_websocket_error_event(&payload)
            .expect("expected websocket error payload to be parsed");
        let api_error = map_wrapped_websocket_error_event(wrapped_error, payload);
        assert!(api_error.is_none());
    }

    #[test]
    fn merge_request_headers_matches_http_precedence() {
        let mut provider_headers = HeaderMap::new();
        provider_headers.insert(
            "originator",
            HeaderValue::from_static("provider-originator"),
        );
        provider_headers.insert("x-priority", HeaderValue::from_static("provider"));

        let mut extra_headers = HeaderMap::new();
        extra_headers.insert("x-priority", HeaderValue::from_static("extra"));

        let mut default_headers = HeaderMap::new();
        default_headers.insert("originator", HeaderValue::from_static("default-originator"));
        default_headers.insert("x-priority", HeaderValue::from_static("default"));
        default_headers.insert("x-default-only", HeaderValue::from_static("default-only"));

        let merged = merge_request_headers(&provider_headers, extra_headers, default_headers);

        assert_eq!(
            merged.get("originator"),
            Some(&HeaderValue::from_static("provider-originator"))
        );
        assert_eq!(
            merged.get("x-priority"),
            Some(&HeaderValue::from_static("extra"))
        );
        assert_eq!(
            merged.get("x-default-only"),
            Some(&HeaderValue::from_static("default-only"))
        );
    }

    #[test]
    fn websocket_safety_buffering_uses_event_before_header_fallback() {
        let metadata: ResponsesStreamEvent = serde_json::from_value(json!({
            "type": "codex.response.metadata",
            "headers": {
                "x-codex-safety-buffering-enabled": "true",
                "x-codex-safety-buffering-faster-model": "gpt-fast-header"
            }
        }))
        .expect("deserialize treatment metadata");
        let event: ResponsesStreamEvent = serde_json::from_value(json!({
            "type": "response.output_text.delta",
            "safety_buffering": {
                "use_cases": ["cyber"],
                "reasons": ["user_risk"],
                "retry_model": "gpt-fast-wire"
            }
        }))
        .expect("deserialize safety buffering event");
        let mut treatment = SafetyBufferingTreatment::default();

        assert!(safety_buffering_for_event(&metadata, &mut treatment).is_none());
        let buffering = safety_buffering_for_event(&event, &mut treatment)
            .expect("expected safety buffering payload");

        assert_eq!(
            buffering,
            crate::common::SafetyBuffering {
                use_cases: vec!["cyber".to_string()],
                reasons: vec!["user_risk".to_string()],
                show_buffering_ui: true,
                faster_model: Some("gpt-fast-wire".to_string()),
            }
        );
    }

    #[test]
    fn websocket_safety_buffering_event_controls_visibility_when_header_disables_it() {
        let metadata: ResponsesStreamEvent = serde_json::from_value(json!({
            "type": "codex.response.metadata",
            "headers": {
                "x-codex-safety-buffering-enabled": "false",
                "x-codex-safety-buffering-faster-model": "gpt-fast-header"
            }
        }))
        .expect("deserialize treatment metadata");
        let event: ResponsesStreamEvent = serde_json::from_value(json!({
            "type": "response.output_text.delta",
            "safety_buffering": {
                "use_cases": ["cyber"],
                "reasons": ["user_risk"]
            }
        }))
        .expect("deserialize safety buffering event");
        let mut treatment = SafetyBufferingTreatment::default();

        assert!(safety_buffering_for_event(&metadata, &mut treatment).is_none());
        let buffering = safety_buffering_for_event(&event, &mut treatment)
            .expect("expected safety buffering payload");

        assert_eq!(
            buffering,
            crate::common::SafetyBuffering {
                use_cases: vec!["cyber".to_string()],
                reasons: vec!["user_risk".to_string()],
                show_buffering_ui: true,
                faster_model: Some("gpt-fast-header".to_string()),
            }
        );
    }
}