Cache executor metadata from initialization (#40343)

## What changed

- Include `environmentInfo` in the exec-server `initialize` response and seed the Rust client's metadata cache from it.
- Fall back to one lazy `environment/info` request when connecting to older servers that omit initialization metadata, then share the cached result across client clones and session recovery.
- Keep app-server `environment/info` requests uncached so they continue to probe the executor.

## Testing

- Cover initialization metadata, the legacy-server fallback, shared caching, accepted WebSocket environments, and repeated uncached app-server probes.

GitOrigin-RevId: 380132d9db183998d73e8c6b0954e4d70c7e8406
This commit is contained in:
cgst-oai authored and copyberry committed 2026-08-24 06:43:07 +00:00
1 parent 068c49f075
commit 2df6705423
14 files changed
+249 -13

No files matched your search

@@ -36,7 +36,7 @@ impl EnvironmentRequestProcessor {
.environment_manager .environment_manager
.get_environment(&environment_id) .get_environment(&environment_id)
.ok_or_else(|| invalid_request(format!("unknown environment id `{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!( internal_error(format!(
"failed to get info for environment `{environment_id}`: {err}" "failed to get info for environment `{environment_id}`: {err}"
)) ))
@@ -11,18 +11,91 @@ use codex_app_server_protocol::JSONRPCErrorError;
use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::JSONRPCResponse;
use codex_app_server_protocol::RequestId; use codex_app_server_protocol::RequestId;
use codex_utils_path_uri::PathUri; use codex_utils_path_uri::PathUri;
use futures::SinkExt;
use pretty_assertions::assert_eq; use pretty_assertions::assert_eq;
use serde_json::json; use serde_json::json;
use tempfile::TempDir; use tempfile::TempDir;
use tokio::net::TcpListener; use tokio::net::TcpListener;
use tokio::time::timeout; 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::accept_exec_server_environment;
use super::exec_server_test_support::read_exec_server_json;
const RPC_TIMEOUT: Duration = Duration::from_secs(10); const RPC_TIMEOUT: Duration = Duration::from_secs(10);
const INVALID_REQUEST_ERROR_CODE: i64 = -32600; const INVALID_REQUEST_ERROR_CODE: i64 = -32600;
const INTERNAL_ERROR_CODE: i64 = -32603; 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)] #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn environment_info_returns_remote_environment_info() -> Result<()> { async fn environment_info_returns_remote_environment_info() -> Result<()> {
let listener = TcpListener::bind("127.0.0.1:0").await?; let listener = TcpListener::bind("127.0.0.1:0").await?;
@@ -332,6 +332,11 @@ async fn legacy_executor_skips_required_browser_and_keeps_host_owned_mcp() -> Re
let mut message = message?; let mut message = message?;
if let Message::Text(text) = &message { if let Message::Text(text) = &message {
let mut response: Value = serde_json::from_str(text)?; 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 if let Some(capabilities) = response
.pointer_mut("/result/capabilities") .pointer_mut("/result/capabilities")
.and_then(Value::as_object_mut) .and_then(Value::as_object_mut)
@@ -83,6 +83,10 @@ pub struct InitializeParams {
#[serde(rename_all = "camelCase")] #[serde(rename_all = "camelCase")]
pub struct InitializeResponse { pub struct InitializeResponse {
pub session_id: String, 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<EnvironmentInfo>,
} }
/// Information about an execution/filesystem environment. /// Information about an execution/filesystem environment.
+15 -2
View File
@@ -153,9 +153,22 @@ Request params:
Response: Response:
```json ```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` ### `initialized`
Handshake acknowledgement notification sent by the client after a successful Handshake acknowledgement notification sent by the client after a successful
@@ -429,7 +442,7 @@ Initialize:
```json ```json
{"id":1,"method":"initialize","params":{"clientName":"example-client"}} {"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":{}} {"method":"initialized","params":{}}
``` ```
+98 -1
View File
@@ -264,6 +264,8 @@ struct Inner {
// Keep admission shared while recovered transports finish older requests. // Keep admission shared while recovered transports finish older requests.
rpc_inbound_request_slots: Arc<Semaphore>, rpc_inbound_request_slots: Arc<Semaphore>,
session_id: OnceLock<String>, session_id: OnceLock<String>,
/// Caches metadata from initialization or the first successful info request for this client's lifetime.
environment_info: OnceCell<EnvironmentInfo>,
reconnect_strategy: Option<ExecServerReconnectStrategy>, reconnect_strategy: Option<ExecServerReconnectStrategy>,
} }
@@ -765,7 +767,18 @@ impl ExecServerClient {
self.call(EXEC_METHOD, &params).await self.call(EXEC_METHOD, &params).await
} }
/// Returns cached executor metadata, fetching it lazily if initialization omitted it.
pub async fn environment_info(&self) -> Result<EnvironmentInfo, ExecServerError> { pub async fn environment_info(&self) -> Result<EnvironmentInfo, ExecServerError> {
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<EnvironmentInfo, ExecServerError> {
let rpc_client = self.rpc_client().await?; let rpc_client = self.rpc_client().await?;
map_rpc_call_result( map_rpc_call_result(
rpc_client rpc_client
@@ -1117,6 +1130,7 @@ impl ExecServerClient {
http_body_stream_next_id: AtomicU64::new(1), http_body_stream_next_id: AtomicU64::new(1),
rpc_inbound_request_slots: Arc::new(Semaphore::new(MAX_IN_FLIGHT_SERVER_CALLS)), rpc_inbound_request_slots: Arc::new(Semaphore::new(MAX_IN_FLIGHT_SERVER_CALLS)),
session_id, session_id,
environment_info: OnceCell::new(),
reconnect_strategy, reconnect_strategy,
}); });
let client = Self { let client = Self {
@@ -1127,7 +1141,13 @@ impl ExecServerClient {
// before initialize returns. Drain them immediately so a burst cannot // before initialize returns. Drain them immediately so a burst cannot
// fill the bounded event channel and block the initialize response. // fill the bounded event channel and block the initialize response.
client.spawn_rpc_reader(&rpc_client, events_rx); 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) Ok(client)
} }
@@ -1740,6 +1760,7 @@ mod tests {
use crate::protocol::EXEC_OUTPUT_DELTA_METHOD; use crate::protocol::EXEC_OUTPUT_DELTA_METHOD;
use crate::protocol::EXEC_READ_METHOD; use crate::protocol::EXEC_READ_METHOD;
use crate::protocol::EXEC_WRITE_METHOD; use crate::protocol::EXEC_WRITE_METHOD;
use crate::protocol::EnvironmentInfo;
use crate::protocol::ExecClosedNotification; use crate::protocol::ExecClosedNotification;
use crate::protocol::ExecExitedNotification; use crate::protocol::ExecExitedNotification;
use crate::protocol::ExecOutputDeltaNotification; use crate::protocol::ExecOutputDeltaNotification;
@@ -1796,6 +1817,7 @@ mod tests {
id: initialize.id, id: initialize.id,
result: serde_json::to_value(InitializeResponse { result: serde_json::to_value(InitializeResponse {
session_id: "trace-test".to_string(), session_id: "trace-test".to_string(),
environment_info: None,
}) })
.expect("initialize response should serialize"), .expect("initialize response should serialize"),
}), }),
@@ -1942,6 +1964,21 @@ mod tests {
websocket: &mut WebSocketStream<TcpStream>, websocket: &mut WebSocketStream<TcpStream>,
session_id: &str, session_id: &str,
expected_resume_session_id: Option<&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<TcpStream>,
session_id: &str,
expected_resume_session_id: Option<&str>,
environment_info: Option<EnvironmentInfo>,
) { ) {
let initialize = read_jsonrpc_websocket(websocket).await; let initialize = read_jsonrpc_websocket(websocket).await;
let request = match initialize { let request = match initialize {
@@ -1961,6 +1998,7 @@ mod tests {
id: request.id, id: request.id,
result: serde_json::to_value(InitializeResponse { result: serde_json::to_value(InitializeResponse {
session_id: session_id.to_string(), session_id: session_id.to_string(),
environment_info,
}) })
.expect("initialize response should serialize"), .expect("initialize response should serialize"),
}), }),
@@ -2184,6 +2222,7 @@ mod tests {
id: request.id, id: request.id,
result: serde_json::to_value(InitializeResponse { result: serde_json::to_value(InitializeResponse {
session_id: "session-1".to_string(), session_id: "session-1".to_string(),
environment_info: None,
}) })
.expect("initialize response should serialize"), .expect("initialize response should serialize"),
}), }),
@@ -2329,6 +2368,7 @@ mod tests {
id: request.id, id: request.id,
result: serde_json::to_value(InitializeResponse { result: serde_json::to_value(InitializeResponse {
session_id: "session-1".to_string(), session_id: "session-1".to_string(),
environment_info: None,
}) })
.expect("initialize response should serialize"), .expect("initialize response should serialize"),
}), }),
@@ -2397,6 +2437,61 @@ mod tests {
server.await.expect("server task should finish"); 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<EnvironmentInfo>,
) -> 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] #[tokio::test]
async fn remote_websocket_client_resumes_session() { async fn remote_websocket_client_resumes_session() {
let listener = TcpListener::bind("127.0.0.1:0") let listener = TcpListener::bind("127.0.0.1:0")
@@ -2615,6 +2710,7 @@ mod tests {
id: request.id, id: request.id,
result: serde_json::to_value(InitializeResponse { result: serde_json::to_value(InitializeResponse {
session_id: "session-1".to_string(), session_id: "session-1".to_string(),
environment_info: None,
}) })
.expect("initialize response should serialize"), .expect("initialize response should serialize"),
}), }),
@@ -2932,6 +3028,7 @@ mod tests {
id: request.id, id: request.id,
result: serde_json::to_value(InitializeResponse { result: serde_json::to_value(InitializeResponse {
session_id: "session-1".to_string(), session_id: "session-1".to_string(),
environment_info: None,
}) })
.expect("initialize response should serialize"), .expect("initialize response should serialize"),
}), }),
+15
View File
@@ -889,6 +889,7 @@ impl Environment {
} }
/// Returns environment information from the selected execution/filesystem environment. /// Returns environment information from the selected execution/filesystem environment.
/// Remote metadata is cached for the current client's lifetime.
#[tracing::instrument( #[tracing::instrument(
name = "exec_server.environment.info", name = "exec_server.environment.info",
skip_all, 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<EnvironmentInfo, ExecServerError> {
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. /// Reads selected executor-local configuration fields for this environment.
pub async fn read_environment_config( pub async fn read_environment_config(
&self, &self,
@@ -614,6 +614,7 @@ async fn complete_websocket_initialize(websocket: &mut WebSocketStream<TcpStream
id: request.id, id: request.id,
result: serde_json::to_value(InitializeResponse { result: serde_json::to_value(InitializeResponse {
session_id: "session-1".to_string(), session_id: "session-1".to_string(),
environment_info: None,
}) })
.expect("initialize response should serialize"), .expect("initialize response should serialize"),
}), }),
+4 -1
View File
@@ -156,7 +156,10 @@ impl ExecServerHandler {
.session .session
.lock() .lock()
.unwrap_or_else(std::sync::PoisonError::into_inner) = Some(session); .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(session);
Ok(InitializeResponse { session_id }) Ok(InitializeResponse {
session_id,
environment_info: Some(EnvironmentInfo::local()),
})
} }
pub(crate) fn initialized(&self) -> Result<(), String> { pub(crate) fn initialized(&self) -> Result<(), String> {
@@ -8,6 +8,7 @@ use axum::extract::State;
use axum::extract::WebSocketUpgrade; use axum::extract::WebSocketUpgrade;
use axum::response::IntoResponse; use axum::response::IntoResponse;
use axum::routing::any; use axum::routing::any;
use codex_exec_server::EnvironmentInfo;
use codex_exec_server::EnvironmentManager; use codex_exec_server::EnvironmentManager;
use codex_exec_server::EnvironmentObservedStatus; use codex_exec_server::EnvironmentObservedStatus;
use codex_exec_server::EnvironmentStatus; use codex_exec_server::EnvironmentStatus;
@@ -77,6 +78,25 @@ async fn accepted_websocket_rejects_initial_resume_session_id() -> Result<()> {
Ok(()) 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] #[tokio::test]
async fn accepted_websocket_environment_is_ready_immediately() -> Result<()> { async fn accepted_websocket_environment_is_ready_immediately() -> Result<()> {
let (websocket_url, mut accepted_sockets, server_task) = start_acceptor().await?; let (websocket_url, mut accepted_sockets, server_task) = start_acceptor().await?;
@@ -430,6 +450,7 @@ where
id, id,
result: serde_json::to_value(InitializeResponse { result: serde_json::to_value(InitializeResponse {
session_id: session_id.to_string(), session_id: session_id.to_string(),
environment_info: Some(EnvironmentInfo::local()),
})?, })?,
}), }),
) )
@@ -1440,6 +1440,7 @@ impl JsonRpcPeer {
request.id, request.id,
InitializeResponse { InitializeResponse {
session_id: "session-1".to_string(), session_id: "session-1".to_string(),
environment_info: None,
}, },
) )
.await?; .await?;
+5
View File
@@ -1,5 +1,6 @@
mod common; mod common;
use codex_exec_server::EnvironmentInfo;
use codex_exec_server::InitializeParams; use codex_exec_server::InitializeParams;
use codex_exec_server::InitializeResponse; use codex_exec_server::InitializeResponse;
use codex_exec_server_protocol::JSONRPCError; use codex_exec_server_protocol::JSONRPCError;
@@ -31,6 +32,10 @@ async fn exec_server_accepts_initialize() -> anyhow::Result<()> {
assert_eq!(id, initialize_id); assert_eq!(id, initialize_id);
let initialize_response: InitializeResponse = serde_json::from_value(result)?; let initialize_response: InitializeResponse = serde_json::from_value(result)?;
Uuid::parse_str(&initialize_response.session_id)?; Uuid::parse_str(&initialize_response.session_id)?;
assert_eq!(
initialize_response.environment_info,
Some(EnvironmentInfo::local())
);
server.shutdown().await?; server.shutdown().await?;
Ok(()) Ok(())
+5 -7
View File
@@ -277,9 +277,11 @@ async fn pending_noise_environment_connects_and_reconnects_after_ready_report()
second_reconnected_websocket, second_reconnected_websocket,
Arc::new(Mutex::new(Vec::new())), Arc::new(Mutex::new(Vec::new())),
)); ));
let recovered_info = timeout(TEST_TIMEOUT, environment.info()) assert_eq!(
.await next_connection_state(&mut connection_state).await?,
.context("pending Noise environment should reconnect")??; EnvironmentConnectionState::Connected
);
let recovered_info = environment.info().await?;
assert_eq!(recovered_info, initial_info); assert_eq!(recovered_info, initial_info);
assert_eq!( assert_eq!(
@@ -287,10 +289,6 @@ async fn pending_noise_environment_connects_and_reconnects_after_ready_report()
selected_capability_roots selected_capability_roots
); );
assert_eq!(provider.calls(), 2); assert_eq!(provider.calls(), 2);
assert_eq!(
next_connection_state(&mut connection_state).await?,
EnvironmentConnectionState::Connected
);
registry.verify().await; registry.verify().await;
second_relay.abort(); second_relay.abort();
@@ -124,7 +124,7 @@ async fn stalled_handshake_unblocks_remote_executor(
tokio::time::timeout( tokio::time::timeout(
Duration::from_secs(2), Duration::from_secs(2),
exec_server.client.environment_info(), exec_server.client.environment_status(),
) )
.await .await
.context("a timed-out handshake must not leave the serial executor blocked")??; .context("a timed-out handshake must not leave the serial executor blocked")??;