From dd5521406a801212f7742432bcacb06e91823b01 Mon Sep 17 00:00:00 2001 From: Chris Ayoub Date: Sun, 25 Oct 2026 00:04:54 +0000 Subject: [PATCH] Enable zstd compression for code-mode gRPC responses (#57212) ## What changed - Configure the code-mode host to send zstd-compressed responses when clients advertise support, and enable both client transports to accept them. - Preserve uncompressed responses for clients that do not advertise zstd support. - Patch `tonic` in Bazel builds to use zstd compression level 1, favoring response latency over compression ratio. ## Testing Extend the TCP session integration test to check response encoding and successful session opening for both zstd-capable and uncompressed-only clients. GitOrigin-RevId: b6cf0da51dd30fc2ac9a1b0af11bec26555c559d --- MODULE.bazel | 8 +++ codex-rs/Cargo.lock | 1 + codex-rs/code-mode-host/Cargo.toml | 2 +- codex-rs/code-mode-host/src/grpc_transport.rs | 1 + codex-rs/code-mode-host/tests/grpc_tcp.rs | 52 ++++++++++++------- codex-rs/code-mode/Cargo.toml | 2 +- .../code-mode/src/grpc_session/transport.rs | 2 + patches/BUILD.bazel | 1 + patches/tonic-zstd-level1.patch | 14 +++++ 9 files changed, 62 insertions(+), 21 deletions(-) create mode 100644 patches/tonic-zstd-level1.patch diff --git a/MODULE.bazel b/MODULE.bazel index 3ffde1809c..3157588c37 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -393,6 +393,14 @@ crate.annotation( }, rustc_flags = ["-Copt-level=3"], ) + +# Match the CodeModeHost response-compression benchmark's zstd level. +crate.annotation( + crate = "tonic", + patch_args = ["-p1"], + patches = ["//patches:tonic-zstd-level1.patch"], + version = "0.14.3", +) crate.annotation( crate = "zstd-sys", gen_build_script = "on", diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index dfa360b42a..c28b3752ea 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -15931,6 +15931,7 @@ dependencies = [ "tower-layer", "tower-service", "tracing", + "zstd", ] [[package]] diff --git a/codex-rs/code-mode-host/Cargo.toml b/codex-rs/code-mode-host/Cargo.toml index d56661f7d0..4b0d51c1a9 100644 --- a/codex-rs/code-mode-host/Cargo.toml +++ b/codex-rs/code-mode-host/Cargo.toml @@ -32,7 +32,7 @@ serde_json = { workspace = true } tokio = { workspace = true, features = ["io-std", "io-util", "macros", "net", "process", "rt", "sync", "time"] } tokio-stream = { workspace = true } tokio-util = { workspace = true, features = ["rt"] } -tonic = { workspace = true, features = ["router", "transport"] } +tonic = { workspace = true, features = ["router", "transport", "zstd"] } tracing = { workspace = true } tracing-subscriber = { workspace = true } uuid = { workspace = true, features = ["v4"] } diff --git a/codex-rs/code-mode-host/src/grpc_transport.rs b/codex-rs/code-mode-host/src/grpc_transport.rs index 4cb59985fa..27e49456f0 100644 --- a/codex-rs/code-mode-host/src/grpc_transport.rs +++ b/codex-rs/code-mode-host/src/grpc_transport.rs @@ -34,6 +34,7 @@ pub(super) async fn run_tcp_listener(bind_address: SocketAddr) -> Result<()> { let routes = Routes::new( CodeModeHostServer::new(GrpcCodeModeHost::new()) + .send_compressed(tonic::codec::CompressionEncoding::Zstd) .max_decoding_message_size(MAX_FRAME_BYTES) .max_encoding_message_size(MAX_FRAME_BYTES), ) diff --git a/codex-rs/code-mode-host/tests/grpc_tcp.rs b/codex-rs/code-mode-host/tests/grpc_tcp.rs index b15699b944..a995119548 100644 --- a/codex-rs/code-mode-host/tests/grpc_tcp.rs +++ b/codex-rs/code-mode-host/tests/grpc_tcp.rs @@ -117,6 +117,7 @@ async fn http1_response(endpoint: &str, request: &[u8]) -> Result { .context("timed out waiting for HTTP/1.1 response from gRPC code-mode host")? } +/// Negotiate zstd for capable clients while preserving identity-only clients. #[tokio::test] async fn tcp_listener_opens_a_grpc_session() -> Result<()> { let mut host = Command::new(codex_utils_cargo_bin::cargo_bin("codex-code-mode-host")?) @@ -139,24 +140,37 @@ async fn tcp_listener_opens_a_grpc_session() -> Result<()> { .context("gRPC code-mode host published an invalid endpoint")? .connect_timeout(TEST_TIMEOUT) .timeout(TEST_TIMEOUT); - let mut client = CodeModeHostClient::connect(endpoint) - .await - .context("failed to connect to gRPC code-mode host")?; - let mut events = client - .open_session(grpc::OpenSessionRequest { - cell_execution_limits: None, - }) - .await - .context("failed to open gRPC code-mode session")? - .into_inner(); - let event = timeout(TEST_TIMEOUT, events.message()) - .await - .context("timed out waiting for gRPC code-mode session event")? - .context("failed to read gRPC code-mode session event")? - .context("gRPC code-mode session ended before opening")?; - assert!(matches!( - event.event, - Some(grpc::session_event::Event::Opened(opened)) if !opened.session_id.is_empty() - )); + for encoding in [None, Some(tonic::codec::CompressionEncoding::Zstd)] { + let mut client = CodeModeHostClient::connect(endpoint.clone()) + .await + .context("failed to connect to gRPC code-mode host")?; + if let Some(encoding) = encoding { + client = client.accept_compressed(encoding); + } + let response = client + .open_session(grpc::OpenSessionRequest { + cell_execution_limits: None, + }) + .await + .context("failed to open gRPC code-mode session")?; + assert_eq!( + response + .metadata() + .get("grpc-encoding") + .map(|value| value.to_str()) + .transpose()?, + encoding.map(|_| "zstd"), + ); + let mut events = response.into_inner(); + let event = timeout(TEST_TIMEOUT, events.message()) + .await + .context("timed out waiting for gRPC code-mode session event")? + .context("failed to read gRPC code-mode session event")? + .context("gRPC code-mode session ended before opening")?; + assert!(matches!( + event.event, + Some(grpc::session_event::Event::Opened(opened)) if !opened.session_id.is_empty() + )); + } Ok(()) } diff --git a/codex-rs/code-mode/Cargo.toml b/codex-rs/code-mode/Cargo.toml index daa2de5471..52cdae6d08 100644 --- a/codex-rs/code-mode/Cargo.toml +++ b/codex-rs/code-mode/Cargo.toml @@ -26,7 +26,7 @@ reqwest = { workspace = true, features = ["stream"] } serde_json = { workspace = true } tokio = { workspace = true, features = ["io-util", "macros", "net", "process", "rt", "sync", "time"] } tokio-util = { workspace = true, features = ["rt"] } -tonic = { workspace = true } +tonic = { workspace = true, features = ["zstd"] } tower = { version = "0.5.3", features = ["util"] } tracing = { workspace = true } uuid = { workspace = true, features = ["v4"] } diff --git a/codex-rs/code-mode/src/grpc_session/transport.rs b/codex-rs/code-mode/src/grpc_session/transport.rs index 0f9b336da1..f16a3fef88 100644 --- a/codex-rs/code-mode/src/grpc_session/transport.rs +++ b/codex-rs/code-mode/src/grpc_session/transport.rs @@ -42,6 +42,7 @@ impl SharedTransport { let transport = channel.map_err(|error| Status::from_error(Box::new(error))); Self::Connected( CodeModeHostClient::new(BoxCloneSyncService::new(transport)) + .accept_compressed(tonic::codec::CompressionEncoding::Zstd) .max_decoding_message_size(MAX_FRAME_BYTES) .max_encoding_message_size(MAX_FRAME_BYTES), ) @@ -132,6 +133,7 @@ impl SharedTransport { CodeModeHostClient::with_origin(BoxCloneSyncService::new(transport), origin) }; let client = client + .accept_compressed(tonic::codec::CompressionEncoding::Zstd) .max_decoding_message_size(MAX_FRAME_BYTES) .max_encoding_message_size(MAX_FRAME_BYTES); *cached = Some((bound_factory, client.clone())); diff --git a/patches/BUILD.bazel b/patches/BUILD.bazel index 1c62ed89f2..715f68299a 100644 --- a/patches/BUILD.bazel +++ b/patches/BUILD.bazel @@ -24,6 +24,7 @@ exports_files([ "ring_windows_msvc_include_dirs.patch", "tikv-jemalloc-sys_make_env.patch", "tikv-jemalloc-sys_dependency_flags.patch", + "tonic-zstd-level1.patch", "v8_array_sort_elements_kind.patch", "v8_source_portability.patch", "webrtc-sys_hermetic_darwin_sysroot.patch", diff --git a/patches/tonic-zstd-level1.patch b/patches/tonic-zstd-level1.patch new file mode 100644 index 0000000000..3b66589907 --- /dev/null +++ b/patches/tonic-zstd-level1.patch @@ -0,0 +1,14 @@ +--- a/src/codec/compression.rs ++++ b/src/codec/compression.rs +@@ -239,8 +239,9 @@ + CompressionEncoding::Zstd => { + let mut zstd_encoder = Encoder::new( + &decompressed_buf[0..len], +- // FIXME: support customizing the compression level +- zstd::DEFAULT_COMPRESSION_LEVEL, ++ // Favor CodeModeHost response latency over compression ratio. ++ // Level 1 is the fastest positive level and matches the benchmark. ++ 1, + )?; + std::io::copy(&mut zstd_encoder, &mut out_writer)?; + }