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 1a29dbd954..f88a9c12a8 100644 --- a/codex-rs/app-server/src/request_processors/environment_processor.rs +++ b/codex-rs/app-server/src/request_processors/environment_processor.rs @@ -36,7 +36,7 @@ impl EnvironmentRequestProcessor { .environment_manager .get_environment(&environment_id) .ok_or_else(|| invalid_request(format!("unknown environment id `{environment_id}`")))?; - let info = environment.info().await.map_err(|err| { + let info = environment.force_info().await.map_err(|err| { internal_error(format!( "failed to get info for environment `{environment_id}`: {err}" )) diff --git a/codex-rs/app-server/tests/suite/v2/environment_info.rs b/codex-rs/app-server/tests/suite/v2/environment_info.rs index 765144968b..4acc494493 100644 --- a/codex-rs/app-server/tests/suite/v2/environment_info.rs +++ b/codex-rs/app-server/tests/suite/v2/environment_info.rs @@ -11,18 +11,91 @@ use codex_app_server_protocol::JSONRPCErrorError; use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::RequestId; use codex_utils_path_uri::PathUri; +use futures::SinkExt; use pretty_assertions::assert_eq; use serde_json::json; use tempfile::TempDir; use tokio::net::TcpListener; use tokio::time::timeout; +use tokio_tungstenite::accept_async; +use tokio_tungstenite::tungstenite::Message; use super::exec_server_test_support::accept_exec_server_environment; +use super::exec_server_test_support::read_exec_server_json; const RPC_TIMEOUT: Duration = Duration::from_secs(10); const INVALID_REQUEST_ERROR_CODE: i64 = -32600; const INTERNAL_ERROR_CODE: i64 = -32603; +#[tokio::test] +async fn environment_info_probes_executor_even_when_metadata_is_cached() -> Result<()> { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let exec_server_url = format!("ws://{}", listener.local_addr()?); + let exec_server = tokio::spawn(async move { + let (stream, _) = listener.accept().await?; + let mut websocket = accept_async(stream).await?; + let initialize = read_exec_server_json(&mut websocket).await?; + assert_eq!(initialize["method"], "initialize"); + let response = json!({ + "id": initialize["id"], + "result": { + "sessionId": "test-session", + "environmentInfo": {"shell": {"name": "bash", "path": "/bin/bash"}}, + }, + }); + websocket + .send(Message::Text(response.to_string().into())) + .await?; + let initialized = read_exec_server_json(&mut websocket).await?; + assert_eq!(initialized["method"], "initialized"); + + for _ in 0..2 { + let request = read_exec_server_json(&mut websocket).await?; + assert_eq!(request["method"], "environment/info"); + let response = json!({ + "id": request["id"], + "error": {"code": INTERNAL_ERROR_CODE, "message": "executor unavailable"}, + }); + websocket + .send(Message::Text(response.to_string().into())) + .await?; + } + Ok::<_, anyhow::Error>(()) + }); + let mut app_server = TestAppServer::builder().without_auto_env().build().await?; + app_server.initialize().await?; + add_environment( + &mut app_server, + &exec_server_url, + /*connect_timeout_ms*/ None, + ) + .await?; + + for _ in 0..2 { + let request_id = app_server + .send_raw_request( + "environment/info", + Some(json!({"environmentId": "remote-a"})), + ) + .await?; + let error = timeout( + RPC_TIMEOUT, + app_server.read_stream_until_error_message(RequestId::Integer(request_id)), + ) + .await??; + assert_eq!(error, JSONRPCError { + id: RequestId::Integer(request_id), + error: JSONRPCErrorError { + code: INTERNAL_ERROR_CODE, + message: "failed to get info for environment `remote-a`: exec-server rejected request (-32603): executor unavailable".to_string(), + data: None, + }, + }); + } + timeout(RPC_TIMEOUT, exec_server).await???; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn environment_info_returns_remote_environment_info() -> Result<()> { let listener = TcpListener::bind("127.0.0.1:0").await?; diff --git a/codex-rs/app-server/tests/suite/v2/executor_mcp.rs b/codex-rs/app-server/tests/suite/v2/executor_mcp.rs index 5348bf5b5a..2fddde479a 100644 --- a/codex-rs/app-server/tests/suite/v2/executor_mcp.rs +++ b/codex-rs/app-server/tests/suite/v2/executor_mcp.rs @@ -332,6 +332,11 @@ async fn legacy_executor_skips_required_browser_and_keeps_host_owned_mcp() -> Re let mut message = message?; if let Message::Text(text) = &message { let mut response: Value = serde_json::from_str(text)?; + if let Some(result) = response.get_mut("result").and_then(Value::as_object_mut) + { + result.remove("environmentInfo"); + message = Message::Text(response.to_string().into()); + } if let Some(capabilities) = response .pointer_mut("/result/capabilities") .and_then(Value::as_object_mut) diff --git a/codex-rs/exec-server-protocol/src/protocol.rs b/codex-rs/exec-server-protocol/src/protocol.rs index c46978f966..07d7b586c4 100644 --- a/codex-rs/exec-server-protocol/src/protocol.rs +++ b/codex-rs/exec-server-protocol/src/protocol.rs @@ -83,6 +83,10 @@ pub struct InitializeParams { #[serde(rename_all = "camelCase")] pub struct InitializeResponse { pub session_id: String, + /// Executor metadata at initialization, with the same shape as `environment/info`. + // TODO: Make this required once all supported exec-server versions return environmentInfo. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub environment_info: Option, } /// Information about an execution/filesystem environment. diff --git a/codex-rs/exec-server/README.md b/codex-rs/exec-server/README.md index 7f5d1247e5..c5fd2c2d59 100644 --- a/codex-rs/exec-server/README.md +++ b/codex-rs/exec-server/README.md @@ -153,9 +153,22 @@ Request params: Response: ```json -{} +{ + "sessionId": "00000000-0000-4000-8000-000000000001", + "environmentInfo": { + "shell": { "name": "bash", "path": "/bin/bash" }, + "cwd": "file:///workspace" + } +} ``` +`environmentInfo` contains the same executor metadata returned by +`environment/info`, so clients can use it without a second request. + +Rust clients cache this metadata for the client's lifetime, including session +resumption. If initialization omits it, the first metadata request fetches and +caches `environment/info`. + ### `initialized` Handshake acknowledgement notification sent by the client after a successful @@ -429,7 +442,7 @@ Initialize: ```json {"id":1,"method":"initialize","params":{"clientName":"example-client"}} -{"id":1,"result":{}} +{"id":1,"result":{"sessionId":"00000000-0000-4000-8000-000000000001","environmentInfo":{"shell":{"name":"bash","path":"/bin/bash"},"cwd":"file:///tmp"}}} {"method":"initialized","params":{}} ``` diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 6f794dd171..63ed561c4a 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -264,6 +264,8 @@ struct Inner { // Keep admission shared while recovered transports finish older requests. rpc_inbound_request_slots: Arc, session_id: OnceLock, + /// Caches metadata from initialization or the first successful info request for this client's lifetime. + environment_info: OnceCell, reconnect_strategy: Option, } @@ -765,7 +767,18 @@ impl ExecServerClient { self.call(EXEC_METHOD, ¶ms).await } + /// Returns cached executor metadata, fetching it lazily if initialization omitted it. pub async fn environment_info(&self) -> Result { + self.inner + .environment_info + .get_or_try_init(|| self.force_environment_info()) + .await + .cloned() + } + + /// Fetches executor metadata over RPC without reading or updating the cache. + // TODO: Remove after app-server migrates off this call. + pub async fn force_environment_info(&self) -> Result { let rpc_client = self.rpc_client().await?; map_rpc_call_result( rpc_client @@ -1117,6 +1130,7 @@ impl ExecServerClient { http_body_stream_next_id: AtomicU64::new(1), rpc_inbound_request_slots: Arc::new(Semaphore::new(MAX_IN_FLIGHT_SERVER_CALLS)), session_id, + environment_info: OnceCell::new(), reconnect_strategy, }); let client = Self { @@ -1127,7 +1141,13 @@ impl ExecServerClient { // before initialize returns. Drain them immediately so a burst cannot // fill the bounded event channel and block the initialize response. client.spawn_rpc_reader(&rpc_client, events_rx); - client.initialize_rpc(&rpc_client, options).await?; + let initialize_response = client.initialize_rpc(&rpc_client, options).await?; + if let Some(info) = initialize_response.environment_info { + assert!( + client.inner.environment_info.set(info).is_ok(), + "new client metadata cache must be empty" + ); + } Ok(client) } @@ -1740,6 +1760,7 @@ mod tests { 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; @@ -1796,6 +1817,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"), }), @@ -1942,6 +1964,21 @@ mod tests { websocket: &mut WebSocketStream, session_id: &str, expected_resume_session_id: Option<&str>, + ) { + complete_websocket_initialize_with_environment_info( + websocket, + session_id, + expected_resume_session_id, + /*environment_info*/ None, + ) + .await; + } + + async fn complete_websocket_initialize_with_environment_info( + websocket: &mut WebSocketStream, + session_id: &str, + expected_resume_session_id: Option<&str>, + environment_info: Option, ) { let initialize = read_jsonrpc_websocket(websocket).await; let request = match initialize { @@ -1961,6 +1998,7 @@ mod tests { id: request.id, result: serde_json::to_value(InitializeResponse { session_id: session_id.to_string(), + environment_info, }) .expect("initialize response should serialize"), }), @@ -2184,6 +2222,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"), }), @@ -2329,6 +2368,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"), }), @@ -2397,6 +2437,61 @@ mod tests { server.await.expect("server task should finish"); } + #[test_case::test_case(Some(EnvironmentInfo::local()); "from_initialize")] + #[test_case::test_case(None; "legacy_server")] + #[tokio::test] + async fn environment_info_is_cached( + initial_environment_info: Option, + ) -> anyhow::Result<()> { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let websocket_url = format!("ws://{}", listener.local_addr()?); + let expected_info = initial_environment_info + .clone() + .unwrap_or_else(EnvironmentInfo::local); + let server_info = expected_info.clone(); + let server = tokio::spawn(async move { + let mut websocket = accept_websocket(&listener).await; + complete_websocket_initialize_with_environment_info( + &mut websocket, + "session-1", + /*expected_resume_session_id*/ None, + initial_environment_info.clone(), + ) + .await; + if initial_environment_info.is_none() { + let JSONRPCMessage::Request(request) = read_jsonrpc_websocket(&mut websocket).await + else { + panic!("expected environment info request"); + }; + assert_eq!(request.method, "environment/info"); + write_jsonrpc_websocket( + &mut websocket, + JSONRPCMessage::Response(JSONRPCResponse { + id: request.id, + result: serde_json::to_value(server_info) + .expect("environment info should serialize"), + }), + ) + .await; + } + }); + let client = ExecServerClient::connect_websocket(RemoteExecServerConnectArgs { + websocket_url, + client_name: "metadata-test".to_string(), + connect_timeout: Duration::from_secs(1), + initialize_timeout: Duration::from_secs(1), + resume_session_id: None, + http_client_factory: HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + }) + .await?; + + assert_eq!(client.environment_info().await?, expected_info); + server.await?; + // The server is gone, so a cloned client must use the shared cache. + assert_eq!(client.clone().environment_info().await?, expected_info); + Ok(()) + } + #[tokio::test] async fn remote_websocket_client_resumes_session() { let listener = TcpListener::bind("127.0.0.1:0") @@ -2615,6 +2710,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"), }), @@ -2932,6 +3028,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/environment.rs b/codex-rs/exec-server/src/environment.rs index e1ea7674da..e3f69de24f 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -889,6 +889,7 @@ impl Environment { } /// Returns environment information from the selected execution/filesystem environment. + /// Remote metadata is cached for the current client's lifetime. #[tracing::instrument( name = "exec_server.environment.info", skip_all, @@ -901,6 +902,20 @@ impl Environment { } } + /// Fetches uncached metadata, connecting or waiting for recovery as needed. + // TODO: Remove after app-server migrates off of force_environment_info. + #[tracing::instrument( + name = "exec_server.environment.force_info", + skip_all, + fields(remote = self.is_remote()) + )] + pub async fn force_info(&self) -> Result { + match &self.remote_client { + Some(client) => client.get().await?.force_environment_info().await, + None => Ok(EnvironmentInfo::local()), + } + } + /// Reads selected executor-local configuration fields for this environment. pub async fn read_environment_config( &self, 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 adba21c78c..4956f4bfdb 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 @@ -614,6 +614,7 @@ async fn complete_websocket_initialize(websocket: &mut WebSocketStream Result<(), String> { diff --git a/codex-rs/exec-server/tests/accepted_websocket.rs b/codex-rs/exec-server/tests/accepted_websocket.rs index 217ad8e74c..5334dfd7fd 100644 --- a/codex-rs/exec-server/tests/accepted_websocket.rs +++ b/codex-rs/exec-server/tests/accepted_websocket.rs @@ -8,6 +8,7 @@ use axum::extract::State; use axum::extract::WebSocketUpgrade; use axum::response::IntoResponse; use axum::routing::any; +use codex_exec_server::EnvironmentInfo; use codex_exec_server::EnvironmentManager; use codex_exec_server::EnvironmentObservedStatus; use codex_exec_server::EnvironmentStatus; @@ -77,6 +78,25 @@ async fn accepted_websocket_rejects_initial_resume_session_id() -> Result<()> { Ok(()) } +#[tokio::test] +async fn accepted_websocket_environment_info_uses_initialization_metadata() -> Result<()> { + let (websocket_url, mut accepted_sockets, server_task) = start_acceptor().await?; + let (_socket, manager) = + connect_executor(&websocket_url, &mut accepted_sockets, "session-1").await?; + let environment = manager + .default_environment() + .context("accepted environment")?; + + assert_eq!( + timeout(TEST_TIMEOUT, environment.info()).await??, + EnvironmentInfo::local() + ); + + server_task.abort(); + let _ = server_task.await; + Ok(()) +} + #[tokio::test] async fn accepted_websocket_environment_is_ready_immediately() -> Result<()> { let (websocket_url, mut accepted_sockets, server_task) = start_acceptor().await?; @@ -430,6 +450,7 @@ where id, result: serde_json::to_value(InitializeResponse { session_id: session_id.to_string(), + environment_info: Some(EnvironmentInfo::local()), })?, }), ) diff --git a/codex-rs/exec-server/tests/http_client.rs b/codex-rs/exec-server/tests/http_client.rs index b10cda7e7d..ebe8471dee 100644 --- a/codex-rs/exec-server/tests/http_client.rs +++ b/codex-rs/exec-server/tests/http_client.rs @@ -1440,6 +1440,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 27bd9a7eb1..64ca63f551 100644 --- a/codex-rs/exec-server/tests/initialize.rs +++ b/codex-rs/exec-server/tests/initialize.rs @@ -1,5 +1,6 @@ mod common; +use codex_exec_server::EnvironmentInfo; use codex_exec_server::InitializeParams; use codex_exec_server::InitializeResponse; use codex_exec_server_protocol::JSONRPCError; @@ -31,6 +32,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(()) diff --git a/codex-rs/exec-server/tests/relay.rs b/codex-rs/exec-server/tests/relay.rs index e54dd2e9f0..3201538366 100644 --- a/codex-rs/exec-server/tests/relay.rs +++ b/codex-rs/exec-server/tests/relay.rs @@ -277,9 +277,11 @@ async fn pending_noise_environment_connects_and_reconnects_after_ready_report() second_reconnected_websocket, Arc::new(Mutex::new(Vec::new())), )); - let recovered_info = timeout(TEST_TIMEOUT, environment.info()) - .await - .context("pending Noise environment should reconnect")??; + assert_eq!( + next_connection_state(&mut connection_state).await?, + EnvironmentConnectionState::Connected + ); + let recovered_info = environment.info().await?; assert_eq!(recovered_info, initial_info); assert_eq!( @@ -287,10 +289,6 @@ async fn pending_noise_environment_connects_and_reconnects_after_ready_report() selected_capability_roots ); assert_eq!(provider.calls(), 2); - assert_eq!( - next_connection_state(&mut connection_state).await?, - EnvironmentConnectionState::Connected - ); registry.verify().await; second_relay.abort(); diff --git a/codex-rs/rmcp-client/tests/streamable_http_remote.rs b/codex-rs/rmcp-client/tests/streamable_http_remote.rs index 070b3c815d..7df871976b 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_remote.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_remote.rs @@ -124,7 +124,7 @@ async fn stalled_handshake_unblocks_remote_executor( tokio::time::timeout( Duration::from_secs(2), - exec_server.client.environment_info(), + exec_server.client.environment_status(), ) .await .context("a timed-out handshake must not leave the serial executor blocked")??;