Honor advertised WebSocket message size limits (#60841)

## Why

Requests that exceed the server's advertised WebSocket message size limit can fall back to HTTP without first uploading the oversized message.

## What changed

Read `x-codex-websocket-max-message-bytes` from the handshake response and check the serialized request's byte length before sending. Route oversized requests through the existing HTTP fallback. Ignore missing or invalid limits, and remove the close-code-specific wording from the size-limit error.

## Testing

Add a fallback test that advertises a one-byte limit and verifies zero WebSocket generation requests and one HTTP request.

GitOrigin-RevId: 03b42012ca77d6b046e937fef8e6335d855a1c10

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<usize>,
     routing_affinity: Option<codex_http_client::ChatGptRoutingAffinity>,
     tx_command: mpsc::Sender<WsCommand>,
     rx_message: mpsc::UnboundedReceiver<Result<Message, WsError>>,
@@ -105,7 +106,7 @@ async fn send_message(
 }
 
 impl WsStream {
-    fn new(inner: WebSocketConnection) -> Self {
+    fn new(inner: WebSocketConnection, max_message_bytes: Option<usize>) -> Self {
         let routing_affinity = inner.routing_affinity().cloned();
         let (tx_command, mut rx_command) = mpsc::channel::<WsCommand>(32);
         let (tx_message, rx_message) = mpsc::unbounded_channel::<Result<Message, WsError>>();
@@ -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::<usize>().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<dyn WebsocketTelemetry>>,
     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,

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,

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;

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(