diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index c28b3752ea..962dfa6219 100644 --- a/codex-rs/Cargo.lock +++ b/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", diff --git a/codex-rs/app-server-client/Cargo.toml b/codex-rs/app-server-client/Cargo.toml index daa6ef661f..272173b83a 100644 --- a/codex-rs/app-server-client/Cargo.toml +++ b/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 } diff --git a/codex-rs/app-server-client/src/lib.rs b/codex-rs/app-server-client/src/lib.rs index 4d42819eb4..e7df794ea8 100644 --- a/codex-rs/app-server-client/src/lib.rs +++ b/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; diff --git a/codex-rs/app-server-client/src/remote.rs b/codex-rs/app-server-client/src/remote.rs index c5d9b0321f..ad8d28bcec 100644 --- a/codex-rs/app-server-client/src/remote.rs +++ b/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( channel_capacity: usize, endpoint: String, - stream: WebSocketStream, + stream: WebSocketHandshake, initialize_params: InitializeParams, ) -> IoResult 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, -) -> IoResult<(String, WebSocketStream>)> { +) -> IoResult<(String, WebSocketHandshake>)> { 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)> { +) -> IoResult<(String, WebSocketHandshake)> { 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 { + 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( - stream: &mut WebSocketStream, + stream: &mut RemoteWebSocket, endpoint: &str, params: InitializeParams, initialize_timeout: Duration, @@ -1030,7 +1038,7 @@ fn jsonrpc_notification_from_client_notification( } async fn write_jsonrpc_message( - stream: &mut WebSocketStream, + stream: &mut RemoteWebSocket, 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 { diff --git a/codex-rs/app-server-client/src/remote_websocket.rs b/codex-rs/app-server-client/src/remote_websocket.rs new file mode 100644 index 0000000000..5caf0db845 --- /dev/null +++ b/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 = (WebSocketStream, Response); + +pub(crate) struct RemoteWebSocket { + stream: WebSocketStream, + codec: Option<(ZstdEncoder, ZstdDecoder)>, + max_frame_bytes: usize, + write_failed: bool, +} + +impl RemoteWebSocket { + pub(crate) fn from_handshake( + (stream, response): WebSocketHandshake, + max_message_size: usize, + ) -> io::Result { + 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::().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> { + 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 + } +} diff --git a/codex-rs/app-server-client/tests/remote_compression.rs b/codex-rs/app-server-client/tests/remote_compression.rs new file mode 100644 index 0000000000..cbab159c1d --- /dev/null +++ b/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> +{ + 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>(()) + }; + 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>(()) + }; + tokio::try_join!(server, client)?; + Ok(()) +}