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
This commit is contained in:
thorvald-oai
2026-10-17 23:34:04 +00:00
committed by copyberry
parent ffd3429d75
commit 7b44ac1661
11 changed files with 238 additions and 82 deletions
@@ -48,6 +48,7 @@ pub async fn start_control_socket_acceptor(
transport_event_tx: mpsc::Sender<TransportEvent>,
shutdown_token: CancellationToken,
daemon_shutdown_access: DaemonShutdownAccess,
compression: WebSocketCompression,
) -> IoResult<JoinHandle<()>> {
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;
});
@@ -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");
@@ -88,6 +88,7 @@ fn print_websocket_startup_banner(addr: SocketAddr) {
struct WebSocketListenerState {
transport_event_tx: mpsc::Sender<TransportEvent>,
auth_policy: Arc<WebsocketAuthPolicy>,
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<TransportEvent>,
shutdown_token: CancellationToken,
auth_policy: WebsocketAuthPolicy,
compression: WebSocketCompression,
) -> IoResult<JoinHandle<()>> {
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,
@@ -23,87 +23,110 @@ use tokio_tungstenite::tungstenite::client::IntoClientRequest;
#[tokio::test]
async fn tcp_compression_negotiates_and_retains_history() -> Result<(), Box<dyn Error>> {
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::<SocketAddr>(),
)
.await
});
let server = tokio::spawn(async move {
axum::serve(
listener,
router.into_make_service_with_connect_info::<SocketAddr>(),
)
.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<dyn Error>> {
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(())
}
+11
View File
@@ -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::<JoinHandle<()>>::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);
@@ -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(())
}
@@ -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)]
+6
View File
@@ -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"
},
+8
View File
@@ -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",
+1
View File
@@ -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"
+3
View File
@@ -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",