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
This commit is contained in:
Charlie Marsh
2026-11-11 18:58:56 +00:00
committed by copyberry
parent bc3ccba6e1
commit ed7ca22bb6
4 changed files with 34 additions and 6 deletions
@@ -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;