From 79177c3e2071204f3ab080e7c2d36f17312f16a3 Mon Sep 17 00:00:00 2001 From: jif Date: Wed, 15 Jul 2026 20:58:44 +0000 Subject: [PATCH] Propagate deferred environment capability roots to MCP (#33427) ## What changed - Let deferred environments provide selected capability roots with their ready signal. - Validate that those roots have unique, non-empty IDs, belong to the registering environment, and stay within the root limit. - Include roots from ready turn environments when resolving MCP contributions, and refresh the MCP runtime when the selected root set changes. - Expose the exact ready root set to MCP contributors so executor plugins become available with their environment. ## Testing - Cover ready-root propagation, validation failures, replacement isolation, reconnection, and MCP plugin availability refresh. GitOrigin-RevId: ec3498aab1164824025094e96a9b1063b7b731ad --- codex-rs/core/src/mcp.rs | 5 +- codex-rs/core/src/session/mcp.rs | 98 +++++++++++++------ codex-rs/core/src/session/mcp_runtime.rs | 11 ++- codex-rs/core/src/session/session.rs | 4 +- codex-rs/core/src/session/tests.rs | 77 ++++++++++++--- codex-rs/core/src/state/service.rs | 8 +- codex-rs/core/src/thread_manager_tests.rs | 4 +- codex-rs/core/tests/suite/remote_env.rs | 3 +- codex-rs/exec-server/src/environment.rs | 94 ++++++++++++++++-- codex-rs/exec-server/src/lib.rs | 2 + .../exec-server/tests/deferred_environment.rs | 62 +++++++++++- codex-rs/exec-server/tests/relay.rs | 22 ++++- .../ext/extension-api/src/contributors/mcp.rs | 20 ++-- codex-rs/ext/mcp/src/executor_plugin.rs | 18 +--- codex-rs/ext/mcp/src/lib_tests.rs | 2 +- codex-rs/ext/mcp/tests/executor_plugin_mcp.rs | 9 +- 16 files changed, 333 insertions(+), 106 deletions(-) diff --git a/codex-rs/core/src/mcp.rs b/codex-rs/core/src/mcp.rs index 399a0912a3..b33111373a 100644 --- a/codex-rs/core/src/mcp.rs +++ b/codex-rs/core/src/mcp.rs @@ -24,6 +24,7 @@ use codex_mcp::codex_apps_mcp_server_config; use codex_mcp::configured_mcp_servers; use codex_mcp::effective_mcp_servers; use codex_plugin::AppConnectorId; +use codex_protocol::capabilities::SelectedCapabilityRoot; const LEGACY_CODEX_APPS_REGISTRATION_ID: &str = "legacy_codex_apps"; @@ -107,7 +108,7 @@ impl McpManager { thread_init: &ExtensionDataInit, thread_store: &ExtensionData, originator: &str, - available_environment_ids: &[String], + ready_selected_capability_roots: &[SelectedCapabilityRoot], ) -> McpRuntimeProjection { self.runtime_config_with_context( McpServerContributionContext::for_step( @@ -115,7 +116,7 @@ impl McpManager { thread_init, thread_store, originator, - available_environment_ids, + ready_selected_capability_roots, ), Some(originator), ) diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index 61f967d955..1f8c752161 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -1,10 +1,12 @@ use super::*; use crate::mcp::McpRuntimeProjection; +use codex_exec_server::MAX_SELECTED_CAPABILITY_ROOTS; use codex_exec_server::ResolvedSelectedCapabilityRoot; use codex_mcp::ElicitationReviewRequest; use codex_mcp::ElicitationReviewer; use codex_mcp::ElicitationReviewerHandle; use codex_protocol::capabilities::CapabilityRootLocation; +use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_protocol::config_types::ApprovalsReviewer; use codex_protocol::mcp_approval_meta::APPROVAL_KIND_KEY as MCP_ELICITATION_APPROVAL_KIND_KEY; use codex_protocol::mcp_approval_meta::APPROVAL_KIND_MCP_TOOL_CALL as MCP_ELICITATION_APPROVAL_KIND_MCP_TOOL_CALL; @@ -83,8 +85,8 @@ impl Session { let selected_capability_roots = self .resolve_selected_capability_roots_for_step(&environments) .await; - let available_environment_ids = - Self::available_selected_environment_ids(&selected_capability_roots); + let ready_selected_capability_roots = + Self::ready_selected_capability_roots(&selected_capability_roots); self.services .mcp_manager .runtime_config_for_step( @@ -92,7 +94,7 @@ impl Session { &self.services.mcp_thread_init, &self.services.thread_extension_data, &originator, - &available_environment_ids, + &ready_selected_capability_roots, ) .await .config @@ -116,16 +118,18 @@ impl Session { environments: &TurnEnvironmentSnapshot, selected_capability_roots: &[ResolvedSelectedCapabilityRoot], ) -> Arc { + let ready_selected_capability_roots = + Self::ready_selected_capability_roots(selected_capability_roots); let available_environment_ids = Self::available_selected_environment_ids(selected_capability_roots); let current = self.services.latest_mcp_runtime(); - if current.available_environment_ids() == available_environment_ids { + if current.ready_selected_capability_roots() == ready_selected_capability_roots { return current; } let _guard = self.services.mcp_projection_lock.lock().await; let current = self.services.latest_mcp_runtime(); - if current.available_environment_ids() == available_environment_ids { + if current.ready_selected_capability_roots() == ready_selected_capability_roots { return current; } let mcp_projection = self @@ -136,7 +140,7 @@ impl Session { &self.services.mcp_thread_init, &self.services.thread_extension_data, &turn_context.originator, - &available_environment_ids, + &ready_selected_capability_roots, ) .await; let mcp_config = &mcp_projection.config; @@ -146,8 +150,13 @@ impl Session { .values() .any(|server| { let was_available = current - .available_environment_ids() - .contains(&server.environment_id); + .ready_selected_capability_roots() + .iter() + .any(|root| { + let CapabilityRootLocation::Environment { environment_id, .. } = + &root.location; + environment_id == &server.environment_id + }); let is_available = available_environment_ids.contains(&server.environment_id); server.enabled && was_available != is_available }); @@ -158,15 +167,15 @@ impl Session { .has_same_servers(&mcp_config.mcp_server_catalog) && current.config().connector_snapshot == mcp_config.connector_snapshot { - // Availability is only an input to the MCP projection. When that input changes but - // the projected servers and connectors do not, advance the input key without + // Selected roots are only an input to the MCP projection. When they change but the + // projected servers and connectors do not, advance the input key without // replacing the live manager and restarting its processes. let runtime = Arc::new(McpRuntimeSnapshot::new( Arc::new(current.config().clone()), mcp_projection.plugins_available, current.manager_arc(), current.runtime_context().clone(), - available_environment_ids, + ready_selected_capability_roots, )); self.services.mcp_runtime.store(Some(Arc::clone(&runtime))); return runtime; @@ -175,7 +184,7 @@ impl Session { turn_context, mcp_projection, environments, - &available_environment_ids, + &ready_selected_capability_roots, Some(self.mcp_elicitation_reviewer()), ) .await @@ -185,11 +194,37 @@ impl Session { &self, environments: &TurnEnvironmentSnapshot, ) -> Vec { + let mut seen_root_ids = self + .services + .selected_capability_roots + .iter() + .map(|root| root.id.as_str()) + .collect::>(); + let mut selected_capability_roots = self.services.selected_capability_roots.clone(); + let mut ready_environment_root_count = 0; + for root in environments + .turn_environments + .iter() + .flat_map(|environment| environment.environment.selected_capability_roots()) + { + if !seen_root_ids.insert(root.id.as_str()) { + continue; + } + if ready_environment_root_count == MAX_SELECTED_CAPABILITY_ROOTS { + tracing::warn!( + max_root_count = MAX_SELECTED_CAPABILITY_ROOTS, + "ignoring excess selected capability roots from ready environments" + ); + break; + } + selected_capability_roots.push(root.clone()); + ready_environment_root_count += 1; + } self.services .turn_environments .environment_manager() .resolve_selected_capability_roots( - &self.services.selected_capability_roots, + &selected_capability_roots, &environments.captured_environments(), ) .await @@ -326,7 +361,7 @@ impl Session { turn_context: &TurnContext, mcp_projection: McpRuntimeProjection, environments: &TurnEnvironmentSnapshot, - available_environment_ids: &[String], + ready_selected_capability_roots: &[SelectedCapabilityRoot], elicitation_reviewer: Option, ) -> Arc { let auth = self.services.auth_manager.auth().await; @@ -394,7 +429,7 @@ impl Session { mcp_config, plugins_available, mcp_runtime_context, - available_environment_ids.to_vec(), + ready_selected_capability_roots.to_vec(), refreshed_manager, ) } @@ -460,11 +495,9 @@ impl Session { } let _guard = self.services.mcp_projection_lock.lock().await; - let available_environment_ids = self - .services - .latest_mcp_runtime() - .available_environment_ids() - .to_vec(); + let current_runtime = self.services.latest_mcp_runtime(); + let ready_selected_capability_roots = + current_runtime.ready_selected_capability_roots().to_vec(); let mut mcp_projection = self .services .mcp_manager @@ -473,7 +506,7 @@ impl Session { &self.services.mcp_thread_init, &self.services.thread_extension_data, &turn_context.originator, - &available_environment_ids, + &ready_selected_capability_roots, ) .await; mcp_projection.config.mcp_server_catalog = mcp_projection @@ -484,7 +517,7 @@ impl Session { turn_context, mcp_projection, &turn_context.environments, - &available_environment_ids, + &ready_selected_capability_roots, elicitation_reviewer, ) .await; @@ -529,11 +562,9 @@ impl Session { elicitation_reviewer: Option, ) { let _guard = self.services.mcp_projection_lock.lock().await; - let available_environment_ids = self - .services - .latest_mcp_runtime() - .available_environment_ids() - .to_vec(); + let current_runtime = self.services.latest_mcp_runtime(); + let ready_selected_capability_roots = + current_runtime.ready_selected_capability_roots().to_vec(); let mcp_projection = self .services .mcp_manager @@ -542,14 +573,14 @@ impl Session { &self.services.mcp_thread_init, &self.services.thread_extension_data, &turn_context.originator, - &available_environment_ids, + &ready_selected_capability_roots, ) .await; self.refresh_mcp_servers_inner( turn_context, mcp_projection, &turn_context.environments, - &available_environment_ids, + &ready_selected_capability_roots, elicitation_reviewer, ) .await; @@ -569,6 +600,15 @@ impl Session { available } + fn ready_selected_capability_roots( + selected_capability_roots: &[ResolvedSelectedCapabilityRoot], + ) -> Vec { + selected_capability_roots + .iter() + .map(|root| root.selected_root().clone()) + .collect() + } + #[cfg(test)] pub(crate) async fn mcp_startup_cancellation_token(&self) -> CancellationToken { self.services diff --git a/codex-rs/core/src/session/mcp_runtime.rs b/codex-rs/core/src/session/mcp_runtime.rs index 5e7aeffcee..a89921ff15 100644 --- a/codex-rs/core/src/session/mcp_runtime.rs +++ b/codex-rs/core/src/session/mcp_runtime.rs @@ -4,6 +4,7 @@ use std::sync::Arc; use codex_mcp::McpConfig; use codex_mcp::McpConnectionManager; use codex_mcp::McpRuntimeContext; +use codex_protocol::capabilities::SelectedCapabilityRoot; /// MCP config, plugin availability, exact environment bindings, and manager for one request. pub struct McpRuntimeSnapshot { @@ -11,7 +12,7 @@ pub struct McpRuntimeSnapshot { plugins_available: bool, manager: Arc, runtime_context: McpRuntimeContext, - available_environment_ids: Vec, + ready_selected_capability_roots: Vec, } impl McpRuntimeSnapshot { @@ -20,14 +21,14 @@ impl McpRuntimeSnapshot { plugins_available: bool, manager: Arc, runtime_context: McpRuntimeContext, - available_environment_ids: Vec, + ready_selected_capability_roots: Vec, ) -> Self { Self { config, plugins_available, manager, runtime_context, - available_environment_ids, + ready_selected_capability_roots, } } @@ -51,8 +52,8 @@ impl McpRuntimeSnapshot { &self.runtime_context } - pub(crate) fn available_environment_ids(&self) -> &[String] { - &self.available_environment_ids + pub(crate) fn ready_selected_capability_roots(&self) -> &[SelectedCapabilityRoot] { + &self.ready_selected_capability_roots } #[cfg(test)] diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 0726e2e3b4..5c5167f9d2 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -685,7 +685,7 @@ impl Session { mcp_thread_init_for_startup, thread_extension_data_for_mcp, &mcp_originator, - /*available_environment_ids*/ &[], + /*ready_selected_capability_roots*/ &[], ) .await; let mcp_config = &mcp_projection.config; @@ -1234,7 +1234,7 @@ impl Session { Arc::new(mcp_projection.config), mcp_projection.plugins_available, mcp_runtime_context, - /*available_environment_ids*/ Vec::new(), + /*ready_selected_capability_roots*/ Vec::new(), mcp_connection_manager, ) .await?; diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index fbea8f551e..e453a545bc 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -7686,9 +7686,26 @@ async fn refresh_mcp_servers_keeps_the_previous_runtime_alive() { } #[tokio::test] -async fn plugin_availability_change_reuses_the_mcp_manager() { +async fn deferred_environment_roots_refresh_plugin_availability() { struct ReadyPluginContributor; + struct PendingNoiseConnectProvider; + + impl codex_exec_server::NoiseRendezvousConnectProvider for PendingNoiseConnectProvider { + fn connect_bundle( + &self, + _: codex_exec_server::NoiseChannelPublicKey, + ) -> futures::future::BoxFuture< + '_, + Result< + codex_exec_server::NoiseRendezvousConnectBundle, + codex_exec_server::ExecServerError, + >, + > { + Box::pin(futures::future::pending()) + } + } + impl codex_extension_api::McpServerContributor for ReadyPluginContributor { fn id(&self) -> &'static str { "ready_plugin_test" @@ -7701,8 +7718,8 @@ async fn plugin_availability_change_reuses_the_mcp_manager() { { Box::pin(async move { let available = context - .available_environment_ids() - .is_some_and(|ids| ids.iter().any(|id| id == "executor")); + .ready_selected_capability_roots() + .is_some_and(|roots| roots.iter().any(|root| root.id == "skill-only")); available .then( || codex_extension_api::McpServerContribution::SelectedPluginPackage { @@ -7746,23 +7763,53 @@ async fn plugin_availability_change_reuses_the_mcp_manager() { .expect("selected capability root URI"), }, }; - let environment = Arc::clone( - &turn_context - .environments - .primary() - .expect("ready test environment") - .environment, + let environment_manager = session.services.turn_environments.environment_manager(); + let registration = environment_manager + .register_deferred_noise_environment( + "executor".to_string(), + Arc::new(PendingNoiseConnectProvider), + ) + .expect("register deferred environment"); + let environment = environment_manager + .get_environment("executor") + .expect("deferred environment"); + assert!(environment.selected_capability_roots().is_empty()); + registration + .complete(Ok(codex_exec_server::EnvironmentReadyInfo { + selected_capability_roots: vec![selected_root.clone()], + })) + .expect("complete deferred environment"); + assert_eq!( + environment.selected_capability_roots(), + std::slice::from_ref(&selected_root) ); - let captured_environments = HashMap::from([("executor".to_string(), Some(environment))]); + let local_environment = turn_context + .environments + .primary() + .expect("ready local environment"); + let environments = TurnEnvironmentSnapshot { + turn_environments: vec![TurnEnvironment::new( + "executor".to_string(), + environment, + local_environment.cwd().clone(), + local_environment.workspace_roots().to_vec(), + local_environment.shell.clone(), + )], + starting: Vec::new(), + }; let resolved_roots = session - .services - .turn_environments - .environment_manager() - .resolve_selected_capability_roots(&[selected_root], &captured_environments) + .resolve_selected_capability_roots_for_step(&environments) .await; + assert_eq!( + resolved_roots + .iter() + .map(|root| root.selected_root().clone()) + .collect::>(), + vec![selected_root] + ); let new_runtime = session - .mcp_runtime_for_step(&turn_context, &turn_context.environments, &resolved_roots) + .mcp_runtime_for_step(&turn_context, &environments, &resolved_roots) .await; assert!(!old_runtime.plugins_available()); diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index 43f411e31b..64aee2d497 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -111,14 +111,14 @@ impl SessionServices { config: Arc, plugins_available: bool, runtime_context: McpRuntimeContext, - available_environment_ids: Vec, + ready_selected_capability_roots: Vec, manager: McpConnectionManager, ) -> Result<()> { let runtime = self.publish_mcp_runtime( config, plugins_available, runtime_context, - available_environment_ids, + ready_selected_capability_roots, manager, ); runtime.manager().validate_required_servers().await @@ -129,7 +129,7 @@ impl SessionServices { config: Arc, plugins_available: bool, runtime_context: McpRuntimeContext, - available_environment_ids: Vec, + ready_selected_capability_roots: Vec, manager: McpConnectionManager, ) -> Arc { let manager = Arc::new(manager); @@ -141,7 +141,7 @@ impl SessionServices { plugins_available, manager, runtime_context, - available_environment_ids, + ready_selected_capability_roots, )); self.mcp_runtime.store(Some(Arc::clone(&runtime))); runtime diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index 09b8b4f994..df2f00e839 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -668,7 +668,7 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() &first_session.services.mcp_thread_init, &first_session.services.thread_extension_data, &first_originator, - /*available_environment_ids*/ &[], + /*ready_selected_capability_roots*/ &[], ) .await; let second_session = &second_thread.thread.session; @@ -681,7 +681,7 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() &second_session.services.mcp_thread_init, &second_session.services.thread_extension_data, &second_originator, - /*available_environment_ids*/ &[], + /*ready_selected_capability_roots*/ &[], ) .await; diff --git a/codex-rs/core/tests/suite/remote_env.rs b/codex-rs/core/tests/suite/remote_env.rs index abaace0c8d..e19bab7eac 100644 --- a/codex-rs/core/tests/suite/remote_env.rs +++ b/codex-rs/core/tests/suite/remote_env.rs @@ -7,6 +7,7 @@ use codex_core::compact::SUMMARIZATION_PROMPT; use codex_core::config::Constrained; use codex_exec_server::CopyOptions; use codex_exec_server::CreateDirectoryOptions; +use codex_exec_server::EnvironmentReadyInfo; use codex_exec_server::ExecServerError; use codex_exec_server::FileSystemSandboxContext; use codex_exec_server::LOCAL_ENVIRONMENT_ID; @@ -585,7 +586,7 @@ async fn deferred_executor_starts_noise_connection_after_registration() -> Resul wait_for_response_request_count(&response_mock, /*expected_count*/ 1).await; assert_eq!(response_mock.requests().len(), 1); assert_eq!(provider.calls.load(Ordering::Relaxed), 0); - registration.complete(Ok(()))?; + registration.complete(Ok(EnvironmentReadyInfo::default()))?; wait_for_event(&test.codex, |event| { matches!(event, EventMsg::TurnComplete(_)) }) diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index 7928a39bb9..58cf5193cf 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -1,8 +1,12 @@ use std::collections::HashMap; +use std::collections::HashSet; use std::sync::Arc; use std::sync::Mutex; +use std::sync::OnceLock; use std::sync::RwLock; +use codex_protocol::capabilities::CapabilityRootLocation; +use codex_protocol::capabilities::SelectedCapabilityRoot; use futures::FutureExt; use crate::ExecServerError; @@ -72,9 +76,23 @@ pub struct EnvironmentManager { local_runtime_paths: Option, } +/// Information supplied by the environment owner when a deferred environment is ready. +#[derive(Clone, Debug, Default, Eq, PartialEq)] +pub struct EnvironmentReadyInfo { + /// Ordered capability roots selected for this environment. + pub selected_capability_roots: Vec, +} + /// The one-shot capability to complete a deferred environment registration. #[must_use = "the deferred environment cannot connect until registration is completed"] -pub struct DeferredEnvironmentRegistration(oneshot::Sender>); +pub struct DeferredEnvironmentRegistration { + completion: oneshot::Sender>, + environment_id: String, + ready_info: Arc>, +} + +/// Maximum capability roots accepted from deferred environment ready information. +pub const MAX_SELECTED_CAPABILITY_ROOTS: usize = 256; pub const LOCAL_ENVIRONMENT_ID: &str = "local"; pub const REMOTE_ENVIRONMENT_ID: &str = "remote"; @@ -367,15 +385,22 @@ impl EnvironmentManager { validate_environment_id(&environment_id)?; let identity = noise_channel_identity()?; let (completion, readiness) = oneshot::channel(); - let environment = Arc::new(Environment::remote_with_transport( + let ready_info = Arc::new(OnceLock::new()); + let mut environment = Environment::remote_with_transport( ExecServerTransportParams::Deferred(Box::new(crate::client_api::Deferred { readiness: readiness.shared(), transport: ExecServerTransportParams::NoiseRendezvous { provider, identity }, })), self.local_runtime_paths.clone(), - )); - self.insert_environment(environment_id, environment); - Ok(DeferredEnvironmentRegistration(completion)) + ); + environment.ready_info = Some(Arc::clone(&ready_info)); + let environment = Arc::new(environment); + self.insert_environment(environment_id.clone(), environment); + Ok(DeferredEnvironmentRegistration { + completion, + environment_id, + ready_info, + }) } /// Adds or replaces a named remote environment that connects through an @@ -408,9 +433,48 @@ impl EnvironmentManager { } impl DeferredEnvironmentRegistration { - /// Completes provisioning with readiness or a terminal error message. - pub fn complete(self, result: Result<(), String>) -> Result<(), ExecServerError> { - self.0.send(result).map_err(|_| { + /// Completes provisioning with ready information or a terminal error message. + pub fn complete( + self, + result: Result, + ) -> Result<(), ExecServerError> { + let result = match result { + Ok(ready_info) => { + if ready_info.selected_capability_roots.len() > MAX_SELECTED_CAPABILITY_ROOTS { + let error = ExecServerError::Protocol(format!( + "environment ready info contains more than {MAX_SELECTED_CAPABILITY_ROOTS} selected capability roots" + )); + let _ = self.completion.send(Err(error.to_string())); + return Err(error); + } + let mut root_ids = + HashSet::with_capacity(ready_info.selected_capability_roots.len()); + for root in &ready_info.selected_capability_roots { + let CapabilityRootLocation::Environment { environment_id, .. } = &root.location; + if root.id.trim().is_empty() + || environment_id != &self.environment_id + || !root_ids.insert(root.id.as_str()) + { + let error = ExecServerError::Protocol(format!( + "selected capability roots must have unique non-empty IDs and belong to environment `{}`", + self.environment_id + )); + let _ = self.completion.send(Err(error.to_string())); + return Err(error); + } + } + if self.ready_info.set(ready_info).is_err() { + let error = ExecServerError::Protocol( + "deferred environment ready info was already set".to_string(), + ); + let _ = self.completion.send(Err(error.to_string())); + return Err(error); + } + Ok(()) + } + Err(message) => Err(message), + }; + self.completion.send(result).map_err(|_| { ExecServerError::Disconnected("deferred environment registration is inactive".into()) }) } @@ -499,6 +563,7 @@ fn optional_environment_value(name: &str) -> Option { #[derive(Clone)] pub struct Environment { remote_client: Option, + ready_info: Option>>, // Dropping the environment stops unfinished background startup work. startup_task: Arc>>>, exec_backend: Arc, @@ -512,6 +577,7 @@ impl Environment { pub fn default_for_tests() -> Self { Self { remote_client: None, + ready_info: None, startup_task: Arc::new(Mutex::new(None)), exec_backend: Arc::new(LocalProcess::default()), filesystem: Arc::new(LocalFileSystem::unsandboxed()), @@ -566,6 +632,7 @@ impl Environment { pub(crate) fn local(local_runtime_paths: ExecServerRuntimePaths) -> Self { Self { remote_client: None, + ready_info: None, startup_task: Arc::new(Mutex::new(None)), exec_backend: Arc::new(LocalProcess::with_local_runtime_paths( local_runtime_paths.clone(), @@ -602,6 +669,7 @@ impl Environment { Self { remote_client: Some(client.clone()), + ready_info: None, startup_task: Arc::new(Mutex::new(None)), exec_backend, filesystem, @@ -614,6 +682,16 @@ impl Environment { self.remote_client.is_some() } + /// Returns capability roots supplied with the deferred environment's ready signal. + pub fn selected_capability_roots(&self) -> &[SelectedCapabilityRoot] { + self.ready_info + .as_ref() + .and_then(|ready_info| ready_info.get()) + .map_or(&[], |ready_info| { + ready_info.selected_capability_roots.as_slice() + }) + } + /// Subscribes to the current connection state for this remote environment. pub fn subscribe_connection_state( &self, diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index ab0cf149ed..c192eebd1e 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -70,7 +70,9 @@ pub use environment::Environment; pub use environment::EnvironmentConnectionState; pub use environment::EnvironmentManager; pub use environment::EnvironmentObservedStatus; +pub use environment::EnvironmentReadyInfo; pub use environment::LOCAL_ENVIRONMENT_ID; +pub use environment::MAX_SELECTED_CAPABILITY_ROOTS; pub use environment::REMOTE_ENVIRONMENT_ID; pub use environment_provider::DefaultEnvironmentProvider; pub use environment_provider::EnvironmentProvider; diff --git a/codex-rs/exec-server/tests/deferred_environment.rs b/codex-rs/exec-server/tests/deferred_environment.rs index 8efff56225..11de320b3c 100644 --- a/codex-rs/exec-server/tests/deferred_environment.rs +++ b/codex-rs/exec-server/tests/deferred_environment.rs @@ -3,10 +3,14 @@ use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use codex_exec_server::EnvironmentManager; +use codex_exec_server::EnvironmentReadyInfo; use codex_exec_server::ExecServerError; use codex_exec_server::NoiseChannelPublicKey; use codex_exec_server::NoiseRendezvousConnectBundle; use codex_exec_server::NoiseRendezvousConnectProvider; +use codex_protocol::capabilities::CapabilityRootLocation; +use codex_protocol::capabilities::SelectedCapabilityRoot; +use codex_utils_path_uri::PathUri; use futures::FutureExt; use futures::future::BoxFuture; use futures::poll; @@ -38,6 +42,18 @@ impl NoiseRendezvousConnectProvider for FailingNoiseConnectProvider { } } +fn ready_info(root_id: &str, environment_id: &str) -> anyhow::Result { + Ok(EnvironmentReadyInfo { + selected_capability_roots: vec![SelectedCapabilityRoot { + id: root_id.to_string(), + location: CapabilityRootLocation::Environment { + environment_id: environment_id.to_string(), + path: PathUri::parse("file:///plugins/root")?, + }, + }], + }) +} + #[tokio::test] async fn deferred_environment_waits_before_connecting() -> anyhow::Result<()> { let manager = EnvironmentManager::without_environments(); @@ -52,8 +68,14 @@ async fn deferred_environment_waits_before_connecting() -> anyhow::Result<()> { assert!(poll!(&mut readiness).is_pending()); assert_eq!(provider.calls(), 0); + assert!(environment.selected_capability_roots().is_empty()); - registration.complete(Ok(()))?; + let ready_info = ready_info("selected-root", "tools")?; + registration.complete(Ok(ready_info.clone()))?; + assert_eq!( + environment.selected_capability_roots(), + ready_info.selected_capability_roots + ); let error = readiness.await.unwrap_err(); assert!(error.to_string().contains("test Noise provider called")); assert_eq!(provider.calls(), 1); @@ -94,6 +116,29 @@ async fn failure_and_dropped_registration_are_terminal() -> anyhow::Result<()> { Ok(()) } +#[tokio::test] +async fn invalid_ready_info_is_terminal() -> anyhow::Result<()> { + let manager = EnvironmentManager::without_environments(); + let provider = Arc::new(FailingNoiseConnectProvider::default()); + let registration = + manager.register_deferred_noise_environment("tools".to_string(), provider.clone())?; + let environment = manager.get_environment("tools").expect("environment"); + + let error = registration + .complete(Ok(ready_info("selected-root", "other")?)) + .unwrap_err(); + assert!(matches!(error, ExecServerError::Protocol(_))); + let readiness_error = environment.wait_until_ready().await.unwrap_err(); + assert!( + readiness_error + .to_string() + .contains("belong to environment") + ); + assert!(environment.selected_capability_roots().is_empty()); + assert_eq!(provider.calls(), 0); + Ok(()) +} + #[tokio::test] async fn late_completion_is_isolated_from_replacement() -> anyhow::Result<()> { let manager = EnvironmentManager::without_environments(); @@ -106,7 +151,13 @@ async fn late_completion_is_isolated_from_replacement() -> anyhow::Result<()> { .register_deferred_noise_environment("tools".to_string(), current_provider.clone())?; let current = manager.get_environment("tools").expect("current"); - old_registration.complete(Ok(()))?; + let old_ready_info = ready_info("old-root", "tools")?; + old_registration.complete(Ok(old_ready_info.clone()))?; + assert_eq!( + old_environment.selected_capability_roots(), + old_ready_info.selected_capability_roots + ); + assert!(current.selected_capability_roots().is_empty()); let old_error = old_environment.wait_until_ready().await.unwrap_err(); assert!(old_error.to_string().contains("test Noise provider called")); assert_eq!(old_provider.calls(), 1); @@ -114,7 +165,12 @@ async fn late_completion_is_isolated_from_replacement() -> anyhow::Result<()> { assert!(poll!(&mut current_readiness).is_pending()); assert_eq!(current_provider.calls(), 0); - current_registration.complete(Ok(()))?; + let current_ready_info = ready_info("current-root", "tools")?; + current_registration.complete(Ok(current_ready_info.clone()))?; + assert_eq!( + current.selected_capability_roots(), + current_ready_info.selected_capability_roots + ); let current_error = current_readiness.await.unwrap_err(); assert!( current_error diff --git a/codex-rs/exec-server/tests/relay.rs b/codex-rs/exec-server/tests/relay.rs index 50aeff73ff..e5c21d468c 100644 --- a/codex-rs/exec-server/tests/relay.rs +++ b/codex-rs/exec-server/tests/relay.rs @@ -17,6 +17,7 @@ use base64::engine::general_purpose::STANDARD; use codex_api::AuthProvider; use codex_exec_server::EnvironmentConnectionState; use codex_exec_server::EnvironmentManager; +use codex_exec_server::EnvironmentReadyInfo; use codex_exec_server::ExecParams; use codex_exec_server::ExecResponse; use codex_exec_server::ExecServerClient; @@ -30,6 +31,8 @@ use codex_exec_server::NoiseRendezvousConnectBundle; use codex_exec_server::NoiseRendezvousConnectProvider; use codex_exec_server::ProcessId; use codex_exec_server::RemoteEnvironmentConfig; +use codex_protocol::capabilities::CapabilityRootLocation; +use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_utils_path_uri::PathUri; use futures::SinkExt; use futures::StreamExt; @@ -164,7 +167,16 @@ async fn deferred_noise_environment_connects_and_reconnects_with_fresh_bundle() .context("remote environment connection state")?; assert_eq!(provider.calls(), 0); - registration.complete(Ok(()))?; + let selected_capability_roots = vec![SelectedCapabilityRoot { + id: "executor-plugin".to_string(), + location: CapabilityRootLocation::Environment { + environment_id: ENVIRONMENT_ID.to_string(), + path: PathUri::parse("file:///plugins/executor-plugin")?, + }, + }]; + registration.complete(Ok(EnvironmentReadyInfo { + selected_capability_roots: selected_capability_roots.clone(), + }))?; let harness_websocket = accept_websocket(&listener, "harness").await?; let first_relay = tokio::spawn(proxy_relay_frames( environment_websocket, @@ -174,6 +186,10 @@ async fn deferred_noise_environment_connects_and_reconnects_with_fresh_bundle() let initial_info = timeout(TEST_TIMEOUT, environment.info()) .await .context("deferred Noise environment should become ready")??; + assert_eq!( + environment.selected_capability_roots(), + selected_capability_roots + ); assert_eq!(provider.calls(), 1); assert_eq!( next_connection_state(&mut connection_state).await?, @@ -198,6 +214,10 @@ async fn deferred_noise_environment_connects_and_reconnects_with_fresh_bundle() .context("deferred Noise environment should reconnect")??; assert_eq!(recovered_info, initial_info); + assert_eq!( + environment.selected_capability_roots(), + selected_capability_roots + ); assert_eq!(provider.calls(), 2); assert_eq!( next_connection_state(&mut connection_state).await?, diff --git a/codex-rs/ext/extension-api/src/contributors/mcp.rs b/codex-rs/ext/extension-api/src/contributors/mcp.rs index 8f7598e264..30c67918d3 100644 --- a/codex-rs/ext/extension-api/src/contributors/mcp.rs +++ b/codex-rs/ext/extension-api/src/contributors/mcp.rs @@ -1,4 +1,5 @@ use codex_config::McpServerConfig; +use codex_protocol::capabilities::SelectedCapabilityRoot; use crate::ExtensionData; use crate::ExtensionDataInit; @@ -17,8 +18,8 @@ pub struct McpServerContributionContext<'a, C> { thread_init: Option<&'a ExtensionDataInit>, /// Effective request originator for the active thread, when resolution is thread-scoped. originator: Option<&'a str>, - /// Environment IDs whose selected roots may contribute to this exact step. - available_environment_ids: Option<&'a [String]>, + /// Selected roots resolved against ready environments for this exact step. + ready_selected_capability_roots: Option<&'a [SelectedCapabilityRoot]>, } impl Clone for McpServerContributionContext<'_, C> { @@ -37,7 +38,7 @@ impl<'a, C> McpServerContributionContext<'a, C> { thread_store: None, thread_init: None, originator: None, - available_environment_ids: None, + ready_selected_capability_roots: None, } } @@ -47,14 +48,14 @@ impl<'a, C> McpServerContributionContext<'a, C> { thread_init: &'a ExtensionDataInit, thread_store: &'a ExtensionData, originator: &'a str, - available_environment_ids: &'a [String], + ready_selected_capability_roots: &'a [SelectedCapabilityRoot], ) -> Self { Self { config, thread_store: Some(thread_store), thread_init: Some(thread_init), originator: Some(originator), - available_environment_ids: Some(available_environment_ids), + ready_selected_capability_roots: Some(ready_selected_capability_roots), } } @@ -78,12 +79,9 @@ impl<'a, C> McpServerContributionContext<'a, C> { self.originator } - /// Returns the exact environment availability projection for a model step. - /// - /// `Some` means contributors must omit selected roots whose environment ID is absent from the - /// slice. Global resolution returns `None` because it has no thread environments. - pub fn available_environment_ids(&self) -> Option<&'a [String]> { - self.available_environment_ids + /// Returns selected roots resolved against the ready environments for this model step. + pub fn ready_selected_capability_roots(&self) -> Option<&'a [SelectedCapabilityRoot]> { + self.ready_selected_capability_roots } } diff --git a/codex-rs/ext/mcp/src/executor_plugin.rs b/codex-rs/ext/mcp/src/executor_plugin.rs index 12c0a99686..2d63ab551f 100644 --- a/codex-rs/ext/mcp/src/executor_plugin.rs +++ b/codex-rs/ext/mcp/src/executor_plugin.rs @@ -6,7 +6,6 @@ use codex_extension_api::ExtensionFuture; use codex_extension_api::McpServerContribution; use codex_extension_api::McpServerContributionContext; use codex_extension_api::McpServerContributor; -use codex_protocol::capabilities::CapabilityRootLocation; use codex_protocol::capabilities::SelectedCapabilityRoot; use std::collections::HashMap; use std::sync::Arc; @@ -148,31 +147,16 @@ impl McpServerContributor for SelectedExecutorPluginMcpContributor { context: McpServerContributionContext<'a, Config>, ) -> ExtensionFuture<'a, Vec> { Box::pin(async move { - let Some(thread_init) = context.thread_init() else { - return Vec::new(); - }; let Some(thread_store) = context.thread_store() else { return Vec::new(); }; - let Some(selected_roots) = thread_init.get::>() else { + let Some(selected_roots) = context.ready_selected_capability_roots() else { return Vec::new(); }; let state = thread_store.get_or_init(SelectedExecutorPluginMcpState::default); let mut contributions = Vec::new(); for (selection_order, selected_root) in selected_roots.iter().enumerate() { - let CapabilityRootLocation::Environment { environment_id, .. } = - &selected_root.location; - if context - .available_environment_ids() - .is_some_and(|available| { - !available - .iter() - .any(|available| available == environment_id) - }) - { - continue; - } let Some(plugin) = self.metadata_for_root(&state, selected_root).await else { continue; }; diff --git a/codex-rs/ext/mcp/src/lib_tests.rs b/codex-rs/ext/mcp/src/lib_tests.rs index 7848dce7d0..220aec1e6b 100644 --- a/codex-rs/ext/mcp/src/lib_tests.rs +++ b/codex-rs/ext/mcp/src/lib_tests.rs @@ -27,7 +27,7 @@ async fn hosted_plugin_runtime_forwards_thread_originator() -> Result<(), Box