Advertise the control socket's WebSocket message size limit (#46548)

## Why

Clients need to know the incoming size limit for single-frame requests so they can reject oversized requests before the socket closes.

## What changed

Add `x-codex-websocket-max-unfragmented-message-bytes` to successful control socket WebSocket handshake responses. Derive the advertised byte limit from the minimum configured frame and message size limits, using the same `WebSocketConfig` for the connection.

## Testing

Extend the control socket upgrade test to verify that the header contains a numeric value matching the effective default limit.

GitOrigin-RevId: c4f46d12e00e20f56d4c14a3e84f9fe6e2f020fa
This commit is contained in:
konsti-openai
2026-09-18 21:20:50 +00:00
committed by copyberry
parent c026e7a622
commit 98a8d4ea9c
2 changed files with 47 additions and 3 deletions

View File

@@ -15,10 +15,13 @@ use futures::StreamExt;
use tokio::sync::mpsc;
use tokio::task::JoinHandle;
use tokio::time::Duration;
use tokio_tungstenite::accept_hdr_async;
use tokio_tungstenite::accept_hdr_async_with_config;
use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::tungstenite::handshake::server::Response as HandshakeResponse;
use tokio_tungstenite::tungstenite::http::HeaderValue;
use tokio_tungstenite::tungstenite::http::Response;
use tokio_tungstenite::tungstenite::http::StatusCode;
use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
use tokio_util::sync::CancellationToken;
use tracing::error;
use tracing::info;
@@ -26,6 +29,10 @@ use tracing::warn;
#[cfg(unix)]
const CONTROL_SOCKET_MODE: u32 = 0o600;
// Advertise the effective incoming cap for single-frame messages so clients can
// reject oversized requests before the socket closes.
const MAX_UNFRAGMENTED_MESSAGE_BYTES_HEADER: &str =
"x-codex-websocket-max-unfragmented-message-bytes";
#[derive(Clone, Copy)]
pub enum DaemonShutdownAccess {
@@ -150,9 +157,18 @@ async fn run_control_socket_acceptor(
let transport_event_tx = transport_event_tx.clone();
tokio::spawn(async move {
let mut shutdown_request = false;
let websocket_stream = match accept_hdr_async(
let websocket_config = WebSocketConfig::default();
let max_unfragmented_message_bytes = [
websocket_config.max_frame_size,
websocket_config.max_message_size,
]
.into_iter()
.flatten()
.min();
let websocket_stream = match accept_hdr_async_with_config(
stream,
|request: &tokio_tungstenite::tungstenite::handshake::server::Request, response| {
|request: &tokio_tungstenite::tungstenite::handshake::server::Request,
mut response: HandshakeResponse| {
if request.uri().path() == "/daemon/shutdown" {
if !matches!(daemon_shutdown_access, DaemonShutdownAccess::Managed) {
let mut rejection = Response::new(Some("unmanaged server".to_string()));
@@ -161,8 +177,15 @@ async fn run_control_socket_acceptor(
}
shutdown_request = true;
}
if let Some(max_bytes) = max_unfragmented_message_bytes {
response.headers_mut().insert(
MAX_UNFRAGMENTED_MESSAGE_BYTES_HEADER,
HeaderValue::from(max_bytes),
);
}
Ok(response)
},
Some(websocket_config),
)
.await
{

View File

@@ -21,6 +21,7 @@ use tokio::time::timeout;
use tokio_tungstenite::client_async;
use tokio_tungstenite::tungstenite::Bytes;
use tokio_tungstenite::tungstenite::Message as WebSocketMessage;
use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
use tokio_util::sync::CancellationToken;
#[test]
@@ -77,6 +78,26 @@ async fn control_socket_acceptor_upgrades_and_forwards_websocket_text_messages_a
.await
.expect("websocket upgrade should complete");
assert_eq!(response.status().as_u16(), 101);
let advertised_max = response
.headers()
.get("x-codex-websocket-max-unfragmented-message-bytes")
.expect("byte cap header should be advertised")
.to_str()
.expect("byte cap header should be ASCII")
.parse::<usize>()
.expect("byte cap header should be a number");
let websocket_config = WebSocketConfig::default();
assert_eq!(
advertised_max,
[
websocket_config.max_frame_size,
websocket_config.max_message_size,
]
.into_iter()
.flatten()
.min()
.expect("default websocket config should have an incoming size limit")
);
let opened = timeout(Duration::from_secs(1), transport_event_rx.recv())
.await