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

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

codex-rs/Cargo.lock

@@ -15931,6 +15931,7 @@ dependencies = [
  "tower-layer",
  "tower-service",
  "tracing",
+ "zstd",
 ]
 
 [[package]]

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"] }

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

codex-rs/code-mode-host/tests/grpc_tcp.rs

@@ -117,6 +117,7 @@ async fn http1_response(endpoint: &str, request: &[u8]) -> Result<String> {
     .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(())
 }

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"] }

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()));

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

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)?;
+         }