From f886e33e5a55fcb6637d1d05c0f44cab9fcb933f Mon Sep 17 00:00:00 2001 From: sayan-oai Date: Thu, 18 Jun 2026 22:27:45 -0700 Subject: [PATCH] [3/3] app-server: configure environment connection timeout (#29025) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Why Remote environments registered through `environment/add` currently use the fixed 10-second WebSocket connection timeout. Slow-starting executors need a caller-selected connection window, but this should not add retry policy or couple exec-server behavior to Core’s `deferred_executor` feature. Make the timeout an optional part of the existing experimental request. Existing clients continue using the current default, while callers that know an executor may take longer can request a larger window explicitly. Depends on #28683. ## What changed - Add optional `connectTimeoutMs` to `EnvironmentAddParams` and document it in the app-server README. - Pass the optional timeout through `EnvironmentRequestProcessor` into one `EnvironmentManager::upsert_environment()` path; the manager applies the existing default when it is omitted. - Preserve the existing single-attempt lifecycle. The configured value controls WebSocket connection and handshake time for both initial connection and later reconnects; initialization retains its separate timeout. - Add an app-server integration test that sends the real JSON-RPC request and verifies a stalled handshake observes the requested timeout. ## Test plan - `just test -p codex-app-server-protocol` - `just test -p codex-exec-server` - `just test -p codex-app-server environment_add_applies_connect_timeout` ## Rollout This is additive and does not enable `deferred_executor`. Callers should send a non-default timeout only after a compatible app-server is deployed; omitted or `null` values retain the existing 10-second default. --- .../src/protocol/common.rs | 6 ++- .../src/protocol/v2/environment.rs | 4 ++ codex-rs/app-server/README.md | 2 +- .../environment_processor.rs | 7 ++- .../tests/suite/v2/environment_add.rs | 52 +++++++++++++++++++ codex-rs/app-server/tests/suite/v2/mod.rs | 1 + codex-rs/core/src/environment_selection.rs | 3 ++ codex-rs/exec-server/src/client_api.rs | 4 +- codex-rs/exec-server/src/environment.rs | 38 +++++++++++--- .../src/remote_file_system_path_uri_tests.rs | 6 ++- 10 files changed, 110 insertions(+), 13 deletions(-) create mode 100644 codex-rs/app-server/tests/suite/v2/environment_add.rs diff --git a/codex-rs/app-server-protocol/src/protocol/common.rs b/codex-rs/app-server-protocol/src/protocol/common.rs index 4e9a78539..643afcc58 100644 --- a/codex-rs/app-server-protocol/src/protocol/common.rs +++ b/codex-rs/app-server-protocol/src/protocol/common.rs @@ -2011,6 +2011,7 @@ mod tests { params: v2::EnvironmentAddParams { environment_id: "remote-a".to_string(), exec_server_url: "ws://127.0.0.1:8765".to_string(), + connect_timeout_ms: None, }, }; assert_eq!( @@ -2958,6 +2959,7 @@ mod tests { params: v2::EnvironmentAddParams { environment_id: "remote-a".to_string(), exec_server_url: "ws://127.0.0.1:8765".to_string(), + connect_timeout_ms: Some(300_000), }, }; assert_eq!( @@ -2966,7 +2968,8 @@ mod tests { "id": 9, "params": { "environmentId": "remote-a", - "execServerUrl": "ws://127.0.0.1:8765" + "execServerUrl": "ws://127.0.0.1:8765", + "connectTimeoutMs": 300000 } }), serde_json::to_value(&request)?, @@ -3392,6 +3395,7 @@ mod tests { params: v2::EnvironmentAddParams { environment_id: "remote-a".to_string(), exec_server_url: "ws://127.0.0.1:8765".to_string(), + connect_timeout_ms: None, }, }; let reason = crate::experimental_api::ExperimentalApi::experimental_reason(&request); diff --git a/codex-rs/app-server-protocol/src/protocol/v2/environment.rs b/codex-rs/app-server-protocol/src/protocol/v2/environment.rs index 294ae736f..ccffd5813 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/environment.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/environment.rs @@ -9,6 +9,10 @@ use ts_rs::TS; pub struct EnvironmentAddParams { pub environment_id: String, pub exec_server_url: String, + /// Optional WebSocket connection timeout. The server default applies when omitted. + #[ts(type = "number | null")] + #[ts(optional = nullable)] + pub connect_timeout_ms: Option, } #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index 5035ebf71..f4f3bbc3b 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -204,7 +204,7 @@ Example with notification opt-out: - `experimentalFeature/list` — list feature flags with stage metadata (`beta`, `underDevelopment`, `stable`, etc.), enabled/default-enabled state, and cursor pagination. Pass `threadId` when showing feature state for an existing loaded thread so `enabled` is computed from that thread's refreshed config, including project-local config for the thread's cwd; if omitted, the server uses its default config resolution context. For non-beta flags, `displayName`/`description`/`announcement` are `null`. - `permissionProfile/list` — beta; list available permission profile ids with optional display `description` text, using cursor pagination. Pass `cwd` when the caller needs project-local `[permissions.]` entries to be included in the current catalog view. - `experimentalFeature/enablement/set` — patch the in-memory process-wide runtime feature enablement for currently supported feature keys. For each feature, precedence is: cloud requirements > --enable > config.toml > experimentalFeature/enablement/set (new) > code default. Invalid keys will be ignored. -- `environment/add` — experimental; add or replace a named remote environment by `environmentId` and `execServerUrl` for later selection by `thread/start` or `turn/start`; returns `{}` and does not change the default environment. +- `environment/add` — experimental; add or replace a named remote environment by `environmentId` and `execServerUrl` for later selection by `thread/start` or `turn/start`; optional `connectTimeoutMs` overrides the WebSocket connection timeout; returns `{}` and does not change the default environment. - `collaborationMode/list` — list available collaboration mode presets (experimental, no pagination). Built-in presets do not select a model; the Plan preset selects medium reasoning effort. This response omits built-in developer instructions; clients should either pass `settings.developer_instructions: null` when setting a mode to use Codex's built-in instructions, or provide their own instructions explicitly. - `skills/list` — list skills for one or more `cwd` values (optional `forceReload`). - `skills/extraRoots/set` — replace the app-server process runtime extra standalone skill roots. The roots are not persisted; missing directories are accepted and simply load no skills. diff --git a/codex-rs/app-server/src/request_processors/environment_processor.rs b/codex-rs/app-server/src/request_processors/environment_processor.rs index eb9b283f7..36533aa23 100644 --- a/codex-rs/app-server/src/request_processors/environment_processor.rs +++ b/codex-rs/app-server/src/request_processors/environment_processor.rs @@ -1,4 +1,5 @@ use super::*; +use std::time::Duration; #[derive(Clone)] pub(crate) struct EnvironmentRequestProcessor { @@ -17,7 +18,11 @@ impl EnvironmentRequestProcessor { params: EnvironmentAddParams, ) -> Result, JSONRPCErrorError> { self.environment_manager - .upsert_environment(params.environment_id, params.exec_server_url) + .upsert_environment( + params.environment_id, + params.exec_server_url, + params.connect_timeout_ms.map(Duration::from_millis), + ) .map_err(|err| invalid_request(err.to_string()))?; Ok(Some(EnvironmentAddResponse {}.into())) } diff --git a/codex-rs/app-server/tests/suite/v2/environment_add.rs b/codex-rs/app-server/tests/suite/v2/environment_add.rs new file mode 100644 index 000000000..bf44b2760 --- /dev/null +++ b/codex-rs/app-server/tests/suite/v2/environment_add.rs @@ -0,0 +1,52 @@ +use std::time::Duration; + +use anyhow::Result; +use app_test_support::TestAppServer; +use app_test_support::to_response; +use codex_app_server_protocol::EnvironmentAddResponse; +use codex_app_server_protocol::JSONRPCResponse; +use codex_app_server_protocol::RequestId; +use serde_json::json; +use tempfile::TempDir; +use tokio::io::AsyncReadExt; +use tokio::net::TcpListener; +use tokio::time::timeout; + +const RPC_TIMEOUT: Duration = Duration::from_secs(10); +const CONNECTION_CLOSE_TIMEOUT: Duration = Duration::from_secs(5); + +#[tokio::test] +async fn environment_add_applies_connect_timeout() -> Result<()> { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let exec_server_url = format!("ws://{}", listener.local_addr()?); + let stalled_server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await?; + let mut request = Vec::new(); + socket.read_to_end(&mut request).await?; + anyhow::ensure!(!request.is_empty(), "expected a WebSocket handshake"); + Ok::<_, anyhow::Error>(()) + }); + let codex_home = TempDir::new()?; + let mut app_server = TestAppServer::new(codex_home.path()).await?; + timeout(RPC_TIMEOUT, app_server.initialize()).await??; + + let request_id = app_server + .send_raw_request( + "environment/add", + Some(json!({ + "environmentId": "remote-a", + "execServerUrl": exec_server_url, + "connectTimeoutMs": 1_000, + })), + ) + .await?; + let response: JSONRPCResponse = timeout( + RPC_TIMEOUT, + app_server.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + let _: EnvironmentAddResponse = to_response(response)?; + + timeout(CONNECTION_CLOSE_TIMEOUT, stalled_server).await???; + Ok(()) +} diff --git a/codex-rs/app-server/tests/suite/v2/mod.rs b/codex-rs/app-server/tests/suite/v2/mod.rs index 2745e8ee0..f4400da46 100644 --- a/codex-rs/app-server/tests/suite/v2/mod.rs +++ b/codex-rs/app-server/tests/suite/v2/mod.rs @@ -13,6 +13,7 @@ mod connection_handling_websocket; mod connection_handling_websocket_unix; mod current_time; mod dynamic_tools; +mod environment_add; #[cfg(not(target_os = "windows"))] mod executor_mcp; mod executor_skills; diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index 0d82a147b..1bed53fcf 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -700,6 +700,7 @@ url = "ws://127.0.0.1:8765" .upsert_environment( REMOTE_ENVIRONMENT_ID.to_string(), format!("ws://{}", listener.local_addr().expect("listener address")), + /*connect_timeout*/ None, ) .expect("replacement environment"); environments.update_selections(std::slice::from_ref(&selection)); @@ -753,6 +754,7 @@ url = "ws://127.0.0.1:8765" .local_addr() .expect("second listener address") ), + /*connect_timeout*/ None, ) .expect("replace environment"); @@ -802,6 +804,7 @@ url = "ws://127.0.0.1:8765" .upsert_environment( REMOTE_ENVIRONMENT_ID.to_string(), "ws://127.0.0.1:9876".to_string(), + /*connect_timeout*/ None, ) .expect("replacement environment"); let environments = ThreadEnvironments::new( diff --git a/codex-rs/exec-server/src/client_api.rs b/codex-rs/exec-server/src/client_api.rs index 60dc638d4..0c6a3576d 100644 --- a/codex-rs/exec-server/src/client_api.rs +++ b/codex-rs/exec-server/src/client_api.rs @@ -136,10 +136,10 @@ impl std::fmt::Debug for ExecServerTransportParams { } impl ExecServerTransportParams { - pub(crate) fn websocket_url(websocket_url: String) -> Self { + pub(crate) fn websocket_url(websocket_url: String, connect_timeout: Duration) -> Self { Self::WebSocketUrl { websocket_url, - connect_timeout: DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT, + connect_timeout, initialize_timeout: DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT, } } diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index 03d0d996c..9496bb8be 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -11,6 +11,7 @@ use crate::NoiseChannelIdentity; use crate::NoiseRendezvousConnectProvider; use crate::client::LazyRemoteExecServerClient; use crate::client::http_client::ReqwestHttpClient; +use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT; use crate::client_api::ExecServerTransportParams; use crate::environment_provider::DefaultEnvironmentProvider; use crate::environment_provider::EnvironmentDefault; @@ -288,11 +289,13 @@ impl EnvironmentManager { } /// Adds or replaces a named remote environment without changing the - /// manager's default environment selection. + /// manager's default environment selection. Uses the default WebSocket + /// connection timeout when none is provided. pub fn upsert_environment( &self, environment_id: String, exec_server_url: String, + connect_timeout: Option, ) -> Result<(), ExecServerError> { if environment_id.is_empty() { return Err(ExecServerError::Protocol( @@ -310,8 +313,11 @@ impl EnvironmentManager { "remote environment requires an exec-server url".to_string(), )); }; - let environment = Arc::new(Environment::remote_inner( - exec_server_url, + let environment = Arc::new(Environment::remote_with_transport( + ExecServerTransportParams::websocket_url( + exec_server_url, + connect_timeout.unwrap_or(DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT), + ), self.local_runtime_paths.clone(), )); environment.start_connecting(); @@ -496,7 +502,10 @@ impl Environment { local_runtime_paths: Option, ) -> Self { Self::remote_with_transport( - ExecServerTransportParams::websocket_url(exec_server_url), + ExecServerTransportParams::websocket_url( + exec_server_url, + DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT, + ), local_runtime_paths, ) } @@ -980,7 +989,11 @@ mod tests { let manager = EnvironmentManager::without_environments(); manager - .upsert_environment("executor-a".to_string(), "ws://127.0.0.1:8765".to_string()) + .upsert_environment( + "executor-a".to_string(), + "ws://127.0.0.1:8765".to_string(), + /*connect_timeout*/ None, + ) .expect("remote environment"); let first = manager .get_environment("executor-a") @@ -990,7 +1003,11 @@ mod tests { assert_eq!(manager.default_environment_id(), None); manager - .upsert_environment("executor-a".to_string(), "ws://127.0.0.1:9876".to_string()) + .upsert_environment( + "executor-a".to_string(), + "ws://127.0.0.1:9876".to_string(), + /*connect_timeout*/ None, + ) .expect("updated remote environment"); let second = manager .get_environment("executor-a") @@ -1011,6 +1028,7 @@ mod tests { .upsert_environment( "executor-a".to_string(), format!("ws://{}", listener.local_addr().expect("listener address")), + /*connect_timeout*/ None, ) .expect("remote environment"); @@ -1066,6 +1084,7 @@ mod tests { "ws://{}", first_listener.local_addr().expect("first listener address") ), + /*connect_timeout*/ None, ) .expect("first remote environment"); let environment = manager @@ -1090,6 +1109,7 @@ mod tests { .local_addr() .expect("second listener address") ), + /*connect_timeout*/ None, ) .expect("replacement remote environment"); @@ -1107,7 +1127,11 @@ mod tests { let manager = EnvironmentManager::without_environments(); let err = manager - .upsert_environment("executor-a".to_string(), String::new()) + .upsert_environment( + "executor-a".to_string(), + String::new(), + /*connect_timeout*/ None, + ) .expect_err("empty URL should fail"); assert_eq!( diff --git a/codex-rs/exec-server/src/remote_file_system_path_uri_tests.rs b/codex-rs/exec-server/src/remote_file_system_path_uri_tests.rs index 29284e6ad..1205c6d7e 100644 --- a/codex-rs/exec-server/src/remote_file_system_path_uri_tests.rs +++ b/codex-rs/exec-server/src/remote_file_system_path_uri_tests.rs @@ -23,6 +23,7 @@ use tokio_tungstenite::accept_async; use tokio_tungstenite::tungstenite::Message; use super::*; +use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT; use crate::client_api::ExecServerTransportParams; use crate::protocol::FS_READ_FILE_METHOD; use crate::protocol::FsReadFileParams; @@ -36,7 +37,10 @@ async fn remote_file_system_sends_path_and_sandbox_cwd_uris_without_native_conve let (websocket_url, captured_params, server) = record_read_file_params(/*expected_requests*/ 2).await; let file_system = RemoteFileSystem::new(LazyRemoteExecServerClient::new( - ExecServerTransportParams::websocket_url(websocket_url), + ExecServerTransportParams::websocket_url( + websocket_url, + DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT, + ), )); let paths = vec![ PathUri::parse("file:///C:/Users/Alice/src/main.rs").expect("valid drive URI"),