From 4882aa428fada274fd8b6d4420ac0bc1aea7031c Mon Sep 17 00:00:00 2001 From: Adam Perry Date: Thu, 9 Jul 2026 10:33:34 +0000 Subject: [PATCH] perf(exec-server): return environment info during initialize (cherry picked from commit a27fea9943e5f9b7c83ebf79e3eb1ef8f0859c18) --- codex-rs/exec-server-protocol/src/protocol.rs | 20 +++ codex-rs/exec-server/README.md | 14 +- codex-rs/exec-server/src/client.rs | 163 +++++++++++++++++- .../src/remote_file_system_path_uri_tests.rs | 1 + codex-rs/exec-server/src/server/handler.rs | 5 +- codex-rs/exec-server/tests/http_client.rs | 1 + codex-rs/exec-server/tests/initialize.rs | 5 + 7 files changed, 200 insertions(+), 9 deletions(-) diff --git a/codex-rs/exec-server-protocol/src/protocol.rs b/codex-rs/exec-server-protocol/src/protocol.rs index a265aaed0c..78cd69ce1c 100644 --- a/codex-rs/exec-server-protocol/src/protocol.rs +++ b/codex-rs/exec-server-protocol/src/protocol.rs @@ -69,6 +69,9 @@ pub struct InitializeParams { #[serde(rename_all = "camelCase")] pub struct InitializeResponse { pub session_id: String, + /// Connection-stable execution environment metadata. Older servers omit this field. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub environment_info: Option, } /// Information about an execution/filesystem environment. @@ -570,6 +573,7 @@ mod tests { use super::ExecParams; use super::FsReadFileParams; use super::HttpRequestParams; + use super::InitializeResponse; use super::ProcessId; use super::ShellInfo; use codex_file_system::FileSystemSandboxContext; @@ -579,6 +583,22 @@ mod tests { use pretty_assertions::assert_eq; use std::collections::HashMap; + #[test] + fn initialize_response_accepts_legacy_response_without_environment_info() { + let response: InitializeResponse = serde_json::from_value(serde_json::json!({ + "sessionId": "legacy-session" + })) + .expect("legacy initialize response should deserialize"); + + assert_eq!( + response, + InitializeResponse { + session_id: "legacy-session".to_string(), + environment_info: None, + } + ); + } + #[test] fn exec_params_managed_network_context_round_trips_and_defaults_for_legacy_peers() { let cwd = diff --git a/codex-rs/exec-server/README.md b/codex-rs/exec-server/README.md index a250ba5886..41b764665d 100644 --- a/codex-rs/exec-server/README.md +++ b/codex-rs/exec-server/README.md @@ -138,9 +138,19 @@ Request params: Response: ```json -{} +{ + "sessionId": "08e2d8e7-8fb7-4495-a8e1-cd729dd68820", + "environmentInfo": { + "shell": { "name": "bash", "path": "/bin/bash" }, + "cwd": "file:///workspace" + } +} ``` +`environmentInfo` describes connection-stable properties of the execution +environment. Clients must fall back to `environment/info` when talking to an +older server that omits it. + ### `initialized` Handshake acknowledgement notification sent by the client after a successful @@ -412,7 +422,7 @@ Initialize: ```json {"id":1,"method":"initialize","params":{"clientName":"example-client"}} -{"id":1,"result":{}} +{"id":1,"result":{"sessionId":"08e2d8e7-8fb7-4495-a8e1-cd729dd68820","environmentInfo":{"shell":{"name":"bash","path":"/bin/bash"},"cwd":"file:///workspace"}}} {"method":"initialized","params":{}} ``` diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 73e49b2f54..26afaf8036 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -207,6 +207,7 @@ struct Inner { http_body_streams_write_lock: Mutex<()>, http_body_stream_next_id: AtomicU64, session_id: OnceLock, + environment_info: OnceCell, reconnect_strategy: Option, } @@ -575,6 +576,9 @@ impl ExecServerClient { response.session_id ))); } + if let Some(environment_info) = response.environment_info.as_ref() { + let _ = self.inner.environment_info.set(environment_info.clone()); + } rpc_client .notify(INITIALIZED_METHOD, &serde_json::json!({})) .await?; @@ -591,12 +595,18 @@ impl ExecServerClient { } pub async fn environment_info(&self) -> Result { - let rpc_client = self.rpc_client().await?; - map_rpc_call_result( - rpc_client - .call_with_timeout(ENVIRONMENT_INFO_METHOD, &(), ENVIRONMENT_INFO_TIMEOUT) - .await, - ) + self.inner + .environment_info + .get_or_try_init(|| async { + let rpc_client = self.inner.rpc_client().await?; + map_rpc_call_result( + rpc_client + .call_with_timeout(ENVIRONMENT_INFO_METHOD, &(), ENVIRONMENT_INFO_TIMEOUT) + .await, + ) + }) + .await + .cloned() } pub async fn read(&self, params: ReadParams) -> Result { @@ -853,6 +863,7 @@ impl ExecServerClient { http_body_streams_write_lock: Mutex::new(()), http_body_stream_next_id: AtomicU64::new(1), session_id, + environment_info: OnceCell::new(), reconnect_strategy, }); let client = Self { @@ -1431,12 +1442,14 @@ mod tests { use crate::client_api::StdioExecServerConnectArgs; use crate::connection::JsonRpcConnection; use crate::process::ExecProcessEvent; + use crate::protocol::ENVIRONMENT_INFO_METHOD; use crate::protocol::EXEC_CLOSED_METHOD; use crate::protocol::EXEC_EXITED_METHOD; use crate::protocol::EXEC_METHOD; use crate::protocol::EXEC_OUTPUT_DELTA_METHOD; use crate::protocol::EXEC_READ_METHOD; use crate::protocol::EXEC_WRITE_METHOD; + use crate::protocol::EnvironmentInfo; use crate::protocol::ExecClosedNotification; use crate::protocol::ExecExitedNotification; use crate::protocol::ExecOutputDeltaNotification; @@ -1448,6 +1461,7 @@ mod tests { use crate::protocol::InitializeResponse; use crate::protocol::ProcessOutputChunk; use crate::protocol::ReadResponse; + use crate::protocol::ShellInfo; use crate::protocol::WriteParams; use crate::protocol::WriteResponse; use crate::protocol::WriteStatus; @@ -1475,6 +1489,137 @@ mod tests { .expect("json-rpc line should write"); } + fn sample_environment_info() -> EnvironmentInfo { + EnvironmentInfo { + shell: ShellInfo { + name: "bash".to_string(), + path: "/bin/bash".to_string(), + }, + cwd: None, + } + } + + #[tokio::test(flavor = "current_thread")] + async fn environment_info_uses_initialize_response_without_follow_up_rpc() { + let (client_stdin, server_reader) = duplex(1 << 20); + let (mut server_writer, client_stdout) = duplex(1 << 20); + let expected = sample_environment_info(); + let server_expected = expected.clone(); + let server = tokio::spawn(async move { + let mut lines = BufReader::new(server_reader).lines(); + let initialize = read_jsonrpc_line(&mut lines).await; + let request = match initialize { + JSONRPCMessage::Request(request) if request.method == INITIALIZE_METHOD => request, + other => panic!("expected initialize request, got {other:?}"), + }; + write_jsonrpc_line( + &mut server_writer, + JSONRPCMessage::Response(JSONRPCResponse { + id: request.id, + result: serde_json::to_value(InitializeResponse { + session_id: "embedded-info".to_string(), + environment_info: Some(server_expected), + }) + .expect("initialize response should serialize"), + }), + ) + .await; + match read_jsonrpc_line(&mut lines).await { + JSONRPCMessage::Notification(notification) + if notification.method == INITIALIZED_METHOD => {} + other => panic!("expected initialized notification, got {other:?}"), + } + }); + + let client = ExecServerClient::connect( + JsonRpcConnection::from_stdio( + client_stdout, + client_stdin, + "embedded-info-client".to_string(), + ), + ExecServerClientConnectOptions::default(), + ) + .await + .expect("client should connect"); + server.await.expect("server task should finish"); + + assert_eq!( + client.environment_info().await.expect("first info"), + expected + ); + assert_eq!( + client.environment_info().await.expect("cached info"), + expected + ); + } + + #[tokio::test(flavor = "current_thread")] + async fn legacy_initialize_falls_back_to_environment_info_once() { + let (client_stdin, server_reader) = duplex(1 << 20); + let (mut server_writer, client_stdout) = duplex(1 << 20); + let expected = sample_environment_info(); + let server_expected = expected.clone(); + let server = tokio::spawn(async move { + let mut lines = BufReader::new(server_reader).lines(); + let initialize = read_jsonrpc_line(&mut lines).await; + let request = match initialize { + JSONRPCMessage::Request(request) if request.method == INITIALIZE_METHOD => request, + other => panic!("expected initialize request, got {other:?}"), + }; + write_jsonrpc_line( + &mut server_writer, + JSONRPCMessage::Response(JSONRPCResponse { + id: request.id, + result: serde_json::json!({ "sessionId": "legacy-info" }), + }), + ) + .await; + match read_jsonrpc_line(&mut lines).await { + JSONRPCMessage::Notification(notification) + if notification.method == INITIALIZED_METHOD => {} + other => panic!("expected initialized notification, got {other:?}"), + } + let info = read_jsonrpc_line(&mut lines).await; + let request = match info { + JSONRPCMessage::Request(request) if request.method == ENVIRONMENT_INFO_METHOD => { + request + } + other => panic!("expected one environment/info request, got {other:?}"), + }; + write_jsonrpc_line( + &mut server_writer, + JSONRPCMessage::Response(JSONRPCResponse { + id: request.id, + result: serde_json::to_value(server_expected) + .expect("environment info should serialize"), + }), + ) + .await; + }); + + let client = ExecServerClient::connect( + JsonRpcConnection::from_stdio( + client_stdout, + client_stdin, + "legacy-info-client".to_string(), + ), + ExecServerClientConnectOptions::default(), + ) + .await + .expect("client should connect"); + let (first, second) = tokio::join!(client.environment_info(), client.environment_info()); + assert_eq!(first.expect("first fallback info"), expected); + assert_eq!(second.expect("shared fallback info"), expected); + server.await.expect("server task should finish"); + assert_eq!( + client + .environment_info() + .await + .expect("cached fallback info"), + expected + ); + } + #[tokio::test(flavor = "current_thread")] async fn process_start_propagates_caller_trace_context_across_background_task() { let (client_stdin, server_reader) = duplex(1 << 20); @@ -1492,6 +1637,7 @@ mod tests { id: initialize.id, result: serde_json::to_value(InitializeResponse { session_id: "trace-test".to_string(), + environment_info: None, }) .expect("initialize response should serialize"), }), @@ -1647,6 +1793,7 @@ mod tests { id: request.id, result: serde_json::to_value(InitializeResponse { session_id: session_id.to_string(), + environment_info: None, }) .expect("initialize response should serialize"), }), @@ -1867,6 +2014,7 @@ mod tests { id: request.id, result: serde_json::to_value(InitializeResponse { session_id: "session-1".to_string(), + environment_info: None, }) .expect("initialize response should serialize"), }), @@ -2012,6 +2160,7 @@ mod tests { id: request.id, result: serde_json::to_value(InitializeResponse { session_id: "session-1".to_string(), + environment_info: None, }) .expect("initialize response should serialize"), }), @@ -2292,6 +2441,7 @@ mod tests { id: request.id, result: serde_json::to_value(InitializeResponse { session_id: "session-1".to_string(), + environment_info: None, }) .expect("initialize response should serialize"), }), @@ -2522,6 +2672,7 @@ mod tests { id: request.id, result: serde_json::to_value(InitializeResponse { session_id: "session-1".to_string(), + environment_info: None, }) .expect("initialize response should serialize"), }), 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 562f5d5643..f3cf0ef3bc 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 @@ -153,6 +153,7 @@ async fn complete_websocket_initialize(websocket: &mut WebSocketStream Result<(), String> { diff --git a/codex-rs/exec-server/tests/http_client.rs b/codex-rs/exec-server/tests/http_client.rs index b3f0182c63..a927488feb 100644 --- a/codex-rs/exec-server/tests/http_client.rs +++ b/codex-rs/exec-server/tests/http_client.rs @@ -1038,6 +1038,7 @@ impl JsonRpcPeer { request.id, InitializeResponse { session_id: "session-1".to_string(), + environment_info: None, }, ) .await?; diff --git a/codex-rs/exec-server/tests/initialize.rs b/codex-rs/exec-server/tests/initialize.rs index daf4306db6..249108ad2f 100644 --- a/codex-rs/exec-server/tests/initialize.rs +++ b/codex-rs/exec-server/tests/initialize.rs @@ -2,6 +2,7 @@ mod common; +use codex_exec_server::EnvironmentInfo; use codex_exec_server::InitializeParams; use codex_exec_server::InitializeResponse; use codex_exec_server_protocol::JSONRPCMessage; @@ -30,6 +31,10 @@ async fn exec_server_accepts_initialize() -> anyhow::Result<()> { assert_eq!(id, initialize_id); let initialize_response: InitializeResponse = serde_json::from_value(result)?; Uuid::parse_str(&initialize_response.session_id)?; + assert_eq!( + initialize_response.environment_info, + Some(EnvironmentInfo::local()) + ); server.shutdown().await?; Ok(())