diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index c80268e2f9..257a32733b 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -4,6 +4,7 @@ use crate::elicitation::ElicitationRegistration; use crate::session::Codex; use crate::session::SessionSettingsUpdate; use crate::session::SteerInputError; +use codex_exec_server::SelectedCapabilityRootsStatus; use codex_features::Feature; use codex_otel::SessionTelemetry; use codex_protocol::ThreadId; @@ -629,6 +630,18 @@ impl CodexThread { self.codex.thread_environment_selections().await } + /// Passively inspects the selected capability roots whose environments are ready now. + pub fn inspect_selected_capability_roots(&self) -> SelectedCapabilityRootsStatus { + self.codex + .session + .services + .turn_environments + .environment_manager() + .inspect_selected_capability_roots( + &self.codex.session.services.selected_capability_roots, + ) + } + pub async fn read_mcp_resource( &self, server: &str, diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index d83765d459..73e49b2f54 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -221,9 +221,16 @@ enum ConnectionStatus { Failed(String), } +#[derive(Clone, Copy)] +enum RecoveryPolicy { + Wait, + FailFast, +} + #[derive(Clone)] pub struct ExecServerClient { inner: Arc, + recovery_policy: RecoveryPolicy, } struct ActiveProcessStart { @@ -242,6 +249,7 @@ type ConnectionAttempt = OnceCell; #[derive(Clone)] pub(crate) struct LazyRemoteExecServerClient { transport_params: ExecServerTransportParams, + recovery_policy: RecoveryPolicy, // Saves the first startup result so callers share it and failures remain final. startup: Arc, // The latest successful client, replaced whenever reconnecting succeeds. @@ -253,6 +261,7 @@ impl LazyRemoteExecServerClient { pub(crate) fn new(transport_params: ExecServerTransportParams) -> Self { Self { transport_params, + recovery_policy: RecoveryPolicy::Wait, startup: Arc::new(ConnectionAttempt::new()), current_client: Arc::new(StdMutex::new(None)), reconnect: Arc::new(StdMutex::new(None)), @@ -279,11 +288,45 @@ impl LazyRemoteExecServerClient { self.startup.get().is_some() } + pub(crate) fn readiness_result(&self) -> Option> { + if let Some(client) = self.cached_client() { + return client.readiness_result(); + } + self.startup.get().and_then(|result| match result { + Ok(client) => client.readiness_result(), + Err(error) => Some(Err(ExecServerError::ConnectionAttempt(Arc::clone(error)))), + }) + } + + pub(crate) fn fail_fast(&self) -> Self { + Self { + recovery_policy: RecoveryPolicy::FailFast, + ..self.clone() + } + } + pub(crate) async fn wait_until_ready(&self) -> Result<(), ExecServerError> { self.initial_client().await.map(drop) } pub(crate) async fn get(&self) -> Result { + if matches!(self.recovery_policy, RecoveryPolicy::FailFast) { + let client = match self.cached_client() { + Some(client) => client, + None => match self.startup.get() { + Some(Ok(client)) => client.clone(), + Some(Err(error)) => { + return Err(ExecServerError::ConnectionAttempt(Arc::clone(error))); + } + None => { + return Err(ExecServerError::Disconnected( + "exec-server environment is not ready".to_string(), + )); + } + }, + }; + return client.fail_fast(); + } if let Some(client) = self.connected_client() { return Ok(client); } @@ -459,11 +502,45 @@ pub enum ExecServerError { } impl ExecServerClient { + fn fail_fast(&self) -> Result { + self.rpc_client_without_recovery()?; + Ok(Self { + inner: Arc::clone(&self.inner), + recovery_policy: RecoveryPolicy::FailFast, + }) + } + + fn rpc_client_without_recovery(&self) -> Result, ExecServerError> { + let connection = self + .inner + .connection + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + match &connection.status { + ConnectionStatus::Connected(rpc_client) if !rpc_client.is_disconnected() => { + Ok(Arc::clone(rpc_client)) + } + ConnectionStatus::Connected(_) | ConnectionStatus::Recovering => Err( + ExecServerError::Disconnected("exec-server environment is recovering".to_string()), + ), + ConnectionStatus::Failed(message) => { + Err(ExecServerError::Disconnected(message.clone())) + } + } + } + + async fn rpc_client(&self) -> Result, ExecServerError> { + match self.recovery_policy { + RecoveryPolicy::Wait => self.inner.rpc_client().await, + RecoveryPolicy::FailFast => self.rpc_client_without_recovery(), + } + } + pub async fn initialize( &self, options: ExecServerClientConnectOptions, ) -> Result { - let rpc_client = self.inner.rpc_client().await?; + let rpc_client = self.rpc_client().await?; self.initialize_rpc(&rpc_client, options).await } @@ -514,7 +591,7 @@ impl ExecServerClient { } pub async fn environment_info(&self) -> Result { - let rpc_client = self.inner.rpc_client().await?; + let rpc_client = self.rpc_client().await?; map_rpc_call_result( rpc_client .call_with_timeout(ENVIRONMENT_INFO_METHOD, &(), ENVIRONMENT_INFO_TIMEOUT) @@ -653,7 +730,7 @@ impl ExecServerClient { params: ExecParams, ) -> Result { loop { - let rpc_client = self.inner.rpc_client().await?; + let rpc_client = self.rpc_client().await?; if !self.inner.begin_process_start(&rpc_client) { continue; } @@ -730,6 +807,23 @@ impl ExecServerClient { self.inner.is_failed() } + fn readiness_result(&self) -> Option> { + let connection = self + .inner + .connection + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + match &connection.status { + ConnectionStatus::Connected(rpc_client) if !rpc_client.is_disconnected() => { + Some(Ok(())) + } + ConnectionStatus::Connected(_) | ConnectionStatus::Recovering => None, + ConnectionStatus::Failed(message) => { + Some(Err(ExecServerError::Disconnected(message.clone()))) + } + } + } + pub(crate) async fn connect( connection: JsonRpcConnection, options: ExecServerClientConnectOptions, @@ -761,7 +855,10 @@ impl ExecServerClient { session_id, reconnect_strategy, }); - let client = Self { inner }; + let client = Self { + inner, + recovery_policy: RecoveryPolicy::Wait, + }; // An explicit resume can redirect notifications from running processes // before initialize returns. Drain them immediately so a burst cannot // fill the bounded event channel and block the initialize response. @@ -775,7 +872,7 @@ impl ExecServerClient { P: serde::Serialize, T: serde::de::DeserializeOwned, { - let rpc_client = self.inner.rpc_client().await?; + let rpc_client = self.rpc_client().await?; self.call_rpc(&rpc_client, method, params).await } diff --git a/codex-rs/exec-server/src/client/rpc_http_client.rs b/codex-rs/exec-server/src/client/rpc_http_client.rs index d243a8d42f..c48ff8d29f 100644 --- a/codex-rs/exec-server/src/client/rpc_http_client.rs +++ b/codex-rs/exec-server/src/client/rpc_http_client.rs @@ -42,7 +42,7 @@ impl ExecServerClient { &self, mut params: HttpRequestParams, ) -> Result<(HttpRequestResponse, HttpResponseBodyStream), ExecServerError> { - let rpc_client = self.inner.rpc_client().await?; + let rpc_client = self.rpc_client().await?; params.stream_response = true; let request_id = self.inner.next_http_body_stream_request_id(); params.request_id = request_id.clone(); diff --git a/codex-rs/exec-server/src/client_recovery.rs b/codex-rs/exec-server/src/client_recovery.rs index 562a845cf1..a6075f3739 100644 --- a/codex-rs/exec-server/src/client_recovery.rs +++ b/codex-rs/exec-server/src/client_recovery.rs @@ -15,6 +15,7 @@ use super::ExecServerClient; use super::ExecServerError; use super::Inner; use super::OrderedSessionEvents; +use super::RecoveryPolicy; use super::SessionState; use super::disconnected_message; use super::fail_all_in_flight_work; @@ -437,6 +438,7 @@ impl Inner { let rpc_client = Arc::new(rpc_client); let client = ExecServerClient { inner: Arc::clone(self), + recovery_policy: RecoveryPolicy::Wait, }; // Resuming a session redirects notifications from its running processes // to this connection during initialize. Drain them immediately so a diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index 66f7381103..d070be6786 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -608,6 +608,14 @@ impl Environment { } } + /// Returns whether the environment can serve a request without waiting or reconnecting. + pub(crate) fn readiness_result(&self) -> Option> { + match &self.remote_client { + Some(client) => client.readiness_result(), + None => Some(Ok(())), + } + } + pub fn get_exec_backend(&self) -> Arc { Arc::clone(&self.exec_backend) } @@ -619,6 +627,14 @@ impl Environment { pub fn get_filesystem(&self) -> Arc { Arc::clone(&self.filesystem) } + + /// Returns a filesystem view that fails instead of starting or waiting for a connection. + pub fn get_filesystem_without_reconnect(&self) -> Arc { + match &self.remote_client { + Some(client) => Arc::new(RemoteFileSystem::new(client.fail_fast())), + None => Arc::clone(&self.filesystem), + } + } } #[cfg(test)] @@ -1083,6 +1099,72 @@ mod tests { assert!(environment.startup_finished()); } + #[tokio::test] + async fn selected_capability_inspection_keeps_stdio_environment_lazy() { + use codex_protocol::capabilities::CapabilityRootLocation; + use codex_protocol::capabilities::SelectedCapabilityRoot; + + let environment = Environment::remote_with_transport( + ExecServerTransportParams::StdioCommand { + command: StdioExecServerCommand { + program: "codex-missing-exec-server-for-test".to_string(), + args: Vec::new(), + env: HashMap::new(), + cwd: None, + }, + initialize_timeout: Duration::from_secs(1), + }, + /*local_runtime_paths*/ None, + ); + let manager = EnvironmentManager::from_snapshot( + EnvironmentProviderSnapshot { + environments: vec![("stdio".to_string(), environment)], + default: EnvironmentDefault::Disabled, + include_local: false, + }, + /*local_runtime_paths*/ None, + ) + .expect("environment manager"); + let environment = manager.get_environment("stdio").expect("stdio environment"); + let selected_root = SelectedCapabilityRoot { + id: "demo@1".to_string(), + location: CapabilityRootLocation::Environment { + environment_id: "stdio".to_string(), + path: PathUri::parse("file:///plugins/demo").expect("plugin path URI"), + }, + }; + + let status = + manager.inspect_selected_capability_roots(std::slice::from_ref(&selected_root)); + assert!(status.ready_roots.is_empty()); + assert_eq!(status.warnings, Vec::::new()); + assert!(!environment.startup_finished()); + + let missing_root = SelectedCapabilityRoot { + id: "missing@1".to_string(), + location: CapabilityRootLocation::Environment { + environment_id: "missing".to_string(), + path: PathUri::parse("file:///plugins/missing").expect("missing plugin path URI"), + }, + }; + let status = manager.inspect_selected_capability_roots(&[missing_root]); + assert!(status.ready_roots.is_empty()); + assert_eq!( + status.warnings, + vec![ + "selected capability root `missing@1` references unavailable environment `missing`" + .to_string() + ] + ); + + assert!(environment.wait_until_ready().await.is_err()); + + let status = manager.inspect_selected_capability_roots(&[selected_root]); + assert!(status.ready_roots.is_empty()); + assert_eq!(status.warnings.len(), 1); + assert!(status.warnings[0].contains("environment `stdio` is unavailable")); + } + #[tokio::test] async fn replacing_environment_stops_its_startup_task() { let first_listener = TcpListener::bind("127.0.0.1:0") diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index 882c83230c..c2c8cc5194 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -148,6 +148,7 @@ pub use protocol::WriteStatus; pub use remote::RemoteEnvironmentConfig; pub use remote::run_remote_environment; pub use resolved_capability::ResolvedSelectedCapabilityRoot; +pub use resolved_capability::SelectedCapabilityRootsStatus; pub use runtime_paths::ExecServerRuntimePaths; pub use server::DEFAULT_LISTEN_URL; pub use server::ExecServerListenUrlParseError; diff --git a/codex-rs/exec-server/src/resolved_capability.rs b/codex-rs/exec-server/src/resolved_capability.rs index 2902cf9eef..f1fa92c2f0 100644 --- a/codex-rs/exec-server/src/resolved_capability.rs +++ b/codex-rs/exec-server/src/resolved_capability.rs @@ -20,6 +20,15 @@ pub struct ResolvedSelectedCapabilityRoot { environment: Arc, } +/// A passive view of selected capability roots and unavailable environments. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct SelectedCapabilityRootsStatus { + /// Selected roots whose environments are ready. + pub ready_roots: Vec, + /// Missing environments and terminal connection failures. + pub warnings: Vec, +} + impl ResolvedSelectedCapabilityRoot { pub fn selected_root(&self) -> &SelectedCapabilityRoot { &self.selected_root @@ -35,6 +44,74 @@ impl ResolvedSelectedCapabilityRoot { } impl EnvironmentManager { + /// Inspects selected roots without starting or waiting for an environment. + /// + /// Starting or recovering environments are omitted. Missing environments and terminal + /// connection failures are returned as warnings so read-only catalog clients can distinguish + /// them from an empty catalog. + /// + /// Environment IDs are stable identities, so callers can safely resolve a returned root by + /// ID when they read it. + pub fn inspect_selected_capability_roots( + &self, + selected_roots: &[SelectedCapabilityRoot], + ) -> SelectedCapabilityRootsStatus { + let (candidates, mut warnings) = { + let environments = self + .environments + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let mut candidates = Vec::with_capacity(selected_roots.len()); + let mut warnings = Vec::new(); + for selected_root in selected_roots { + let CapabilityRootLocation::Environment { environment_id, .. } = + &selected_root.location; + let Some(environment) = environments.get(environment_id) else { + warnings.push(format!( + "selected capability root `{}` references unavailable environment `{environment_id}`", + selected_root.id + )); + continue; + }; + candidates.push((selected_root.clone(), Arc::clone(environment))); + } + (candidates, warnings) + }; + let mut readiness = HashMap::new(); + for (selected_root, environment) in &candidates { + let CapabilityRootLocation::Environment { environment_id, .. } = + &selected_root.location; + if readiness.contains_key(environment_id) { + continue; + } + let ready = match environment.readiness_result() { + Some(Ok(())) => true, + Some(Err(error)) => { + warnings.push(format!( + "selected capability environment `{environment_id}` is unavailable: {error}" + )); + false + } + None => false, + }; + readiness.insert(environment_id.clone(), ready); + } + + let ready_roots = candidates + .into_iter() + .filter(|(selected_root, _)| { + let CapabilityRootLocation::Environment { environment_id, .. } = + &selected_root.location; + readiness.get(environment_id).copied().unwrap_or(false) + }) + .map(|(selected_root, _)| selected_root) + .collect(); + SelectedCapabilityRootsStatus { + ready_roots, + warnings, + } + } + /// Resolves selected roots whose stable environments are ready for the current model step. /// /// Environment identity comes from the selected root's stable environment ID. A ready diff --git a/codex-rs/exec-server/tests/environment.rs b/codex-rs/exec-server/tests/environment.rs new file mode 100644 index 0000000000..4b97361ac2 --- /dev/null +++ b/codex-rs/exec-server/tests/environment.rs @@ -0,0 +1,84 @@ +mod common; + +use std::time::Duration; + +use anyhow::Context; +use codex_exec_server::EnvironmentManager; +use codex_exec_server::REMOTE_ENVIRONMENT_ID; +use codex_exec_server::SelectedCapabilityRootsStatus; +use codex_protocol::capabilities::CapabilityRootLocation; +use codex_protocol::capabilities::SelectedCapabilityRoot; +use codex_utils_path_uri::PathUri; +use common::exec_server::exec_server; +use pretty_assertions::assert_eq; +use tokio::time::sleep; +use tokio::time::timeout; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial_test::serial(remote_exec_server)] +async fn selected_capability_inspection_tracks_connection_recovery() -> anyhow::Result<()> { + let server = exec_server().await?; + let mut proxy = server.disconnectable_websocket_proxy().await?; + let manager = EnvironmentManager::create_for_tests( + Some(proxy.websocket_url().to_string()), + /*local_runtime_paths*/ None, + ) + .await; + let environment = manager + .default_environment() + .context("remote environment")?; + environment.info().await?; + + let skill_root_path = PathUri::parse("file:///plugins/demo")?; + let selected_root = SelectedCapabilityRoot { + id: "demo@1".to_string(), + location: CapabilityRootLocation::Environment { + environment_id: REMOTE_ENVIRONMENT_ID.to_string(), + path: skill_root_path.clone(), + }, + }; + assert_eq!( + manager.inspect_selected_capability_roots(std::slice::from_ref(&selected_root)), + SelectedCapabilityRootsStatus { + ready_roots: vec![selected_root.clone()], + warnings: Vec::new(), + } + ); + let file_system = environment.get_filesystem_without_reconnect(); + + proxy.pause_and_disconnect().await?; + assert_eq!( + manager.inspect_selected_capability_roots(std::slice::from_ref(&selected_root)), + SelectedCapabilityRootsStatus::default() + ); + let read_result = timeout( + Duration::from_secs(1), + file_system.read_directory(&skill_root_path, /*sandbox*/ None), + ) + .await + .context("passive filesystem read waited for recovery")?; + assert!(read_result.is_err()); + + proxy.resume()?; + let recovered_status = timeout(Duration::from_secs(5), async { + loop { + let status = + manager.inspect_selected_capability_roots(std::slice::from_ref(&selected_root)); + if !status.ready_roots.is_empty() { + break status; + } + sleep(Duration::from_millis(10)).await; + } + }) + .await + .context("environment did not recover")?; + assert_eq!( + recovered_status, + SelectedCapabilityRootsStatus { + ready_roots: vec![selected_root], + warnings: Vec::new(), + } + ); + + Ok(()) +}