mirror of
https://github.com/openai/codex.git
synced 2026-10-19 17:28:46 +00:00
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:
@@ -264,6 +264,8 @@ struct Inner {
|
||||
// Keep admission shared while recovered transports finish older requests.
|
||||
rpc_inbound_request_slots: Arc<Semaphore>,
|
||||
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>,
|
||||
}
|
||||
|
||||
@@ -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<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?;
|
||||
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<TcpStream>,
|
||||
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<TcpStream>,
|
||||
session_id: &str,
|
||||
expected_resume_session_id: Option<&str>,
|
||||
environment_info: Option<EnvironmentInfo>,
|
||||
) {
|
||||
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<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]
|
||||
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"),
|
||||
}),
|
||||
|
||||
@@ -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<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.
|
||||
pub async fn read_environment_config(
|
||||
&self,
|
||||
|
||||
@@ -614,6 +614,7 @@ async fn complete_websocket_initialize(websocket: &mut WebSocketStream<TcpStream
|
||||
id: request.id,
|
||||
result: serde_json::to_value(InitializeResponse {
|
||||
session_id: "session-1".to_string(),
|
||||
environment_info: None,
|
||||
})
|
||||
.expect("initialize response should serialize"),
|
||||
}),
|
||||
|
||||
@@ -156,7 +156,10 @@ impl ExecServerHandler {
|
||||
.session
|
||||
.lock()
|
||||
.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> {
|
||||
|
||||
Reference in New Issue
Block a user