diff --git a/codex-rs/codex-api/src/endpoint/responses_websocket.rs b/codex-rs/codex-api/src/endpoint/responses_websocket.rs index 8a4c6275a0..bced460a6b 100644 --- a/codex-rs/codex-api/src/endpoint/responses_websocket.rs +++ b/codex-rs/codex-api/src/endpoint/responses_websocket.rs @@ -61,6 +61,7 @@ use tungstenite::protocol::WebSocketConfig; use url::Url; struct WsStream { + max_message_bytes: Option, routing_affinity: Option, tx_command: mpsc::Sender, rx_message: mpsc::UnboundedReceiver>, @@ -105,7 +106,7 @@ async fn send_message( } impl WsStream { - fn new(inner: WebSocketConnection) -> Self { + fn new(inner: WebSocketConnection, max_message_bytes: Option) -> Self { let routing_affinity = inner.routing_affinity().cloned(); let (tx_command, mut rx_command) = mpsc::channel::(32); let (tx_message, rx_message) = mpsc::unbounded_channel::>(); @@ -164,6 +165,7 @@ impl WsStream { }); Self { + max_message_bytes, routing_affinity, tx_command, rx_message, @@ -204,6 +206,7 @@ impl Drop for WsStream { } const X_CODEX_TURN_STATE_HEADER: &str = "x-codex-turn-state"; +const X_CODEX_WEBSOCKET_MAX_MESSAGE_BYTES_HEADER: &str = "x-codex-websocket-max-message-bytes"; 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"; @@ -649,6 +652,11 @@ async fn connect_websocket( } }; + let max_message_bytes = response + .headers() + .get(X_CODEX_WEBSOCKET_MAX_MESSAGE_BYTES_HEADER) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.parse::().ok()); let reasoning_included = response.headers().contains_key(X_REASONING_INCLUDED_HEADER); let server_model = response .headers() @@ -664,7 +672,7 @@ async fn connect_websocket( let _ = turn_state.set(header_value.to_string()); } Ok(( - WsStream::new(stream), + WsStream::new(stream, max_message_bytes), response, reasoning_included, server_model, @@ -1101,6 +1109,12 @@ async fn send_websocket_request( telemetry: Option<&Arc>, connection_reused: bool, ) -> Result<(), ApiError> { + if ws_stream + .max_message_bytes + .is_some_and(|limit| request_text.len() > limit) + { + return Err(ApiError::WebsocketMessageTooLarge); + } let request_start = Instant::now(); let result = tokio::time::timeout( idle_timeout, diff --git a/codex-rs/codex-api/src/error.rs b/codex-rs/codex-api/src/error.rs index 10f818eeb5..a853c7be33 100644 --- a/codex-rs/codex-api/src/error.rs +++ b/codex-rs/codex-api/src/error.rs @@ -15,7 +15,7 @@ pub enum ApiError { Api { status: StatusCode, message: String }, #[error("stream error: {0}")] Stream(String), - #[error("websocket request exceeds the peer's message size limit (close code 1009)")] + #[error("websocket request exceeds the peer's message size limit")] WebsocketMessageTooLarge, #[error("stream error: Incomplete response returned, reason: content_filter")] ContentFilter, diff --git a/codex-rs/core/tests/suite/websocket_fallback_close.rs b/codex-rs/core/tests/suite/websocket_fallback_close.rs index f0c2c9eb39..f963162da7 100644 --- a/codex-rs/core/tests/suite/websocket_fallback_close.rs +++ b/codex-rs/core/tests/suite/websocket_fallback_close.rs @@ -19,12 +19,15 @@ use tokio::net::TcpListener; use tokio::net::TcpStream; use tokio::task::JoinSet; use tokio_tungstenite::tungstenite::Message; +use tokio_tungstenite::tungstenite::handshake::server::Request; +use tokio_tungstenite::tungstenite::handshake::server::Response; use tokio_tungstenite::tungstenite::protocol::CloseFrame; use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode; #[derive(Clone, Copy, PartialEq, Eq)] enum Rejection { Oversized, + AdvertisedLimit, OversizedWarmup, AfterModelMetadata, AfterRateLimits, @@ -32,6 +35,7 @@ enum Rejection { GoingAway, } +#[test_case::test_case(Rejection::AdvertisedLimit, 0, 1; "advertised_limit_skips_websocket_upload")] #[test_case::test_case(Rejection::Oversized, 1, 1; "size_rejection_before_any_text")] #[test_case::test_case(Rejection::OversizedWarmup, 0, 1; "size_rejection_during_prewarm")] #[test_case::test_case(Rejection::AfterModelMetadata, 2, 0; "size_close_after_model_metadata")] @@ -98,7 +102,17 @@ async fn serve_connection( tokio::io::copy_bidirectional(&mut stream, &mut upstream).await?; return Ok(()); } - let mut socket = tokio_tungstenite::accept_async(stream).await?; + let mut socket = + tokio_tungstenite::accept_hdr_async(stream, |_: &Request, mut res: Response| { + if rejection == Rejection::AdvertisedLimit { + res.headers_mut().insert( + "x-codex-websocket-max-message-bytes", + http::HeaderValue::from_static("1"), + ); + } + Ok(res) + }) + .await?; while let Some(message) = socket.next().await.transpose()? { let Message::Text(text) = message else { continue; diff --git a/codex-rs/protocol/src/error.rs b/codex-rs/protocol/src/error.rs index e769caa5bb..9884549607 100644 --- a/codex-rs/protocol/src/error.rs +++ b/codex-rs/protocol/src/error.rs @@ -121,8 +121,8 @@ pub enum CodexErrorDetails { /// The Session loop treats this as a transient error and will automatically retry the turn. #[error("stream disconnected before completion: {0}")] Stream(String), - /// Retry the rejected request over HTTP instead of sending it on another socket. - #[error("websocket request exceeds the peer's message size limit (close code 1009)")] + /// The peer's WebSocket message size limit requires sending this request over HTTP. + #[error("websocket request exceeds the peer's message size limit")] WebsocketMessageTooLarge, /// A response stopped by the content filter. Sampling retries need developer guidance. #[error(