From 4ca84a10696623fcd7c200db3df9c41671d3bad0 Mon Sep 17 00:00:00 2001 From: ostepanian Date: Wed, 14 Oct 2026 23:53:14 +0000 Subject: [PATCH] Allow large file writes over direct exec-server WebSockets (#54135) ## Why Direct mode serves executor requests but applied the client's 8 KiB inbound request limit, causing larger `fs/writeFile` requests to disconnect the executor. ## What changed Use a dedicated direct WebSocket connection constructor that omits the client inbound request limit while retaining WebSocket message and frame limits. ## Testing Add a regression test that runs the direct environment and verifies successful file writes above the client request limit using both text and binary WebSocket messages. GitOrigin-RevId: e8e77a5e8581ab9eb98446e23b77a6da41d64868 --- codex-rs/exec-server/src/connection.rs | 16 +++ codex-rs/exec-server/src/remote/direct.rs | 2 +- .../exec-server/src/remote/direct_tests.rs | 97 +++++++++++++++++++ 3 files changed, 114 insertions(+), 1 deletion(-) diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs index e18faffd1e..62f98dfcc3 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/codex-rs/exec-server/src/connection.rs @@ -481,6 +481,22 @@ impl JsonRpcConnection { ) } + /// Direct mode dials the WebSocket but serves executor requests from the + /// harness. Do not apply the harness client's inbound request limit here; + /// the WebSocket transport still enforces its message and frame limits. + pub(crate) fn from_direct_websocket(stream: T, connection_label: String) -> Self + where + T: Sink + Stream> + Unpin + Send + 'static, + E: std::fmt::Display + Send + 'static, + { + Self::from_websocket_stream( + stream, + connection_label, + /*ping_interval*/ None, + /*max_inbound_request_len*/ None, + ) + } + pub(crate) fn from_axum_websocket(stream: AxumWebSocket, connection_label: String) -> Self { Self::from_websocket_stream( stream, diff --git a/codex-rs/exec-server/src/remote/direct.rs b/codex-rs/exec-server/src/remote/direct.rs index a1446a9ec4..a99454e177 100644 --- a/codex-rs/exec-server/src/remote/direct.rs +++ b/codex-rs/exec-server/src/remote/direct.rs @@ -167,7 +167,7 @@ pub(super) async fn run_direct_environment( ); processor .run_connection( - JsonRpcConnection::from_websocket( + JsonRpcConnection::from_direct_websocket( websocket, format!( "direct exec-server websocket {}", diff --git a/codex-rs/exec-server/src/remote/direct_tests.rs b/codex-rs/exec-server/src/remote/direct_tests.rs index 73e9345b9b..ae40fa5db6 100644 --- a/codex-rs/exec-server/src/remote/direct_tests.rs +++ b/codex-rs/exec-server/src/remote/direct_tests.rs @@ -5,11 +5,14 @@ use std::time::Duration; use std::time::Instant; use anyhow::Result; +use base64::Engine; use codex_api::AuthProvider; use codex_http_client::DestinationPolicy; use codex_http_client::HttpClientFactory; use codex_http_client::NetworkPolicyController; use codex_http_client::OutboundProxyPolicy; +use futures::SinkExt; +use futures::StreamExt; use http::HeaderMap; use http::HeaderValue; use pretty_assertions::assert_eq; @@ -18,6 +21,7 @@ use tokio::io::AsyncWriteExt; use tokio::net::TcpListener; use tokio::time::timeout; use tokio_tungstenite::accept_hdr_async; +use tokio_tungstenite::tungstenite::Message; use tokio_tungstenite::tungstenite::handshake::server::Request as HandshakeRequest; use tokio_tungstenite::tungstenite::handshake::server::Response as HandshakeResponse; use wiremock::Mock; @@ -61,6 +65,99 @@ impl AuthProvider for StaticAuthProvider { #[derive(Debug)] struct QueryAuthProvider; +#[tokio::test] +async fn direct_environment_accepts_large_file_writes() -> Result<()> { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let registry = MockServer::start().await; + Mock::given(method("POST")) + .and(path( + "/cloud/environment/environment-requested/direct/register", + )) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "environment_id": "environment-requested", + "transport": DIRECT_TRANSPORT, + "registration_id": "registration-1", + "url": format!("ws://{}", listener.local_addr()?), + }))) + .expect(1) + .mount(®istry) + .await; + let config = RemoteEnvironmentConfig::new_with_transport( + registry.uri(), + "environment-requested".to_string(), + RemoteEnvironmentTransport::Direct, + Arc::new(StaticAuthProvider), + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + )?; + let runtime_paths = ExecServerRuntimeOptions::new( + std::env::current_exe()?, + /*codex_linux_sandbox_exe*/ None, + )?; + let task = tokio_util::task::AbortOnDropHandle::new(tokio::spawn( + crate::run_remote_environment(config, runtime_paths), + )); + let (socket, _) = timeout(Duration::from_secs(5), listener.accept()).await??; + let mut websocket = tokio_tungstenite::accept_async(socket).await?; + websocket + .send(Message::Text( + serde_json::json!({ + "id": 1, "method": "initialize", "params": {"clientName": "direct-test"}, + }) + .to_string() + .into(), + )) + .await?; + let response = timeout(Duration::from_secs(5), websocket.next()) + .await? + .expect("initialize response")?; + let response: serde_json::Value = serde_json::from_slice(&response.into_data())?; + assert_eq!(response["id"], 1); + assert!(response.get("result").is_some(), "{response}"); + websocket + .send(Message::Text( + serde_json::json!({ + "method": "initialized", "params": {}, + }) + .to_string() + .into(), + )) + .await?; + + let directory = tempfile::tempdir()?; + let output = directory.path().join("report.txt"); + let path = codex_utils_path_uri::PathUri::from_host_native_path(&output)?; + let data = "report contents\n".repeat(1024); + // Exercise both JSON-RPC WebSocket encodings on the production direct runner. + for binary in [false, true] { + let request = serde_json::json!({ + "id": 2, "method": "fs/writeFile", "params": { + "path": path, + "dataBase64": base64::engine::general_purpose::STANDARD.encode(&data), + }, + }) + .to_string(); + assert!( + request.len() > crate::client_inbound_request_limit::MAX_CLIENT_INBOUND_REQUEST_LEN + ); + let frame = if binary { + Message::Binary(request.into_bytes().into()) + } else { + Message::Text(request.into()) + }; + websocket.send(frame).await?; + let response = timeout(Duration::from_secs(5), websocket.next()) + .await? + .expect("large write must not disconnect the executor")?; + let response: serde_json::Value = serde_json::from_slice(&response.into_data())?; + assert_eq!(response["id"], 2); + assert_eq!(response["result"], serde_json::json!({}), "{response}"); + assert_eq!(tokio::fs::read(&output).await?, data.as_bytes()); + } + registry.verify().await; + task.abort(); + Ok(()) +} + impl AuthProvider for QueryAuthProvider { fn add_auth_headers(&self, _headers: &mut HeaderMap) {}