From 7b44ac1661433ae3a6f375fde873fa7123c935fe Mon Sep 17 00:00:00 2001 From: thorvald-oai Date: Sat, 17 Oct 2026 23:28:00 +0000 Subject: [PATCH] Gate app-server WebSocket compression behind a feature flag (#54949) ## What changed Add `features.app_server_websocket_compression`, disabled by default. TCP WebSocket and Unix-domain socket listeners now negotiate zstd compression only when the feature is enabled and the client offers the supported protocol. Otherwise, connections use plain text messages. Include the feature in the configuration schema and allow its override during TUI daemon startup. ## Testing Extend transport tests to cover enabled and disabled compression with supported, unsupported, and absent client offers on both transports. Add an app-server integration test that verifies startup configuration controls negotiation and the message format used for initialization responses. GitOrigin-RevId: 6001d65f66f9c48f099b85380c2cad8ffee04de7 --- .../src/transport/unix_socket.rs | 18 +- .../src/transport/unix_socket_tests.rs | 10 ++ .../src/transport/websocket.rs | 10 +- .../transport/websocket_compression_tests.rs | 167 ++++++++++-------- codex-rs/app-server/src/lib.rs | 11 ++ ...on_handling_websocket_compression_tests.rs | 84 +++++++++ codex-rs/app-server/tests/suite/v2/mod.rs | 2 + codex-rs/core/config.schema.json | 6 + codex-rs/features/src/lib.rs | 8 + codex-rs/tui/src/daemon_startup.rs | 1 + codex-rs/tui/src/daemon_startup_tests.rs | 3 + 11 files changed, 238 insertions(+), 82 deletions(-) create mode 100644 codex-rs/app-server/tests/suite/v2/connection_handling_websocket_compression_tests.rs diff --git a/codex-rs/app-server-transport/src/transport/unix_socket.rs b/codex-rs/app-server-transport/src/transport/unix_socket.rs index 70af70c07a..10fc906796 100644 --- a/codex-rs/app-server-transport/src/transport/unix_socket.rs +++ b/codex-rs/app-server-transport/src/transport/unix_socket.rs @@ -48,6 +48,7 @@ pub async fn start_control_socket_acceptor( transport_event_tx: mpsc::Sender, shutdown_token: CancellationToken, daemon_shutdown_access: DaemonShutdownAccess, + compression: WebSocketCompression, ) -> IoResult> { let (listener, socket_guard) = bind_private_daemon_socket(socket_path.as_path()).await?; info!( @@ -61,6 +62,7 @@ pub async fn start_control_socket_acceptor( shutdown_token, socket_guard, daemon_shutdown_access, + compression, ))) } @@ -149,6 +151,7 @@ async fn run_control_socket_acceptor( shutdown_token: CancellationToken, socket_guard: DaemonSocketFileGuard, daemon_shutdown_access: DaemonShutdownAccess, + compression: WebSocketCompression, ) { let _socket_guard = socket_guard; loop { @@ -178,7 +181,7 @@ async fn run_control_socket_acceptor( let transport_event_tx = transport_event_tx.clone(); tokio::spawn(async move { let mut shutdown_request = false; - let mut compression = WebSocketCompression::Plain; + let mut negotiated_compression = WebSocketCompression::Plain; let websocket_config = WebSocketConfig::default(); let max_unfragmented_message_bytes = [ websocket_config.max_frame_size, @@ -198,12 +201,13 @@ async fn run_control_socket_acceptor( return Err(rejection); } shutdown_request = true; - } else if request - .headers() - .get(WEBSOCKET_COMPRESSION_HEADER) - .is_some_and(|value| value == ZSTD_STREAM_PROTOCOL) + } else if compression == WebSocketCompression::Zstd + && request + .headers() + .get(WEBSOCKET_COMPRESSION_HEADER) + .is_some_and(|value| value == ZSTD_STREAM_PROTOCOL) { - compression = WebSocketCompression::Zstd; + negotiated_compression = WebSocketCompression::Zstd; response.headers_mut().insert( WEBSOCKET_COMPRESSION_HEADER, HeaderValue::from_static(ZSTD_STREAM_PROTOCOL), @@ -236,7 +240,7 @@ async fn run_control_socket_acceptor( websocket_writer, websocket_reader, transport_event_tx, - compression, + negotiated_compression, ) .await; }); diff --git a/codex-rs/app-server-transport/src/transport/unix_socket_tests.rs b/codex-rs/app-server-transport/src/transport/unix_socket_tests.rs index 43dbee22f4..9bfa950a81 100644 --- a/codex-rs/app-server-transport/src/transport/unix_socket_tests.rs +++ b/codex-rs/app-server-transport/src/transport/unix_socket_tests.rs @@ -9,6 +9,7 @@ use super::start_control_socket_acceptor; use super::unix_socket::acquire_removable_app_server_startup_lock; #[cfg(unix)] use super::unix_socket::try_acquire_removable_app_server_startup_lock; +use crate::websocket_compression::WebSocketCompression; use codex_app_server_protocol::JSONRPCMessage; use codex_app_server_protocol::JSONRPCNotification; use codex_core::config::find_codex_home; @@ -79,6 +80,7 @@ async fn long_control_socket_paths_connect_to_distinct_daemons() { transport_event_tx, shutdown_token.clone(), DaemonShutdownAccess::Disabled, + WebSocketCompression::Plain, ) .await .expect("control socket acceptor should start"); @@ -113,6 +115,7 @@ async fn control_socket_acceptor_upgrades_and_forwards_websocket_text_messages_a transport_event_tx, shutdown_token.clone(), DaemonShutdownAccess::Disabled, + WebSocketCompression::Plain, ) .await .expect("control socket acceptor should start"); @@ -221,6 +224,7 @@ async fn shutdown_is_only_accepted_on_managed_local_socket_for_its_own_pid() { tx, shutdown.clone(), DaemonShutdownAccess::Disabled, + WebSocketCompression::Plain, ) .await .expect("acceptor"); @@ -244,6 +248,7 @@ async fn shutdown_is_only_accepted_on_managed_local_socket_for_its_own_pid() { tx, shutdown.clone(), DaemonShutdownAccess::Managed, + WebSocketCompression::Plain, ) .await .expect("managed acceptor"); @@ -365,6 +370,7 @@ async fn control_socket_rejects_writable_parent_without_changing_permissions() { tx, CancellationToken::new(), DaemonShutdownAccess::Disabled, + WebSocketCompression::Plain, ) .await .unwrap_err(); @@ -398,6 +404,7 @@ async fn control_socket_file_is_private_after_bind() { transport_event_tx, shutdown_token.clone(), DaemonShutdownAccess::Disabled, + WebSocketCompression::Plain, ) .await .expect("control socket acceptor should start"); @@ -450,12 +457,14 @@ async fn control_socket_file_is_private_after_bind() { sender.clone(), shutdown.clone(), DaemonShutdownAccess::Disabled, + WebSocketCompression::Plain, ), start_control_socket_acceptor( socket_path.clone(), sender, shutdown.clone(), DaemonShutdownAccess::Disabled, + WebSocketCompression::Plain, ), ); let (acceptor, error) = match (first, second) { @@ -491,6 +500,7 @@ async fn control_socket_pins_directory_until_shutdown() { tx, shutdown.clone(), DaemonShutdownAccess::Disabled, + WebSocketCompression::Plain, ) .await .expect("acceptor"); diff --git a/codex-rs/app-server-transport/src/transport/websocket.rs b/codex-rs/app-server-transport/src/transport/websocket.rs index b4f3cc97d6..f02aed975a 100644 --- a/codex-rs/app-server-transport/src/transport/websocket.rs +++ b/codex-rs/app-server-transport/src/transport/websocket.rs @@ -88,6 +88,7 @@ fn print_websocket_startup_banner(addr: SocketAddr) { struct WebSocketListenerState { transport_event_tx: mpsc::Sender, auth_policy: Arc, + compression: WebSocketCompression, } async fn health_check_handler() -> StatusCode { @@ -125,9 +126,10 @@ async fn websocket_upgrade_handler( return (err.status_code(), err.message()).into_response(); } info!(%peer_addr, "websocket client connected"); - let compression = if headers - .get(WEBSOCKET_COMPRESSION_HEADER) - .is_some_and(|value| value == ZSTD_STREAM_PROTOCOL) + let compression = if state.compression == WebSocketCompression::Zstd + && headers + .get(WEBSOCKET_COMPRESSION_HEADER) + .is_some_and(|value| value == ZSTD_STREAM_PROTOCOL) { WebSocketCompression::Zstd } else { @@ -159,6 +161,7 @@ pub async fn start_websocket_acceptor( transport_event_tx: mpsc::Sender, shutdown_token: CancellationToken, auth_policy: WebsocketAuthPolicy, + compression: WebSocketCompression, ) -> IoResult> { if is_unauthenticated_non_loopback_listener(bind_address, &auth_policy) { return Err(std::io::Error::new( @@ -181,6 +184,7 @@ pub async fn start_websocket_acceptor( .with_state(WebSocketListenerState { transport_event_tx, auth_policy: Arc::new(auth_policy), + compression, }); let server = axum::serve( listener, diff --git a/codex-rs/app-server-transport/src/transport/websocket_compression_tests.rs b/codex-rs/app-server-transport/src/transport/websocket_compression_tests.rs index f2897cdfcc..6885cb7281 100644 --- a/codex-rs/app-server-transport/src/transport/websocket_compression_tests.rs +++ b/codex-rs/app-server-transport/src/transport/websocket_compression_tests.rs @@ -23,87 +23,110 @@ use tokio_tungstenite::tungstenite::client::IntoClientRequest; #[tokio::test] async fn tcp_compression_negotiates_and_retains_history() -> Result<(), Box> { - let listener = TcpListener::bind("127.0.0.1:0").await?; - let address = listener.local_addr()?; - let (transport_event_tx, mut events) = mpsc::channel(CHANNEL_CAPACITY); - let router = Router::new() - .fallback(any(websocket_upgrade_handler)) - .with_state(WebSocketListenerState { - transport_event_tx, - auth_policy: Arc::new(WebsocketAuthPolicy::default()), + for compression in [WebSocketCompression::Plain, WebSocketCompression::Zstd] { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let address = listener.local_addr()?; + let (transport_event_tx, mut events) = mpsc::channel(CHANNEL_CAPACITY); + let router = Router::new() + .fallback(any(websocket_upgrade_handler)) + .with_state(WebSocketListenerState { + transport_event_tx, + auth_policy: Arc::new(WebsocketAuthPolicy::default()), + compression, + }); + let server = tokio::spawn(async move { + axum::serve( + listener, + router.into_make_service_with_connect_info::(), + ) + .await }); - let server = tokio::spawn(async move { - axum::serve( - listener, - router.into_make_service_with_connect_info::(), - ) - .await - }); - for offer in [None, Some(ZSTD_STREAM_PROTOCOL), Some("future-protocol")] { - let mut request = format!("ws://{address}").into_client_request()?; - if let Some(offer) = offer { - request.headers_mut().insert( - WEBSOCKET_COMPRESSION_HEADER, - HeaderValue::from_static(offer), + for offer in [None, Some(ZSTD_STREAM_PROTOCOL), Some("future-protocol")] { + let mut request = format!("ws://{address}").into_client_request()?; + if let Some(offer) = offer { + request.headers_mut().insert( + WEBSOCKET_COMPRESSION_HEADER, + HeaderValue::from_static(offer), + ); + } + let (websocket, response) = connect_async(request).await?; + let expected = offer.filter(|offer| { + compression == WebSocketCompression::Zstd && *offer == ZSTD_STREAM_PROTOCOL + }); + assert_eq!( + response + .headers() + .get(WEBSOCKET_COMPRESSION_HEADER) + .map(HeaderValue::to_str) + .transpose()?, + expected ); + timeout( + Duration::from_secs(/*secs*/ 5), + roundtrip( + websocket, + &mut events, + WebSocketCompression::from_response_header(expected)?, + ), + ) + .await??; } - let (websocket, response) = connect_async(request).await?; - let expected = offer.filter(|offer| *offer == ZSTD_STREAM_PROTOCOL); - assert_eq!( - response - .headers() - .get(WEBSOCKET_COMPRESSION_HEADER) - .map(HeaderValue::to_str) - .transpose()?, - expected - ); - timeout( - Duration::from_secs(/*secs*/ 5), - roundtrip( - websocket, - &mut events, - WebSocketCompression::from_response_header(expected)?, - ), - ) - .await??; + server.abort(); } - server.abort(); Ok(()) } #[tokio::test] async fn local_socket_compression_negotiates_and_retains_history() -> Result<(), Box> { - let temp_dir = tempfile::TempDir::new()?; - let socket_path = AbsolutePathBuf::from_absolute_path( - temp_dir.path().join("app-server-control").join("rpc.sock"), - )?; - let (transport_event_tx, mut events) = mpsc::channel(CHANNEL_CAPACITY); - let shutdown = CancellationToken::new(); - let server = start_control_socket_acceptor( - socket_path.clone(), - transport_event_tx, - shutdown.clone(), - DaemonShutdownAccess::Disabled, - ) - .await?; - let stream = UnixStream::connect(socket_path.as_path()).await?; - let mut request = "ws://localhost/rpc".into_client_request()?; - request.headers_mut().insert( - WEBSOCKET_COMPRESSION_HEADER, - HeaderValue::from_static(ZSTD_STREAM_PROTOCOL), - ); - let (websocket, response) = client_async(request, stream).await?; - assert_eq!( - response.headers().get(WEBSOCKET_COMPRESSION_HEADER), - Some(&HeaderValue::from_static(ZSTD_STREAM_PROTOCOL)) - ); - timeout( - Duration::from_secs(/*secs*/ 5), - roundtrip(websocket, &mut events, WebSocketCompression::Zstd), - ) - .await??; - shutdown.cancel(); - server.await?; + for compression in [WebSocketCompression::Plain, WebSocketCompression::Zstd] { + let temp_dir = tempfile::TempDir::new()?; + let socket_path = AbsolutePathBuf::from_absolute_path( + temp_dir.path().join("app-server-control").join("rpc.sock"), + )?; + let (transport_event_tx, mut events) = mpsc::channel(CHANNEL_CAPACITY); + let shutdown = CancellationToken::new(); + let server = start_control_socket_acceptor( + socket_path.clone(), + transport_event_tx, + shutdown.clone(), + DaemonShutdownAccess::Disabled, + compression, + ) + .await?; + for offer in [None, Some(ZSTD_STREAM_PROTOCOL), Some("future-protocol")] { + let stream = UnixStream::connect(socket_path.as_path()).await?; + let mut request = "ws://localhost/rpc".into_client_request()?; + if let Some(offer) = offer { + request.headers_mut().insert( + WEBSOCKET_COMPRESSION_HEADER, + HeaderValue::from_static(offer), + ); + } + let (websocket, response) = client_async(request, stream).await?; + let expected = offer.filter(|offer| { + compression == WebSocketCompression::Zstd && *offer == ZSTD_STREAM_PROTOCOL + }); + assert_eq!( + response + .headers() + .get(WEBSOCKET_COMPRESSION_HEADER) + .map(HeaderValue::to_str) + .transpose()?, + expected + ); + timeout( + Duration::from_secs(/*secs*/ 5), + roundtrip( + websocket, + &mut events, + WebSocketCompression::from_response_header(expected)?, + ), + ) + .await??; + } + shutdown.cancel(); + server.await?; + } Ok(()) } diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index d34ea297bd..5dbce95300 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -53,6 +53,7 @@ use codex_app_server_protocol::JSONRPCMessage; use codex_app_server_protocol::ServerNotification; use codex_app_server_protocol::TextPosition as AppTextPosition; use codex_app_server_protocol::TextRange as AppTextRange; +use codex_app_server_transport::WebSocketCompression; use codex_app_server_transport::daemon_recovery_file_path; use codex_config::ConfigLayerSource; use codex_config::ConfigLoadError; @@ -792,6 +793,14 @@ pub async fn run_main_with_transport_options( let remote_control_shutdown_token = transport_shutdown_token.child_token(); let mut transport_accept_handles = Vec::>::new(); + let websocket_compression = if config + .features + .enabled(Feature::AppServerWebSocketCompression) + { + WebSocketCompression::Zstd + } else { + WebSocketCompression::Plain + }; let single_client_mode = matches!(&transport, AppServerTransport::Stdio); let graceful_signal_restart_enabled = runtime_options.install_shutdown_signal_handler && !single_client_mode; @@ -822,6 +831,7 @@ pub async fn run_main_with_transport_options( } else { DaemonShutdownAccess::Disabled }, + websocket_compression, ) .await?; transport_accept_handles.push(accept_handle); @@ -832,6 +842,7 @@ pub async fn run_main_with_transport_options( transport_event_tx.clone(), transport_shutdown_token.clone(), policy_from_settings(&auth)?, + websocket_compression, ) .await?; transport_accept_handles.push(accept_handle); diff --git a/codex-rs/app-server/tests/suite/v2/connection_handling_websocket_compression_tests.rs b/codex-rs/app-server/tests/suite/v2/connection_handling_websocket_compression_tests.rs new file mode 100644 index 0000000000..3d1d56feb6 --- /dev/null +++ b/codex-rs/app-server/tests/suite/v2/connection_handling_websocket_compression_tests.rs @@ -0,0 +1,84 @@ +//! Verify that startup configuration controls compression on the public RPC listener. + +use super::connection_handling_websocket::DEFAULT_READ_TIMEOUT; +use super::connection_handling_websocket::create_config_toml; +use super::connection_handling_websocket::spawn_websocket_server_with_args; +use anyhow::Context; +use anyhow::Result; +use anyhow::bail; +use app_test_support::create_mock_responses_server_sequence_unchecked; +use codex_app_server_protocol::JSONRPCMessage; +use codex_app_server_protocol::RequestId; +use codex_app_server_transport::WEBSOCKET_COMPRESSION_HEADER; +use codex_app_server_transport::ZSTD_STREAM_PROTOCOL; +use codex_app_server_transport::ZstdDecoder; +use codex_app_server_transport::ZstdEncoder; +use futures::SinkExt; +use futures::StreamExt; +use pretty_assertions::assert_eq; +use serde_json::json; +use tempfile::TempDir; +use tokio::time::timeout; +use tokio_tungstenite::connect_async; +use tokio_tungstenite::tungstenite::Message; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; +use tokio_tungstenite::tungstenite::http::HeaderValue; + +#[tokio::test] +async fn websocket_compression_requires_server_feature() -> Result<()> { + for enabled in [false, true] { + let server = create_mock_responses_server_sequence_unchecked(Vec::new()).await; + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), &server.uri(), "never")?; + let args = if enabled { + vec![ + "-c".to_string(), + "features.app_server_websocket_compression=true".to_string(), + ] + } else { + Vec::new() + }; + let (mut process, address) = + spawn_websocket_server_with_args(codex_home.path(), "ws://127.0.0.1:0", &args).await?; + let mut request = format!("ws://{address}").into_client_request()?; + request.headers_mut().insert( + WEBSOCKET_COMPRESSION_HEADER, + HeaderValue::from_static(ZSTD_STREAM_PROTOCOL), + ); + let (mut websocket, response) = connect_async(request).await?; + assert_eq!( + response.headers().get(WEBSOCKET_COMPRESSION_HEADER), + enabled + .then(|| HeaderValue::from_static(ZSTD_STREAM_PROTOCOL)) + .as_ref() + ); + let initialize = json!({ + "id": 1, + "method": "initialize", + "params": {"clientInfo": {"name": "compression-test", "version": "1"}} + }) + .to_string(); + let message = if enabled { + Message::Binary(ZstdEncoder::new()?.encode(&initialize)?.into()) + } else { + Message::Text(initialize.into()) + }; + websocket.send(message).await?; + let response = timeout(DEFAULT_READ_TIMEOUT, websocket.next()) + .await? + .context("server closed before initialize response")??; + let text = match (enabled, response) { + (true, Message::Binary(bytes)) => { + ZstdDecoder::new(/*max_message_size*/ 1024 * 1024)?.decode(&bytes)? + } + (false, Message::Text(text)) => text.to_string(), + _ => bail!("initialize response used unexpected WebSocket message type"), + }; + let JSONRPCMessage::Response(response) = serde_json::from_str(&text)? else { + bail!("expected initialize response"); + }; + assert_eq!(response.id, RequestId::Integer(1)); + process.kill().await?; + } + Ok(()) +} diff --git a/codex-rs/app-server/tests/suite/v2/mod.rs b/codex-rs/app-server/tests/suite/v2/mod.rs index 68167089eb..9d3b0a68a6 100644 --- a/codex-rs/app-server/tests/suite/v2/mod.rs +++ b/codex-rs/app-server/tests/suite/v2/mod.rs @@ -37,6 +37,8 @@ mod config_rpc; #[path = "connection_handling_stdio_tests.rs"] mod connection_handling_stdio; mod connection_handling_websocket; +#[path = "connection_handling_websocket_compression_tests.rs"] +mod connection_handling_websocket_compression; #[cfg(unix)] mod connection_handling_websocket_unix; #[cfg(unix)] diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 80c0539b5b..70609d4737 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -753,6 +753,9 @@ "api_key_model_discovery": { "type": "boolean" }, + "app_server_websocket_compression": { + "type": "boolean" + }, "apply_patch_freeform": { "type": "boolean" }, @@ -7142,6 +7145,9 @@ "api_key_model_discovery": { "type": "boolean" }, + "app_server_websocket_compression": { + "type": "boolean" + }, "apply_patch_freeform": { "type": "boolean" }, diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index df12fcddf4..48ddca715d 100644 --- a/codex-rs/features/src/lib.rs +++ b/codex-rs/features/src/lib.rs @@ -117,6 +117,8 @@ pub enum Feature { /// Automatically start the shared local daemon for eligible interactive launches. DaemonAutoStart, + /// Enable negotiated zstd compression for app-server WebSocket and UDS connections. + AppServerWebSocketCompression, // Experimental /// Advertise enabled environment-backed tools and their shell parameters before the @@ -1037,6 +1039,12 @@ pub const FEATURES: &[FeatureSpec] = &[ stage: Stage::Stable, default_enabled: true, }, + FeatureSpec { + id: Feature::AppServerWebSocketCompression, + key: "app_server_websocket_compression", + stage: Stage::UnderDevelopment, + default_enabled: false, + }, FeatureSpec { id: Feature::TranscriptV2, key: "transcript_v2", diff --git a/codex-rs/tui/src/daemon_startup.rs b/codex-rs/tui/src/daemon_startup.rs index 9e25c09904..f766370155 100644 --- a/codex-rs/tui/src/daemon_startup.rs +++ b/codex-rs/tui/src/daemon_startup.rs @@ -133,6 +133,7 @@ fn allowed_feature(name: &str) -> bool { name, // Client gates and per-thread settings already forwarded in thread requests. "daemon_auto_start" | "transcript_v2" | "standalone_web_search" + | "app_server_websocket_compression" // Shared services and threadless MCP operations need daemon compatibility checks. | "api_key_model_discovery" | "code_mode_host" | "auth_elicitation" | "mcp_oauth_refresh_coordination" diff --git a/codex-rs/tui/src/daemon_startup_tests.rs b/codex-rs/tui/src/daemon_startup_tests.rs index 6dfea7cdb6..c8a37f4f28 100644 --- a/codex-rs/tui/src/daemon_startup_tests.rs +++ b/codex-rs/tui/src/daemon_startup_tests.rs @@ -10,6 +10,8 @@ fn audited_overrides_allow_daemon_without_allowing_arbitrary_config() { for (raw, eligible) in [ ("features.transcript_v2=true", true), ("features.transcript_v2=false", true), + ("features.app_server_websocket_compression=true", true), + ("features.app_server_websocket_compression=false", true), ("features={transcript_v2=true}", true), ("features.transcript_v2='true'", false), ("features.worktrees=true", true), @@ -60,6 +62,7 @@ fn monorepo_wrapper_overrides_are_eligible_and_select_only_server_features() { let overrides = codex_utils_cli::CliConfigOverrides { raw_overrides: [ "features.realtime_conversation=true", + "features.app_server_websocket_compression=true", "features.worktrees=true", "features.remote_models=true", "features.api_key_model_discovery=true",