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",