From 3b7c38489802d5fa36318fab76065bfe23651887 Mon Sep 17 00:00:00 2001 From: Richard Lee Date: Tue, 9 Jun 2026 20:15:15 -0700 Subject: [PATCH] refresh remote environment URLs on reconnect --- codex-rs/core-api/src/lib.rs | 1 + codex-rs/exec-server/src/client.rs | 293 +---------- codex-rs/exec-server/src/client/error.rs | 124 +++++ .../exec-server/src/client/lazy_remote.rs | 155 ++++++ .../src/client/lazy_remote_tests.rs | 485 ++++++++++++++++++ codex-rs/exec-server/src/client_transport.rs | 38 +- codex-rs/exec-server/src/environment.rs | 136 ++--- codex-rs/exec-server/src/environment_info.rs | 54 ++ codex-rs/exec-server/src/lib.rs | 1 + codex-rs/exec-server/src/relay.rs | 4 +- 10 files changed, 926 insertions(+), 365 deletions(-) create mode 100644 codex-rs/exec-server/src/client/error.rs create mode 100644 codex-rs/exec-server/src/client/lazy_remote.rs create mode 100644 codex-rs/exec-server/src/client/lazy_remote_tests.rs create mode 100644 codex-rs/exec-server/src/environment_info.rs diff --git a/codex-rs/core-api/src/lib.rs b/codex-rs/core-api/src/lib.rs index 47452eb032..767890afae 100644 --- a/codex-rs/core-api/src/lib.rs +++ b/codex-rs/core-api/src/lib.rs @@ -47,6 +47,7 @@ pub use codex_core::resolve_installation_id; pub use codex_core::skills::SkillsManager; pub use codex_core::thread_store_from_config; pub use codex_exec_server::EnvironmentManager; +pub use codex_exec_server::ExecServerError; pub use codex_exec_server::ExecServerRuntimePaths; pub use codex_extension_api::empty_extension_registry; pub use codex_features::Feature; diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 26b3e46e88..826902d730 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -8,11 +8,8 @@ use std::time::Duration; use arc_swap::ArcSwap; use codex_app_server_protocol::JSONRPCNotification; -use futures::FutureExt; -use futures::future::BoxFuture; use serde_json::Value; use tokio::sync::Mutex; -use tokio::sync::Semaphore; use tokio::sync::mpsc; use tokio::sync::watch; @@ -21,8 +18,6 @@ use tracing::debug; use crate::ProcessId; use crate::client_api::ExecServerClientConnectOptions; -use crate::client_api::ExecServerTransportParams; -use crate::client_api::HttpClient; use crate::client_api::RemoteExecServerConnectArgs; use crate::client_api::StdioExecServerConnectArgs; use crate::connection::JsonRpcConnection; @@ -94,7 +89,14 @@ use crate::rpc::RpcCallError; use crate::rpc::RpcClient; use crate::rpc::RpcClientEvent; +mod error; pub(crate) mod http_client; +mod lazy_remote; + +pub use error::ExecServerError; +pub(crate) use lazy_remote::LazyRemoteExecServerClient; +pub(crate) use lazy_remote::RemoteConnectionSource; +pub(crate) use lazy_remote::RemoteWebSocketUrlProvider; const CONNECT_TIMEOUT: Duration = Duration::from_secs(10); const INITIALIZE_TIMEOUT: Duration = Duration::from_secs(10); @@ -204,133 +206,6 @@ pub struct ExecServerClient { inner: Arc, } -#[derive(Clone)] -pub(crate) struct LazyRemoteExecServerClient { - transport_params: ExecServerTransportParams, - client: Arc>>, - connect_lock: Arc, -} - -impl LazyRemoteExecServerClient { - pub(crate) fn new(transport_params: ExecServerTransportParams) -> Self { - Self { - transport_params, - client: Arc::new(StdMutex::new(None)), - connect_lock: Arc::new(Semaphore::new(/*permits*/ 1)), - } - } - - pub(crate) async fn get(&self) -> Result { - if let Some(client) = self.connected_client() { - return Ok(client); - } - - let _connect_permit = self.connect_lock.acquire().await.map_err(|_| { - ExecServerError::Protocol("exec-server connect lock closed".to_string()) - })?; - if let Some(client) = self.connected_client() { - return Ok(client); - } - - let next_client = match self.cached_client() { - Some(_client) - if matches!( - &self.transport_params, - ExecServerTransportParams::WebSocketUrl { .. } - ) => - { - ExecServerClient::connect_for_transport(self.transport_params.clone()).await? - } - Some(client) => return Ok(client), - None => ExecServerClient::connect_for_transport(self.transport_params.clone()).await?, - }; - - let mut cached_client = self - .client - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - *cached_client = Some(next_client.clone()); - Ok(next_client) - } - - fn connected_client(&self) -> Option { - self.cached_client() - .filter(|client| !client.is_disconnected()) - } - - fn cached_client(&self) -> Option { - self.client - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .clone() - } -} - -impl HttpClient for LazyRemoteExecServerClient { - fn http_request( - &self, - params: crate::HttpRequestParams, - ) -> BoxFuture<'_, Result> { - async move { self.get().await?.http_request(params).await }.boxed() - } - - fn http_request_stream( - &self, - params: crate::HttpRequestParams, - ) -> BoxFuture< - '_, - Result<(crate::HttpRequestResponse, crate::HttpResponseBodyStream), ExecServerError>, - > { - async move { self.get().await?.http_request_stream(params).await }.boxed() - } -} - -impl LazyRemoteExecServerClient { - pub(crate) async fn environment_info(&self) -> Result { - self.get().await?.environment_info().await - } -} - -#[derive(Debug, thiserror::Error)] -pub enum ExecServerError { - #[error("failed to spawn exec-server: {0}")] - Spawn(#[source] std::io::Error), - #[error("timed out connecting to exec-server websocket `{url}` after {timeout:?}")] - WebSocketConnectTimeout { url: String, timeout: Duration }, - #[error("failed to connect to exec-server websocket `{url}`: {source}")] - WebSocketConnect { - url: String, - #[source] - source: tokio_tungstenite::tungstenite::Error, - }, - #[error("timed out waiting for exec-server initialize handshake after {timeout:?}")] - InitializeTimedOut { timeout: Duration }, - #[error("exec-server transport closed")] - Closed, - #[error("{0}")] - Disconnected(String), - #[error("failed to serialize or deserialize exec-server JSON: {0}")] - Json(#[from] serde_json::Error), - #[error("HTTP request failed: {0}")] - HttpRequest(String), - #[error("exec-server protocol error: {0}")] - Protocol(String), - #[error("exec-server rejected request ({code}): {message}")] - Server { code: i64, message: String }, - #[error("environment registry request failed ({status}{code_suffix}): {message}", code_suffix = .code.as_ref().map(|code| format!(", {code}")).unwrap_or_default())] - EnvironmentRegistryHttp { - status: reqwest::StatusCode, - code: Option, - message: String, - }, - #[error("environment registry configuration error: {0}")] - EnvironmentRegistryConfig(String), - #[error("environment registry authentication error: {0}")] - EnvironmentRegistryAuth(String), - #[error("environment registry request failed: {0}")] - EnvironmentRegistryRequest(#[from] reqwest::Error), -} - impl ExecServerClient { pub async fn initialize( &self, @@ -584,7 +459,7 @@ impl ExecServerClient { .client .notify(INITIALIZED_METHOD, &serde_json::json!({})) .await - .map_err(ExecServerError::Json) + .map_err(ExecServerError::from) } async fn call(&self, method: &str, params: &P) -> Result @@ -622,7 +497,7 @@ impl From for ExecServerError { fn from(value: RpcCallError) -> Self { match value { RpcCallError::Closed => Self::Closed, - RpcCallError::Json(err) => Self::Json(err), + RpcCallError::Json(err) => Self::Json(Arc::new(err)), RpcCallError::Server(error) => Self::Server { code: error.code, message: error.message, @@ -976,35 +851,26 @@ mod tests { use codex_app_server_protocol::JSONRPCMessage; use codex_app_server_protocol::JSONRPCNotification; use codex_app_server_protocol::JSONRPCResponse; - use futures::SinkExt; - use futures::StreamExt; use pretty_assertions::assert_eq; use std::collections::HashMap; #[cfg(unix)] use std::path::Path; #[cfg(unix)] use std::process::Command; - use std::sync::Arc; use tokio::io::AsyncBufReadExt; use tokio::io::AsyncWrite; use tokio::io::AsyncWriteExt; use tokio::io::BufReader; use tokio::io::duplex; - use tokio::net::TcpListener; - use tokio::net::TcpStream; use tokio::sync::mpsc; use tokio::sync::oneshot; use tokio::time::Duration; #[cfg(unix)] use tokio::time::sleep; use tokio::time::timeout; - use tokio_tungstenite::WebSocketStream; - use tokio_tungstenite::accept_async; - use tokio_tungstenite::tungstenite::Message; use super::ExecServerClient; use super::ExecServerClientConnectOptions; - use super::LazyRemoteExecServerClient; use crate::ProcessId; #[cfg(not(windows))] use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT; @@ -1048,96 +914,6 @@ mod tests { .expect("json-rpc line should write"); } - async fn accept_websocket(listener: &TcpListener) -> WebSocketStream { - let (stream, _) = listener.accept().await.expect("listener should accept"); - accept_async(stream) - .await - .expect("websocket handshake should succeed") - } - - async fn read_jsonrpc_websocket(websocket: &mut WebSocketStream) -> JSONRPCMessage { - loop { - match timeout(Duration::from_secs(1), websocket.next()) - .await - .expect("json-rpc websocket read should not time out") - .expect("websocket should stay open") - .expect("websocket frame should read") - { - Message::Text(text) => { - return serde_json::from_str(text.as_ref()) - .expect("json-rpc text frame should parse"); - } - Message::Binary(bytes) => { - return serde_json::from_slice(bytes.as_ref()) - .expect("json-rpc binary frame should parse"); - } - Message::Ping(_) | Message::Pong(_) => {} - other => panic!("expected json-rpc websocket frame, got {other:?}"), - } - } - } - - async fn write_jsonrpc_websocket( - websocket: &mut WebSocketStream, - message: JSONRPCMessage, - ) { - let encoded = serde_json::to_string(&message).expect("json-rpc should serialize"); - websocket - .send(Message::Text(encoded.into())) - .await - .expect("json-rpc websocket frame should write"); - } - - async fn complete_websocket_initialize( - websocket: &mut WebSocketStream, - session_id: &str, - expected_resume_session_id: Option<&str>, - ) { - let initialize = read_jsonrpc_websocket(websocket).await; - let request = match initialize { - JSONRPCMessage::Request(request) if request.method == INITIALIZE_METHOD => request, - other => panic!("expected initialize request, got {other:?}"), - }; - let params: crate::protocol::InitializeParams = - serde_json::from_value(request.params.expect("initialize params should exist")) - .expect("initialize params should deserialize"); - assert_eq!( - params.resume_session_id.as_deref(), - expected_resume_session_id - ); - write_jsonrpc_websocket( - websocket, - JSONRPCMessage::Response(JSONRPCResponse { - id: request.id, - result: serde_json::to_value(InitializeResponse { - session_id: session_id.to_string(), - }) - .expect("initialize response should serialize"), - }), - ) - .await; - - let initialized = read_jsonrpc_websocket(websocket).await; - match initialized { - JSONRPCMessage::Notification(notification) - if notification.method == INITIALIZED_METHOD => {} - other => panic!("expected initialized notification, got {other:?}"), - } - } - - async fn wait_for_disconnect(client: &ExecServerClient) { - timeout(Duration::from_secs(1), async { - loop { - if client.is_disconnected() { - return; - } - tokio::task::yield_now().await; - } - }) - .await - .expect("client should observe disconnect"); - } - #[cfg(not(windows))] #[tokio::test] async fn connect_stdio_command_initializes_json_rpc_client() { @@ -1555,57 +1331,6 @@ mod tests { server.await.expect("server task should finish"); } - #[tokio::test] - async fn remote_websocket_client_replaces_disconnected_client_with_fresh_session() { - let listener = TcpListener::bind("127.0.0.1:0") - .await - .expect("listener should bind"); - let websocket_url = format!( - "ws://{}", - listener.local_addr().expect("listener should have address") - ); - let server = tokio::spawn({ - async move { - let mut first = accept_websocket(&listener).await; - complete_websocket_initialize( - &mut first, - "session-1", - /*expected_resume_session_id*/ None, - ) - .await; - first - .close(None) - .await - .expect("first websocket should close"); - - let mut second = accept_websocket(&listener).await; - complete_websocket_initialize( - &mut second, - "session-2", - /*expected_resume_session_id*/ None, - ) - .await; - } - }); - - let client = LazyRemoteExecServerClient::new(ExecServerTransportParams::WebSocketUrl { - websocket_url, - connect_timeout: Duration::from_secs(1), - initialize_timeout: Duration::from_secs(1), - }); - let first = client.get().await.expect("first client should connect"); - wait_for_disconnect(&first).await; - - let (replacement_a, replacement_b) = tokio::join!(client.get(), client.get()); - let replacement_a = replacement_a.expect("first replacement should connect"); - let replacement_b = replacement_b.expect("second replacement should reuse client"); - assert_eq!(replacement_a.session_id().as_deref(), Some("session-2")); - assert_eq!(replacement_b.session_id().as_deref(), Some("session-2")); - assert!(Arc::ptr_eq(&replacement_a.inner, &replacement_b.inner)); - - server.await.expect("server task should finish"); - } - #[tokio::test] async fn wake_notifications_do_not_block_other_sessions() { let (client_stdin, server_reader) = duplex(1 << 20); diff --git a/codex-rs/exec-server/src/client/error.rs b/codex-rs/exec-server/src/client/error.rs new file mode 100644 index 0000000000..4bcd19d336 --- /dev/null +++ b/codex-rs/exec-server/src/client/error.rs @@ -0,0 +1,124 @@ +use std::fmt; +use std::sync::Arc; +use std::time::Duration; + +#[derive(Clone, Debug)] +pub enum ExecServerError { + Spawn(Arc), + WebSocketConnectTimeout { + url: String, + timeout: Duration, + }, + WebSocketConnect { + url: String, + source: Arc, + }, + InitializeTimedOut { + timeout: Duration, + }, + Closed, + Disconnected(String), + Json(Arc), + HttpRequest(String), + Protocol(String), + Server { + code: i64, + message: String, + }, + EnvironmentRegistryHttp { + status: reqwest::StatusCode, + code: Option, + message: String, + }, + EnvironmentRegistryConfig(String), + EnvironmentRegistryAuth(String), + EnvironmentRegistryRequest(Arc), +} + +impl fmt::Display for ExecServerError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Spawn(error) => write!(f, "failed to spawn exec-server: {error}"), + Self::WebSocketConnectTimeout { url, timeout } => write!( + f, + "timed out connecting to exec-server websocket `{url}` after {timeout:?}" + ), + Self::WebSocketConnect { url, source } => { + write!( + f, + "failed to connect to exec-server websocket `{url}`: {source}" + ) + } + Self::InitializeTimedOut { timeout } => write!( + f, + "timed out waiting for exec-server initialize handshake after {timeout:?}" + ), + Self::Closed => f.write_str("exec-server transport closed"), + Self::Disconnected(message) => f.write_str(message), + Self::Json(error) => { + write!( + f, + "failed to serialize or deserialize exec-server JSON: {error}" + ) + } + Self::HttpRequest(message) => write!(f, "HTTP request failed: {message}"), + Self::Protocol(message) => write!(f, "exec-server protocol error: {message}"), + Self::Server { code, message } => { + write!(f, "exec-server rejected request ({code}): {message}") + } + Self::EnvironmentRegistryHttp { + status, + code, + message, + } => { + write!(f, "environment registry request failed ({status}")?; + if let Some(code) = code { + write!(f, ", {code}")?; + } + write!(f, "): {message}") + } + Self::EnvironmentRegistryConfig(message) => { + write!(f, "environment registry configuration error: {message}") + } + Self::EnvironmentRegistryAuth(message) => { + write!(f, "environment registry authentication error: {message}") + } + Self::EnvironmentRegistryRequest(error) => { + write!(f, "environment registry request failed: {error}") + } + } + } +} + +impl std::error::Error for ExecServerError { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + match self { + Self::Spawn(error) => Some(error.as_ref()), + Self::WebSocketConnect { source, .. } => Some(source.as_ref()), + Self::Json(error) => Some(error.as_ref()), + Self::EnvironmentRegistryRequest(error) => Some(error.as_ref()), + Self::WebSocketConnectTimeout { .. } + | Self::InitializeTimedOut { .. } + | Self::Closed + | Self::Disconnected(_) + | Self::HttpRequest(_) + | Self::Protocol(_) + | Self::Server { .. } + | Self::EnvironmentRegistryHttp { .. } + | Self::EnvironmentRegistryConfig(_) + | Self::EnvironmentRegistryAuth(_) => None, + } + } +} + +impl From for ExecServerError { + fn from(error: serde_json::Error) -> Self { + Self::Json(Arc::new(error)) + } +} + +impl From for ExecServerError { + fn from(error: reqwest::Error) -> Self { + Self::EnvironmentRegistryRequest(Arc::new(error)) + } +} diff --git a/codex-rs/exec-server/src/client/lazy_remote.rs b/codex-rs/exec-server/src/client/lazy_remote.rs new file mode 100644 index 0000000000..edd8439f07 --- /dev/null +++ b/codex-rs/exec-server/src/client/lazy_remote.rs @@ -0,0 +1,155 @@ +use std::sync::Arc; +use std::sync::Mutex as StdMutex; + +use futures::FutureExt; +use futures::future::BoxFuture; +use futures::future::Shared; + +use crate::ExecServerClient; +use crate::ExecServerError; +use crate::HttpClient; +use crate::client_api::ExecServerTransportParams; +use crate::protocol::EnvironmentInfo; + +pub(crate) type RemoteWebSocketUrlProvider = + Arc BoxFuture<'static, Result> + Send + Sync>; + +#[derive(Clone)] +pub(crate) enum RemoteConnectionSource { + Fixed(ExecServerTransportParams), + RefreshingWebSocket(RemoteWebSocketUrlProvider), +} + +impl RemoteConnectionSource { + pub(crate) fn configured_url(&self) -> Option<&str> { + match self { + Self::Fixed(ExecServerTransportParams::WebSocketUrl { websocket_url, .. }) => { + Some(websocket_url) + } + Self::Fixed(ExecServerTransportParams::StdioCommand { .. }) + | Self::RefreshingWebSocket(_) => None, + } + } + + fn is_reconnectable(&self) -> bool { + matches!( + self, + Self::Fixed(ExecServerTransportParams::WebSocketUrl { .. }) + | Self::RefreshingWebSocket(_) + ) + } + + async fn connect(&self) -> Result { + let transport_params = match self { + Self::Fixed(transport_params) => transport_params.clone(), + Self::RefreshingWebSocket(provider) => { + ExecServerTransportParams::websocket_url(provider().await?) + } + }; + ExecServerClient::connect_for_transport(transport_params).await + } +} + +#[derive(Clone)] +pub(crate) struct LazyRemoteExecServerClient { + source: RemoteConnectionSource, + state: Arc>, +} + +#[derive(Default)] +struct ConnectionState { + client: Option, + in_flight: Option, + next_attempt_id: u64, +} + +struct InFlightConnection { + attempt_id: u64, + attempt: SharedConnectionAttempt, +} + +type SharedConnectionAttempt = + Shared>>; + +impl LazyRemoteExecServerClient { + pub(crate) fn new(source: RemoteConnectionSource) -> Self { + Self { + source, + state: Arc::new(StdMutex::new(ConnectionState::default())), + } + } + + pub(crate) async fn get(&self) -> Result { + let (attempt_id, attempt) = { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(client) = state.client.as_ref() + && (!client.is_disconnected() || !self.source.is_reconnectable()) + { + return Ok(client.clone()); + } + + if let Some(in_flight) = state.in_flight.as_ref() { + (in_flight.attempt_id, in_flight.attempt.clone()) + } else { + let attempt_id = state.next_attempt_id; + state.next_attempt_id = state.next_attempt_id.wrapping_add(1); + let source = self.source.clone(); + let attempt = async move { source.connect().await }.boxed().shared(); + state.in_flight = Some(InFlightConnection { + attempt_id, + attempt: attempt.clone(), + }); + (attempt_id, attempt) + } + }; + + let result = attempt.await; + { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state + .in_flight + .as_ref() + .is_some_and(|in_flight| in_flight.attempt_id == attempt_id) + { + if let Ok(connected_client) = &result { + state.client = Some(connected_client.clone()); + } + state.in_flight = None; + } + } + result + } + + pub(crate) async fn environment_info(&self) -> Result { + self.get().await?.environment_info().await + } +} + +impl HttpClient for LazyRemoteExecServerClient { + fn http_request( + &self, + params: crate::HttpRequestParams, + ) -> BoxFuture<'_, Result> { + async move { self.get().await?.http_request(params).await }.boxed() + } + + fn http_request_stream( + &self, + params: crate::HttpRequestParams, + ) -> BoxFuture< + '_, + Result<(crate::HttpRequestResponse, crate::HttpResponseBodyStream), ExecServerError>, + > { + async move { self.get().await?.http_request_stream(params).await }.boxed() + } +} + +#[cfg(test)] +#[path = "lazy_remote_tests.rs"] +mod tests; diff --git a/codex-rs/exec-server/src/client/lazy_remote_tests.rs b/codex-rs/exec-server/src/client/lazy_remote_tests.rs new file mode 100644 index 0000000000..6ec4a5828a --- /dev/null +++ b/codex-rs/exec-server/src/client/lazy_remote_tests.rs @@ -0,0 +1,485 @@ +use std::sync::Arc; +use std::sync::Mutex as StdMutex; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::time::Duration; + +use codex_app_server_protocol::JSONRPCMessage; +use codex_app_server_protocol::JSONRPCResponse; +use futures::FutureExt; +use futures::SinkExt; +use futures::StreamExt; +use pretty_assertions::assert_eq; +use tokio::net::TcpListener; +use tokio::net::TcpStream; +use tokio::sync::oneshot; +use tokio::sync::watch; +use tokio::time::timeout; +use tokio_tungstenite::WebSocketStream; +use tokio_tungstenite::accept_async; +use tokio_tungstenite::tungstenite::Message; + +use super::LazyRemoteExecServerClient; +use super::RemoteConnectionSource; +use crate::EnvironmentManager; +use crate::ExecServerClient; +use crate::ExecServerError; +use crate::client_api::ExecServerTransportParams; +use crate::protocol::ENVIRONMENT_INFO_METHOD; +use crate::protocol::EnvironmentInfo; +use crate::protocol::INITIALIZE_METHOD; +use crate::protocol::INITIALIZED_METHOD; +use crate::protocol::InitializeResponse; +use crate::protocol::ShellInfo; + +async fn accept_websocket(listener: &TcpListener) -> WebSocketStream { + let (stream, _) = listener.accept().await.expect("listener should accept"); + accept_async(stream) + .await + .expect("websocket handshake should succeed") +} + +async fn read_jsonrpc_websocket(websocket: &mut WebSocketStream) -> JSONRPCMessage { + loop { + match timeout(Duration::from_secs(1), websocket.next()) + .await + .expect("json-rpc websocket read should not time out") + .expect("websocket should stay open") + .expect("websocket frame should read") + { + Message::Text(text) => { + return serde_json::from_str(text.as_ref()) + .expect("json-rpc text frame should parse"); + } + Message::Binary(bytes) => { + return serde_json::from_slice(bytes.as_ref()) + .expect("json-rpc binary frame should parse"); + } + Message::Ping(_) | Message::Pong(_) => {} + other => panic!("expected json-rpc websocket frame, got {other:?}"), + } + } +} + +async fn write_jsonrpc_websocket( + websocket: &mut WebSocketStream, + message: JSONRPCMessage, +) { + let encoded = serde_json::to_string(&message).expect("json-rpc should serialize"); + websocket + .send(Message::Text(encoded.into())) + .await + .expect("json-rpc websocket frame should write"); +} + +async fn complete_websocket_initialize( + websocket: &mut WebSocketStream, + session_id: &str, +) { + let initialize = read_jsonrpc_websocket(websocket).await; + let request = match initialize { + JSONRPCMessage::Request(request) if request.method == INITIALIZE_METHOD => request, + other => panic!("expected initialize request, got {other:?}"), + }; + let params: crate::protocol::InitializeParams = + serde_json::from_value(request.params.expect("initialize params should exist")) + .expect("initialize params should deserialize"); + assert_eq!(params.resume_session_id, None); + write_jsonrpc_websocket( + websocket, + JSONRPCMessage::Response(JSONRPCResponse { + id: request.id, + result: serde_json::to_value(InitializeResponse { + session_id: session_id.to_string(), + }) + .expect("initialize response should serialize"), + }), + ) + .await; + + let initialized = read_jsonrpc_websocket(websocket).await; + match initialized { + JSONRPCMessage::Notification(notification) if notification.method == INITIALIZED_METHOD => { + } + other => panic!("expected initialized notification, got {other:?}"), + } +} + +async fn complete_environment_info(websocket: &mut WebSocketStream) { + let request = match read_jsonrpc_websocket(websocket).await { + JSONRPCMessage::Request(request) if request.method == ENVIRONMENT_INFO_METHOD => request, + other => panic!("expected environment info request, got {other:?}"), + }; + write_jsonrpc_websocket( + websocket, + JSONRPCMessage::Response(JSONRPCResponse { + id: request.id, + result: serde_json::to_value(EnvironmentInfo { + shell: ShellInfo { + name: "sh".to_string(), + path: "/bin/sh".to_string(), + }, + }) + .expect("environment info response should serialize"), + }), + ) + .await; +} + +async fn wait_for_disconnect(client: &ExecServerClient) { + timeout(Duration::from_secs(1), async { + while !client.is_disconnected() { + tokio::task::yield_now().await; + } + }) + .await + .expect("client should observe disconnect"); +} + +#[tokio::test] +async fn replaces_disconnected_websocket_client_with_one_shared_connection() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let websocket_url = format!( + "ws://{}", + listener.local_addr().expect("listener should have address") + ); + let server = tokio::spawn(async move { + let mut first = accept_websocket(&listener).await; + complete_websocket_initialize(&mut first, "session-1").await; + first + .close(None) + .await + .expect("first websocket should close"); + + let mut second = accept_websocket(&listener).await; + complete_websocket_initialize(&mut second, "session-2").await; + }); + + let client = LazyRemoteExecServerClient::new(RemoteConnectionSource::Fixed( + ExecServerTransportParams::WebSocketUrl { + websocket_url, + connect_timeout: Duration::from_secs(1), + initialize_timeout: Duration::from_secs(1), + }, + )); + let first = client.get().await.expect("first client should connect"); + wait_for_disconnect(&first).await; + + let (replacement_a, replacement_b) = tokio::join!(client.get(), client.get()); + let replacement_a = replacement_a.expect("first replacement should connect"); + let replacement_b = replacement_b.expect("second replacement should reuse client"); + assert_eq!(replacement_a.session_id().as_deref(), Some("session-2")); + assert_eq!(replacement_b.session_id().as_deref(), Some("session-2")); + assert!(Arc::ptr_eq(&replacement_a.inner, &replacement_b.inner)); + + server.await.expect("server task should finish"); +} + +#[tokio::test] +async fn refreshes_websocket_url_for_reconnect() { + let first_listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let second_listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let websocket_urls = Arc::new([ + format!( + "ws://{}", + first_listener + .local_addr() + .expect("listener should have address") + ), + format!( + "ws://{}", + second_listener + .local_addr() + .expect("listener should have address") + ), + ]); + let server = tokio::spawn(async move { + let mut first = accept_websocket(&first_listener).await; + complete_websocket_initialize(&mut first, "session-1").await; + first + .close(None) + .await + .expect("first websocket should close"); + + let mut second = accept_websocket(&second_listener).await; + complete_websocket_initialize(&mut second, "session-2").await; + }); + let provider_calls = Arc::new(AtomicUsize::new(0)); + let client = + LazyRemoteExecServerClient::new(RemoteConnectionSource::RefreshingWebSocket(Arc::new({ + let provider_calls = Arc::clone(&provider_calls); + let websocket_urls = Arc::clone(&websocket_urls); + move || { + let call = provider_calls.fetch_add(1, Ordering::Relaxed); + let websocket_url = websocket_urls.get(call).cloned().ok_or_else(|| { + ExecServerError::Protocol("unexpected URL provider call".to_string()) + }); + async move { websocket_url }.boxed() + } + }))); + + let first = client.get().await.expect("first client should connect"); + wait_for_disconnect(&first).await; + let second = client.get().await.expect("replacement should connect"); + + assert_eq!(second.session_id().as_deref(), Some("session-2")); + assert_eq!(provider_calls.load(Ordering::Relaxed), 2); + server.await.expect("server task should finish"); +} + +#[tokio::test] +async fn environment_manager_provider_is_lazy_and_refreshes_after_disconnect() { + let first_listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let second_listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let websocket_urls = Arc::new([ + format!( + "ws://{}", + first_listener + .local_addr() + .expect("listener should have address") + ), + format!( + "ws://{}", + second_listener + .local_addr() + .expect("listener should have address") + ), + ]); + let server = tokio::spawn(async move { + let mut first = accept_websocket(&first_listener).await; + complete_websocket_initialize(&mut first, "session-1").await; + complete_environment_info(&mut first).await; + first + .close(None) + .await + .expect("first websocket should close"); + + let mut second = accept_websocket(&second_listener).await; + complete_websocket_initialize(&mut second, "session-2").await; + complete_environment_info(&mut second).await; + }); + let provider_calls = Arc::new(AtomicUsize::new(0)); + let manager = EnvironmentManager::without_environments(); + manager + .upsert_environment_with_url_provider("executor-a".to_string(), { + let provider_calls = Arc::clone(&provider_calls); + let websocket_urls = Arc::clone(&websocket_urls); + move || { + let call = provider_calls.fetch_add(1, Ordering::Relaxed); + let websocket_url = websocket_urls.get(call).cloned().ok_or_else(|| { + ExecServerError::Protocol("unexpected URL provider call".to_string()) + }); + async move { websocket_url } + } + }) + .expect("provider-backed environment should be installed"); + let environment = manager + .get_environment("executor-a") + .expect("provider-backed environment should exist"); + assert_eq!(provider_calls.load(Ordering::Relaxed), 0); + + let first_info = environment.info().await.expect("first info should succeed"); + assert_eq!(first_info.shell.name, "sh"); + let second_info = timeout(Duration::from_secs(1), async { + loop { + match environment.info().await { + Ok(info) => break info, + Err(ExecServerError::Closed | ExecServerError::Disconnected(_)) => { + tokio::task::yield_now().await; + } + Err(error) => panic!("unexpected reconnect error: {error}"), + } + } + }) + .await + .expect("provider-backed environment should reconnect"); + + assert_eq!(second_info, first_info); + assert_eq!(provider_calls.load(Ordering::Relaxed), 2); + assert!(Arc::ptr_eq( + &environment, + &manager + .get_environment("executor-a") + .expect("environment should remain installed") + )); + server.await.expect("server task should finish"); +} + +#[tokio::test] +async fn redacts_refreshing_websocket_url_connection_errors() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let address = listener.local_addr().expect("listener should have address"); + let websocket_url = format!("ws://{address}/environment?sig=secret"); + drop(listener); + let client = LazyRemoteExecServerClient::new(RemoteConnectionSource::RefreshingWebSocket( + Arc::new(move || { + let websocket_url = websocket_url.clone(); + async move { Ok(websocket_url) }.boxed() + }), + )); + + let error = match client.get().await { + Ok(_) => panic!("connection should fail"), + Err(error) => error, + }; + let ExecServerError::WebSocketConnect { url, .. } = error else { + panic!("expected websocket connection error"); + }; + assert_eq!(url, format!("ws://{address}/environment")); +} + +#[tokio::test] +async fn coalesces_only_concurrent_provider_failures() { + let provider_calls = Arc::new(AtomicUsize::new(0)); + let (started_tx, started_rx) = oneshot::channel(); + let started_tx = Arc::new(StdMutex::new(Some(started_tx))); + let (release_tx, release_rx) = watch::channel(false); + let client = + LazyRemoteExecServerClient::new(RemoteConnectionSource::RefreshingWebSocket(Arc::new({ + let provider_calls = Arc::clone(&provider_calls); + let started_tx = Arc::clone(&started_tx); + move || { + let call = provider_calls.fetch_add(1, Ordering::Relaxed); + let started_tx = Arc::clone(&started_tx); + let mut release_rx = release_rx.clone(); + async move { + if call == 0 { + if let Some(started_tx) = started_tx + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .take() + { + let _ = started_tx.send(()); + } + release_rx + .wait_for(|released| *released) + .await + .expect("release sender should stay open"); + } + Err(ExecServerError::Protocol( + "environment registry unavailable".to_string(), + )) + } + .boxed() + } + }))); + + let first = tokio::spawn({ + let client = client.clone(); + async move { client.get().await } + }); + started_rx.await.expect("provider should start"); + let second = tokio::spawn({ + let client = client.clone(); + async move { client.get().await } + }); + tokio::task::yield_now().await; + release_tx.send(true).expect("provider should be waiting"); + + let first = first.await.expect("first task should finish"); + let second = second.await.expect("second task should finish"); + match (first, second) { + (Err(ExecServerError::Protocol(first)), Err(ExecServerError::Protocol(second))) => { + assert_eq!(first, second); + } + _ => panic!("both callers should receive the same protocol error"), + } + assert_eq!(provider_calls.load(Ordering::Relaxed), 1); + + assert!(client.get().await.is_err()); + assert_eq!(provider_calls.load(Ordering::Relaxed), 2); +} + +#[tokio::test] +async fn waiter_completes_when_connection_starter_is_cancelled() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let websocket_url = format!( + "ws://{}", + listener.local_addr().expect("listener should have address") + ); + let (server_release_tx, server_release_rx) = oneshot::channel(); + let server = tokio::spawn(async move { + let mut websocket = accept_websocket(&listener).await; + complete_websocket_initialize(&mut websocket, "session-1").await; + let _ = server_release_rx.await; + }); + + let provider_calls = Arc::new(AtomicUsize::new(0)); + let (provider_started_tx, provider_started_rx) = oneshot::channel(); + let provider_started_tx = Arc::new(StdMutex::new(Some(provider_started_tx))); + let (provider_release_tx, provider_release_rx) = watch::channel(false); + let client = + LazyRemoteExecServerClient::new(RemoteConnectionSource::RefreshingWebSocket(Arc::new({ + let provider_calls = Arc::clone(&provider_calls); + move || { + provider_calls.fetch_add(1, Ordering::Relaxed); + let provider_started_tx = Arc::clone(&provider_started_tx); + let mut provider_release_rx = provider_release_rx.clone(); + let websocket_url = websocket_url.clone(); + async move { + if let Some(provider_started_tx) = provider_started_tx + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .take() + { + let _ = provider_started_tx.send(()); + } + provider_release_rx + .wait_for(|released| *released) + .await + .expect("provider release sender should stay open"); + Ok(websocket_url) + } + .boxed() + } + }))); + + let starter = tokio::spawn({ + let client = client.clone(); + async move { client.get().await } + }); + provider_started_rx + .await + .expect("connection provider should start"); + let waiter = tokio::spawn({ + let client = client.clone(); + async move { client.get().await } + }); + tokio::task::yield_now().await; + starter.abort(); + match starter.await { + Err(error) if error.is_cancelled() => {} + _ => panic!("connection starter should be cancelled"), + } + provider_release_tx + .send(true) + .expect("connection provider should be waiting"); + + let connected = waiter + .await + .expect("waiter task should finish") + .expect("waiter should complete the shared connection"); + let cached = client + .get() + .await + .expect("connected client should be cached"); + assert!(Arc::ptr_eq(&connected.inner, &cached.inner)); + assert_eq!(provider_calls.load(Ordering::Relaxed), 1); + + let _ = server_release_tx.send(()); + server.await.expect("server task should finish"); +} diff --git a/codex-rs/exec-server/src/client_transport.rs b/codex-rs/exec-server/src/client_transport.rs index 4bdc09a80e..706cd2ec6d 100644 --- a/codex-rs/exec-server/src/client_transport.rs +++ b/codex-rs/exec-server/src/client_transport.rs @@ -1,4 +1,5 @@ use std::process::Stdio; +use std::sync::Arc; use tokio::io::AsyncBufReadExt; use tokio::io::BufReader; use tokio::process::Command; @@ -58,20 +59,28 @@ impl ExecServerClient { ) -> Result { ensure_rustls_crypto_provider(); let websocket_url = args.websocket_url.clone(); + let parsed_websocket_url = reqwest::Url::parse(&websocket_url).ok(); + let redacted_websocket_url = parsed_websocket_url + .as_ref() + .map(redact_websocket_url) + .unwrap_or_else(|| "".to_string()); let connect_timeout = args.connect_timeout; let (stream, _) = timeout(connect_timeout, connect_async(websocket_url.as_str())) .await .map_err(|_| ExecServerError::WebSocketConnectTimeout { - url: websocket_url.clone(), + url: redacted_websocket_url.clone(), timeout: connect_timeout, })? .map_err(|source| ExecServerError::WebSocketConnect { - url: websocket_url.clone(), - source, + url: redacted_websocket_url.clone(), + source: source.into(), })?; - let connection_label = format!("exec-server websocket {websocket_url}"); - let connection = if is_rendezvous_harness_url(&websocket_url) { + let connection_label = format!("exec-server websocket {redacted_websocket_url}"); + let connection = if parsed_websocket_url + .as_ref() + .is_some_and(is_rendezvous_harness_url) + { harness_connection_from_websocket(stream, connection_label) } else { JsonRpcConnection::from_websocket(stream, connection_label) @@ -87,7 +96,7 @@ impl ExecServerClient { .stdout(Stdio::piped()) .stderr(Stdio::piped()) .spawn() - .map_err(ExecServerError::Spawn)?; + .map_err(|error| ExecServerError::Spawn(Arc::new(error)))?; let stdin = child.stdin.take().ok_or_else(|| { ExecServerError::Protocol("spawned exec-server command has no stdin".to_string()) @@ -120,13 +129,16 @@ impl ExecServerClient { } } -fn is_rendezvous_harness_url(websocket_url: &str) -> bool { - let Some((_path, query)) = websocket_url.split_once('?') else { - return false; - }; - query - .split('&') - .filter_map(|pair| pair.split_once('=')) +fn redact_websocket_url(websocket_url: &reqwest::Url) -> String { + let mut websocket_url = websocket_url.clone(); + websocket_url.set_query(None); + websocket_url.set_fragment(None); + websocket_url.to_string() +} + +fn is_rendezvous_harness_url(websocket_url: &reqwest::Url) -> bool { + websocket_url + .query_pairs() .any(|(key, value)| key == "role" && value == "harness") } diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index 1dfd49d3fd..32a8b6d833 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -1,17 +1,22 @@ use std::collections::HashMap; +use std::future::Future; use std::sync::Arc; use std::sync::RwLock; use futures::FutureExt; -use futures::future::BoxFuture; use crate::ExecServerError; use crate::ExecServerRuntimePaths; use crate::ExecutorFileSystem; use crate::HttpClient; use crate::client::LazyRemoteExecServerClient; +use crate::client::RemoteConnectionSource; +use crate::client::RemoteWebSocketUrlProvider; use crate::client::http_client::ReqwestHttpClient; use crate::client_api::ExecServerTransportParams; +use crate::environment_info::EnvironmentInfoProvider; +use crate::environment_info::LocalEnvironmentInfoProvider; +use crate::environment_info::RemoteEnvironmentInfoProvider; use crate::environment_provider::DefaultEnvironmentProvider; use crate::environment_provider::EnvironmentDefault; use crate::environment_provider::EnvironmentProvider; @@ -22,10 +27,8 @@ use crate::local_file_system::LocalFileSystem; use crate::local_process::LocalProcess; use crate::process::ExecBackend; use crate::protocol::EnvironmentInfo; -use crate::protocol::ShellInfo; use crate::remote_file_system::RemoteFileSystem; use crate::remote_process::RemoteProcess; -use codex_shell_command::shell_detect::DetectedShell; pub const CODEX_EXEC_SERVER_URL_ENV_VAR: &str = "CODEX_EXEC_SERVER_URL"; @@ -282,6 +285,35 @@ impl EnvironmentManager { .insert(environment_id, Arc::new(environment)); Ok(()) } + + /// Adds or replaces a remote environment whose signed WebSocket URL is + /// resolved before every connection attempt. + pub fn upsert_environment_with_url_provider( + &self, + environment_id: String, + websocket_url_provider: F, + ) -> Result<(), ExecServerError> + where + F: Fn() -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, + { + if environment_id.is_empty() { + return Err(ExecServerError::Protocol( + "environment id cannot be empty".to_string(), + )); + } + let websocket_url_provider: RemoteWebSocketUrlProvider = + Arc::new(move || websocket_url_provider().boxed()); + let environment = Environment::remote_with_url_provider( + websocket_url_provider, + self.local_runtime_paths.clone(), + ); + self.environments + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .insert(environment_id, Arc::new(environment)); + Ok(()) + } } /// Concrete execution/filesystem environment selected for a session. @@ -290,8 +322,7 @@ impl EnvironmentManager { /// paths used by filesystem helpers. #[derive(Clone)] pub struct Environment { - exec_server_url: Option, - remote_transport: Option, + connection: EnvironmentConnection, info_provider: Arc, exec_backend: Arc, filesystem: Arc, @@ -299,41 +330,17 @@ pub struct Environment { local_runtime_paths: Option, } -/// Provides environment metadata from either a local environment or a remote exec-server. -trait EnvironmentInfoProvider: Send + Sync { - fn info(&self) -> BoxFuture<'_, Result>; -} - -struct LocalEnvironmentInfoProvider; - -impl EnvironmentInfoProvider for LocalEnvironmentInfoProvider { - fn info(&self) -> BoxFuture<'_, Result> { - std::future::ready(Ok(EnvironmentInfo::local())).boxed() - } -} - -struct RemoteEnvironmentInfoProvider { - client: LazyRemoteExecServerClient, -} - -impl RemoteEnvironmentInfoProvider { - fn new(client: LazyRemoteExecServerClient) -> Self { - Self { client } - } -} - -impl EnvironmentInfoProvider for RemoteEnvironmentInfoProvider { - fn info(&self) -> BoxFuture<'_, Result> { - async move { self.client.environment_info().await }.boxed() - } +#[derive(Clone)] +enum EnvironmentConnection { + Local, + Remote(RemoteConnectionSource), } impl Environment { /// Builds a test-only local environment without configured sandbox helper paths. pub fn default_for_tests() -> Self { Self { - exec_server_url: None, - remote_transport: None, + connection: EnvironmentConnection::Local, info_provider: Arc::new(LocalEnvironmentInfoProvider), exec_backend: Arc::new(LocalProcess::default()), filesystem: Arc::new(LocalFileSystem::unsandboxed()), @@ -346,7 +353,7 @@ impl Environment { impl std::fmt::Debug for Environment { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("Environment") - .field("exec_server_url", &self.exec_server_url) + .field("exec_server_url", &self.exec_server_url()) .finish_non_exhaustive() } } @@ -389,8 +396,7 @@ impl Environment { pub(crate) fn local(local_runtime_paths: ExecServerRuntimePaths) -> Self { Self { - exec_server_url: None, - remote_transport: None, + connection: EnvironmentConnection::Local, info_provider: Arc::new(LocalEnvironmentInfoProvider), exec_backend: Arc::new(LocalProcess::default()), filesystem: Arc::new(LocalFileSystem::with_runtime_paths( @@ -415,21 +421,23 @@ impl Environment { remote_transport: ExecServerTransportParams, local_runtime_paths: Option, ) -> Self { - let exec_server_url = match &remote_transport { - ExecServerTransportParams::WebSocketUrl { - websocket_url: exec_server_url, - .. - } => Some(exec_server_url.clone()), - ExecServerTransportParams::StdioCommand { .. } => None, - }; - let client = LazyRemoteExecServerClient::new(remote_transport.clone()); + Self::remote( + RemoteConnectionSource::Fixed(remote_transport), + local_runtime_paths, + ) + } + + fn remote( + connection_source: RemoteConnectionSource, + local_runtime_paths: Option, + ) -> Self { + let client = LazyRemoteExecServerClient::new(connection_source.clone()); let exec_backend: Arc = Arc::new(RemoteProcess::new(client.clone())); let filesystem: Arc = Arc::new(RemoteFileSystem::new(client.clone())); Self { - exec_server_url, - remote_transport: Some(remote_transport), + connection: EnvironmentConnection::Remote(connection_source), info_provider: Arc::new(RemoteEnvironmentInfoProvider::new(client.clone())), exec_backend, filesystem, @@ -438,13 +446,26 @@ impl Environment { } } + fn remote_with_url_provider( + websocket_url_provider: RemoteWebSocketUrlProvider, + local_runtime_paths: Option, + ) -> Self { + Self::remote( + RemoteConnectionSource::RefreshingWebSocket(websocket_url_provider), + local_runtime_paths, + ) + } + pub fn is_remote(&self) -> bool { - self.remote_transport.is_some() + matches!(&self.connection, EnvironmentConnection::Remote(_)) } /// Returns the remote exec-server URL when this environment is remote. pub fn exec_server_url(&self) -> Option<&str> { - self.exec_server_url.as_deref() + match &self.connection { + EnvironmentConnection::Local => None, + EnvironmentConnection::Remote(connection_source) => connection_source.configured_url(), + } } pub fn local_runtime_paths(&self) -> Option<&ExecServerRuntimePaths> { @@ -469,23 +490,6 @@ impl Environment { } } -impl EnvironmentInfo { - pub(crate) fn local() -> Self { - Self { - shell: codex_shell_command::shell_detect::default_user_shell().into(), - } - } -} - -impl From for ShellInfo { - fn from(shell: DetectedShell) -> Self { - Self { - name: shell.name().to_string(), - path: shell.shell_path.to_string_lossy().into_owned(), - } - } -} - #[cfg(test)] mod tests { use std::sync::Arc; diff --git a/codex-rs/exec-server/src/environment_info.rs b/codex-rs/exec-server/src/environment_info.rs new file mode 100644 index 0000000000..6eb192b433 --- /dev/null +++ b/codex-rs/exec-server/src/environment_info.rs @@ -0,0 +1,54 @@ +use futures::FutureExt; +use futures::future::BoxFuture; + +use crate::ExecServerError; +use crate::client::LazyRemoteExecServerClient; +use crate::protocol::EnvironmentInfo; +use crate::protocol::ShellInfo; +use codex_shell_command::shell_detect::DetectedShell; + +/// Provides environment metadata from either a local environment or a remote exec-server. +pub(crate) trait EnvironmentInfoProvider: Send + Sync { + fn info(&self) -> BoxFuture<'_, Result>; +} + +pub(crate) struct LocalEnvironmentInfoProvider; + +impl EnvironmentInfoProvider for LocalEnvironmentInfoProvider { + fn info(&self) -> BoxFuture<'_, Result> { + std::future::ready(Ok(EnvironmentInfo::local())).boxed() + } +} + +pub(crate) struct RemoteEnvironmentInfoProvider { + client: LazyRemoteExecServerClient, +} + +impl RemoteEnvironmentInfoProvider { + pub(crate) fn new(client: LazyRemoteExecServerClient) -> Self { + Self { client } + } +} + +impl EnvironmentInfoProvider for RemoteEnvironmentInfoProvider { + fn info(&self) -> BoxFuture<'_, Result> { + async move { self.client.environment_info().await }.boxed() + } +} + +impl EnvironmentInfo { + pub(crate) fn local() -> Self { + Self { + shell: codex_shell_command::shell_detect::default_user_shell().into(), + } + } +} + +impl From for ShellInfo { + fn from(shell: DetectedShell) -> Self { + Self { + name: shell.name().to_string(), + path: shell.shell_path.to_string_lossy().into_owned(), + } + } +} diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index 378d12e632..9308881f6d 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -3,6 +3,7 @@ mod client_api; mod client_transport; mod connection; mod environment; +mod environment_info; mod environment_path; mod environment_provider; mod environment_toml; diff --git a/codex-rs/exec-server/src/relay.rs b/codex-rs/exec-server/src/relay.rs index ae07d7652c..648a223a43 100644 --- a/codex-rs/exec-server/src/relay.rs +++ b/codex-rs/exec-server/src/relay.rs @@ -115,7 +115,7 @@ impl RelayMessageFrame { Some(relay_message_frame::Body::Data(data)) => data.payload, _ => Vec::new(), }; - serde_json::from_slice(&payload).map_err(ExecServerError::Json) + serde_json::from_slice(&payload).map_err(ExecServerError::from) } fn into_reset_reason(self) -> Option { @@ -138,7 +138,7 @@ fn decode_relay_message_frame(payload: &[u8]) -> Result Result, ExecServerError> { - serde_json::to_vec(message).map_err(ExecServerError::Json) + serde_json::to_vec(message).map_err(ExecServerError::from) } pub(crate) fn harness_connection_from_websocket(