Support negotiated WebSocket compression in the remote app-server client (#57218)

## What changed

Advertise `zstd-stream-v1` compression for WebSocket and Unix socket connections and use the negotiated codec for JSON-RPC messages. Retain compression history for the connection and keep plain-text messaging when compression is not negotiated.

Fragment compressed writes to respect the server's advertised frame limit, reject text messages in compressed sessions, and stop using the connection after a failed write to avoid mismatched compression history.

## Testing

Add an integration test covering compression negotiation, fragmentation, retained codec history, request/response handling, and disconnection on an uncompressed text message.

GitOrigin-RevId: 510073444ea49376a116c05214bcb79bd0bc3b50

codex-rs/Cargo.lock

@@ -2179,6 +2179,7 @@ version = "0.0.0"
 dependencies = [
  "codex-app-server",
  "codex-app-server-protocol",
+ "codex-app-server-transport",
  "codex-arg0",
  "codex-config",
  "codex-core",

codex-rs/app-server-client/Cargo.toml

@@ -15,6 +15,7 @@ workspace = true
 [dependencies]
 codex-app-server = { workspace = true }
 codex-app-server-protocol = { workspace = true }
+codex-app-server-transport = { workspace = true }
 codex-arg0 = { workspace = true }
 codex-config = { workspace = true }
 codex-core = { workspace = true }

codex-rs/app-server-client/src/lib.rs

@@ -23,6 +23,7 @@
 mod cloud;
 mod path;
 mod remote;
+mod remote_websocket;
 pub use cloud::CloudClient;
 pub use cloud::CloudCommand;
 pub use cloud::CloudOperation;

codex-rs/app-server-client/src/remote.rs

@@ -22,6 +22,8 @@ use crate::RequestResult;
 use crate::SHUTDOWN_TIMEOUT;
 use crate::TypedRequestError;
 use crate::decode_typed_response;
+use crate::remote_websocket::RemoteWebSocket;
+use crate::remote_websocket::WebSocketHandshake;
 use codex_app_server_protocol::ClientInfo;
 use codex_app_server_protocol::ClientNotification;
 use codex_app_server_protocol::ClientRequest;
@@ -37,11 +39,11 @@ use codex_app_server_protocol::RequestId;
 use codex_app_server_protocol::Result as JsonRpcResult;
 use codex_app_server_protocol::ServerNotification;
 use codex_app_server_protocol::ServerRequest;
+use codex_app_server_transport::WEBSOCKET_COMPRESSION_HEADER;
+use codex_app_server_transport::ZSTD_STREAM_PROTOCOL;
 use codex_uds::UnixStream;
 use codex_utils_absolute_path::AbsolutePathBuf;
 use codex_utils_rustls_provider::ensure_rustls_crypto_provider;
-use futures::SinkExt;
-use futures::StreamExt;
 use serde::de::DeserializeOwned;
 use tokio::io::AsyncRead;
 use tokio::io::AsyncWrite;
@@ -50,12 +52,12 @@ use tokio::sync::mpsc;
 use tokio::sync::oneshot;
 use tokio::time::timeout;
 use tokio_tungstenite::MaybeTlsStream;
-use tokio_tungstenite::WebSocketStream;
 use tokio_tungstenite::client_async_with_config;
 use tokio_tungstenite::connect_async_with_config;
 use tokio_tungstenite::tungstenite::Error as TungsteniteError;
 use tokio_tungstenite::tungstenite::Message;
 use tokio_tungstenite::tungstenite::client::IntoClientRequest;
+use tokio_tungstenite::tungstenite::handshake::client::Request;
 use tokio_tungstenite::tungstenite::http::HeaderValue;
 use tokio_tungstenite::tungstenite::http::header::AUTHORIZATION;
 use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
@@ -239,13 +241,14 @@ impl RemoteAppServerClient {
     async fn connect_with_stream<S>(
         channel_capacity: usize,
         endpoint: String,
-        stream: WebSocketStream<S>,
+        stream: WebSocketHandshake<S>,
         initialize_params: InitializeParams,
     ) -> IoResult<Self>
     where
         S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
     {
-        let mut stream = stream;
+        let mut stream =
+            RemoteWebSocket::from_handshake(stream, REMOTE_APP_SERVER_MAX_WEBSOCKET_MESSAGE_SIZE)?;
         let (pending_events, metadata) = initialize_remote_connection(
             &mut stream,
             &endpoint,
@@ -264,7 +267,7 @@ impl RemoteAppServerClient {
                 tokio::select! {
                     command = command_rx.recv() => {
                         let Some(command) = command else {
-                            let _ = stream.close(None).await;
+                            let _ = stream.close().await;
                             break;
                         };
                         match command {
@@ -346,7 +349,7 @@ impl RemoteAppServerClient {
                                 let _ = response_tx.send(result);
                             }
                             RemoteClientCommand::Shutdown { response_tx } => {
-                                let close_result = stream.close(None).await.or_else(|err| {
+                                let close_result = stream.close().await.or_else(|err| {
                                     if websocket_close_error_is_already_closed(&err) {
                                         Ok(())
                                     } else {
@@ -719,7 +722,7 @@ impl RemoteAppServerRequestHandle {
 async fn connect_websocket_endpoint(
     websocket_url: String,
     auth_token: Option<String>,
-) -> IoResult<(String, WebSocketStream<MaybeTlsStream<TcpStream>>)> {
+) -> IoResult<(String, WebSocketHandshake<MaybeTlsStream<TcpStream>>)> {
     let url = Url::parse(&websocket_url).map_err(|err| {
         IoError::new(
             ErrorKind::InvalidInput,
@@ -735,7 +738,7 @@ async fn connect_websocket_endpoint(
         ));
     }
 
-    let mut request = url.as_str().into_client_request().map_err(|err| {
+    let mut request = websocket_request(url.as_str()).map_err(|err| {
         IoError::new(
             ErrorKind::InvalidInput,
             format!("invalid websocket URL `{websocket_url}`: {err}"),
@@ -769,7 +772,6 @@ async fn connect_websocket_endpoint(
             format!("timed out connecting to remote app server at `{websocket_url}`"),
         )
     })?
-    .map(|(stream, _response)| stream)
     .map_err(|err| {
         IoError::other(format!(
             "failed to connect to remote app server at `{websocket_url}`: {err}"
@@ -782,16 +784,14 @@ async fn connect_websocket_endpoint(
 async fn connect_unix_socket_endpoint(
     socket_path: AbsolutePathBuf,
     peer_policy: SocketPeerPolicy,
-) -> IoResult<(String, WebSocketStream<UnixStream>)> {
+) -> IoResult<(String, WebSocketHandshake<UnixStream>)> {
     let endpoint = format!("unix://{}", socket_path.display());
-    let request = UDS_WEBSOCKET_HANDSHAKE_URL
-        .into_client_request()
-        .map_err(|err| {
-            IoError::new(
-                ErrorKind::InvalidInput,
-                format!("invalid UDS websocket handshake URL: {err}"),
-            )
-        })?;
+    let request = websocket_request(UDS_WEBSOCKET_HANDSHAKE_URL).map_err(|err| {
+        IoError::new(
+            ErrorKind::InvalidInput,
+            format!("invalid UDS websocket handshake URL: {err}"),
+        )
+    })?;
     let stream = timeout(CONNECT_TIMEOUT, UnixStream::connect(socket_path.as_path()))
         .await
         .map_err(|_| {
@@ -822,7 +822,6 @@ async fn connect_unix_socket_endpoint(
             format!("timed out upgrading remote app server at `{endpoint}`"),
         )
     })?
-    .map(|(stream, _response)| stream)
     .map_err(|err| {
         IoError::other(format!(
             "failed to upgrade remote app server at `{endpoint}`: {err}"
@@ -832,6 +831,15 @@ async fn connect_unix_socket_endpoint(
     Ok((endpoint, stream))
 }
 
+fn websocket_request(url: &str) -> Result<Request, TungsteniteError> {
+    let mut request = url.into_client_request()?;
+    request.headers_mut().insert(
+        WEBSOCKET_COMPRESSION_HEADER,
+        HeaderValue::from_static(ZSTD_STREAM_PROTOCOL),
+    );
+    Ok(request)
+}
+
 fn remote_websocket_config() -> WebSocketConfig {
     WebSocketConfig::default()
         .max_frame_size(Some(REMOTE_APP_SERVER_MAX_WEBSOCKET_MESSAGE_SIZE))
@@ -839,7 +847,7 @@ fn remote_websocket_config() -> WebSocketConfig {
 }
 
 async fn initialize_remote_connection<S>(
-    stream: &mut WebSocketStream<S>,
+    stream: &mut RemoteWebSocket<S>,
     endpoint: &str,
     params: InitializeParams,
     initialize_timeout: Duration,
@@ -1030,7 +1038,7 @@ fn jsonrpc_notification_from_client_notification(
 }
 
 async fn write_jsonrpc_message<S>(
-    stream: &mut WebSocketStream<S>,
+    stream: &mut RemoteWebSocket<S>,
     message: JSONRPCMessage,
     endpoint: &str,
 ) -> IoResult<()>
@@ -1038,14 +1046,11 @@ where
     S: AsyncRead + AsyncWrite + Unpin,
 {
     let payload = serde_json::to_string(&message).map_err(IoError::other)?;
-    stream
-        .send(Message::Text(payload.into()))
-        .await
-        .map_err(|err| {
-            IoError::other(format!(
-                "failed to write websocket message to `{endpoint}`: {err}"
-            ))
-        })
+    stream.send_text(payload).await.map_err(|err| {
+        IoError::other(format!(
+            "failed to write websocket message to `{endpoint}`: {err}"
+        ))
+    })
 }
 
 fn websocket_close_error_is_already_closed(err: &TungsteniteError) -> bool {

codex-rs/app-server-client/src/remote_websocket.rs

@@ -0,0 +1,117 @@
+//! Negotiated WebSocket compression. Codec history lasts for the connection.
+
+use std::io;
+
+use codex_app_server_transport::WEBSOCKET_COMPRESSION_HEADER;
+use codex_app_server_transport::WebSocketCompression;
+use codex_app_server_transport::ZstdDecoder;
+use codex_app_server_transport::ZstdEncoder;
+use futures::SinkExt;
+use futures::StreamExt;
+use tokio::io::AsyncRead;
+use tokio::io::AsyncWrite;
+use tokio_tungstenite::WebSocketStream;
+use tokio_tungstenite::tungstenite::Bytes;
+use tokio_tungstenite::tungstenite::Error;
+use tokio_tungstenite::tungstenite::Message;
+use tokio_tungstenite::tungstenite::handshake::client::Response;
+use tokio_tungstenite::tungstenite::http::HeaderValue;
+use tokio_tungstenite::tungstenite::protocol::frame::Frame;
+use tokio_tungstenite::tungstenite::protocol::frame::coding::Data;
+use tokio_tungstenite::tungstenite::protocol::frame::coding::OpCode;
+
+// Match the listener's default incoming frame limit when no smaller cap is advertised.
+const DEFAULT_MAX_FRAME_BYTES: usize = 16 << 20;
+const MAX_UNFRAGMENTED_MESSAGE_BYTES_HEADER: &str =
+    "x-codex-websocket-max-unfragmented-message-bytes";
+
+pub(crate) type WebSocketHandshake<S> = (WebSocketStream<S>, Response);
+
+pub(crate) struct RemoteWebSocket<S> {
+    stream: WebSocketStream<S>,
+    codec: Option<(ZstdEncoder, ZstdDecoder)>,
+    max_frame_bytes: usize,
+    write_failed: bool,
+}
+
+impl<S: AsyncRead + AsyncWrite + Unpin> RemoteWebSocket<S> {
+    pub(crate) fn from_handshake(
+        (stream, response): WebSocketHandshake<S>,
+        max_message_size: usize,
+    ) -> io::Result<Self> {
+        let headers = response.headers();
+        let header = headers
+            .get(WEBSOCKET_COMPRESSION_HEADER)
+            .map(HeaderValue::to_str)
+            .transpose()
+            .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?;
+        let codec = match WebSocketCompression::from_response_header(header)? {
+            WebSocketCompression::Plain => None,
+            WebSocketCompression::Zstd => {
+                Some((ZstdEncoder::new()?, ZstdDecoder::new(max_message_size)?))
+            }
+        };
+        let max_frame_bytes = headers
+            .get(MAX_UNFRAGMENTED_MESSAGE_BYTES_HEADER)
+            .and_then(|value| value.to_str().ok())
+            .and_then(|value| value.parse::<usize>().ok())
+            .filter(|value| *value > 0)
+            .unwrap_or(DEFAULT_MAX_FRAME_BYTES)
+            .min(DEFAULT_MAX_FRAME_BYTES);
+        Ok(Self {
+            stream,
+            codec,
+            max_frame_bytes,
+            write_failed: false,
+        })
+    }
+
+    pub(crate) async fn next(&mut self) -> Option<Result<Message, Error>> {
+        if self.write_failed {
+            return Some(Err(Error::Io(io::ErrorKind::BrokenPipe.into())));
+        }
+        let message = self.stream.next().await?;
+        Some(match (&mut self.codec, message) {
+            (Some((_, decoder)), Ok(Message::Binary(bytes))) => decoder
+                .decode(&bytes)
+                .map(|text| Message::Text(text.into()))
+                .map_err(Error::Io),
+            (Some(_), Ok(Message::Text(_))) => Err(Error::Io(io::Error::new(
+                io::ErrorKind::InvalidData,
+                "expected compressed binary app-server message",
+            ))),
+            (_, message) => message,
+        })
+    }
+
+    pub(crate) async fn send_text(&mut self, text: String) -> Result<(), Error> {
+        if self.write_failed {
+            return Err(Error::Io(io::ErrorKind::BrokenPipe.into()));
+        }
+        let result = async {
+            let Some((encoder, _)) = &mut self.codec else {
+                return self.stream.send(Message::Text(text.into())).await;
+            };
+            let mut bytes = Bytes::from(encoder.encode(&text).map_err(Error::Io)?);
+            if bytes.len() <= self.max_frame_bytes {
+                return self.stream.send(Message::Binary(bytes)).await;
+            }
+            let mut opcode = OpCode::Data(Data::Binary);
+            while !bytes.is_empty() {
+                let chunk = bytes.split_to(bytes.len().min(self.max_frame_bytes));
+                let frame = Frame::message(chunk, opcode, bytes.is_empty());
+                self.stream.feed(Message::Frame(frame)).await?;
+                opcode = OpCode::Data(Data::Continue);
+            }
+            self.stream.flush().await
+        }
+        .await;
+        // A failed write may leave the compression history ahead of the receiver.
+        self.write_failed = result.is_err();
+        result
+    }
+
+    pub(crate) async fn close(&mut self) -> Result<(), Error> {
+        self.stream.close(/*msg*/ None).await
+    }
+}

codex-rs/app-server-client/tests/remote_compression.rs

@@ -0,0 +1,115 @@
+//! Exercise a negotiated compressed session through the remote client API.
+
+use std::error::Error;
+
+use codex_app_server_client::AppServerEvent;
+use codex_app_server_client::RemoteAppServerClient;
+use codex_app_server_client::RemoteAppServerConnectArgs;
+use codex_app_server_client::RemoteAppServerEndpoint;
+use codex_app_server_protocol::JSONRPCRequest;
+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 codex_uds::UnixListener;
+use codex_utils_absolute_path::AbsolutePathBuf;
+use futures::SinkExt;
+use futures::StreamExt;
+use pretty_assertions::assert_eq;
+use serde_json::Value;
+use serde_json::json;
+use tempfile::TempDir;
+use tokio_tungstenite::accept_hdr_async_with_config;
+use tokio_tungstenite::tungstenite::Message;
+use tokio_tungstenite::tungstenite::handshake::server::Request;
+use tokio_tungstenite::tungstenite::handshake::server::Response;
+use tokio_tungstenite::tungstenite::http::HeaderValue;
+use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
+
+#[tokio::test]
+async fn compressed_session_negotiates_fragments_and_retains_history() -> Result<(), Box<dyn Error>>
+{
+    const FRAME_LIMIT: usize = 16;
+    let dir = TempDir::new()?;
+    let socket_path = AbsolutePathBuf::from_absolute_path(dir.path().join("codex.sock"))?;
+    let mut listener = UnixListener::bind(socket_path.as_path()).await?;
+    let server = async {
+        let stream = listener.accept().await?;
+        let mut websocket = accept_hdr_async_with_config(
+            stream,
+            |request: &Request, mut response: Response| {
+                assert_eq!(
+                    request.headers()[WEBSOCKET_COMPRESSION_HEADER],
+                    ZSTD_STREAM_PROTOCOL
+                );
+                response.headers_mut().insert(
+                    WEBSOCKET_COMPRESSION_HEADER,
+                    HeaderValue::from_static(ZSTD_STREAM_PROTOCOL),
+                );
+                response.headers_mut().insert(
+                    "x-codex-websocket-max-unfragmented-message-bytes",
+                    HeaderValue::from(FRAME_LIMIT),
+                );
+                Ok(response)
+            },
+            Some(WebSocketConfig::default().max_frame_size(Some(FRAME_LIMIT))),
+        )
+        .await?;
+        let mut decoder = ZstdDecoder::new(/*max_message_size*/ 1024)?;
+        let mut encoder = ZstdEncoder::new()?;
+        for expected_method in ["initialize", "initialized", "echo"] {
+            let Message::Binary(bytes) = websocket.next().await.ok_or("connection closed")?? else {
+                return Err("expected compressed message".into());
+            };
+            if expected_method == "initialize" {
+                assert!(
+                    bytes.len() > FRAME_LIMIT,
+                    "initialize must require fragmentation"
+                );
+            }
+            let request: Value = serde_json::from_str(&decoder.decode(&bytes)?)?;
+            assert_eq!(request["method"], expected_method);
+            if expected_method != "initialized" {
+                let response = json!({"id": request["id"], "result": {"ok": true}}).to_string();
+                websocket
+                    .send(Message::Binary(encoder.encode(&response)?.into()))
+                    .await?;
+            }
+        }
+        websocket.send(Message::Text("{}".into())).await?;
+        Ok::<(), Box<dyn Error>>(())
+    };
+    let client = async {
+        let mut client = RemoteAppServerClient::connect(RemoteAppServerConnectArgs {
+            endpoint: RemoteAppServerEndpoint::UnixSocket { socket_path },
+            client_name: "compression-test".to_string(),
+            client_version: "0.0.0".to_string(),
+            experimental_api: false,
+            mcp_server_openai_form_elicitation: false,
+            opt_out_notification_methods: Vec::new(),
+            channel_capacity: 8,
+        })
+        .await?;
+        assert_eq!(
+            client
+                .request_handle()
+                .request_json_rpc(JSONRPCRequest {
+                    id: RequestId::Integer(1),
+                    method: "echo".to_string(),
+                    params: None,
+                    trace: None,
+                })
+                .await?,
+            Ok(json!({"ok": true}))
+        );
+        let Some(AppServerEvent::Disconnected { message }) = client.next_event().await else {
+            return Err("expected disconnect after uncompressed message".into());
+        };
+        assert!(message.contains("expected compressed binary"));
+        client.shutdown().await?;
+        Ok::<(), Box<dyn Error>>(())
+    };
+    tokio::try_join!(server, client)?;
+    Ok(())
+}