From e1d9a697cdaeb9ec9438cfc6a2f5d0aac562574d Mon Sep 17 00:00:00 2001 From: jif-oai Date: Thu, 25 Jun 2026 01:01:42 +0100 Subject: [PATCH] Share selected capability bindings --- codex-rs/Cargo.lock | 1 - codex-rs/app-server/src/extensions.rs | 10 +---- codex-rs/app-server/src/mcp_refresh.rs | 1 - codex-rs/app-server/src/message_processor.rs | 1 - codex-rs/ext/connectors/Cargo.toml | 1 - .../ext/connectors/src/executor_plugin.rs | 10 +++-- codex-rs/ext/connectors/src/lib.rs | 3 +- codex-rs/ext/connectors/src/selected.rs | 43 +++++++------------ codex-rs/ext/mcp/src/executor_plugin.rs | 43 +++++++------------ .../ext/mcp/src/executor_plugin/provider.rs | 11 +++-- codex-rs/ext/mcp/src/lib.rs | 7 +-- codex-rs/ext/mcp/tests/executor_plugin_mcp.rs | 15 ++++--- 12 files changed, 57 insertions(+), 89 deletions(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index c21d639b5d..0121d173cd 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2603,7 +2603,6 @@ version = "0.0.0" dependencies = [ "codex-connectors", "codex-core-plugins", - "codex-exec-server", "codex-extension-api", "codex-plugin", "codex-protocol", diff --git a/codex-rs/app-server/src/extensions.rs b/codex-rs/app-server/src/extensions.rs index de29d58327..e823831262 100644 --- a/codex-rs/app-server/src/extensions.rs +++ b/codex-rs/app-server/src/extensions.rs @@ -9,7 +9,6 @@ use codex_core::NewThread; use codex_core::StartThreadOptions; use codex_core::ThreadManager; use codex_core::config::Config; -use codex_exec_server::EnvironmentManager; use codex_extension_api::AgentSpawnFuture; use codex_extension_api::AgentSpawner; use codex_extension_api::ExtensionEventSink; @@ -35,7 +34,6 @@ pub(crate) struct ThreadExtensionDependencies { pub(crate) analytics_events_client: AnalyticsEventsClient, pub(crate) thread_manager: Weak, pub(crate) goal_service: Arc, - pub(crate) environment_manager: Arc, pub(crate) executor_skill_provider: Arc, /// Process-scoped persistence backend for extensions that need stored thread history. pub(crate) thread_store: Arc, @@ -55,7 +53,6 @@ where analytics_events_client, thread_manager, goal_service, - environment_manager, executor_skill_provider, thread_store: _thread_store, } = dependencies; @@ -73,12 +70,9 @@ where } codex_guardian::install(&mut builder, guardian_agent_spawner); codex_memories_extension::install(&mut builder, codex_otel::global()); - codex_connectors_extension::install_selected_executor_connectors( - &mut builder, - Arc::clone(&environment_manager), - ); + codex_connectors_extension::install_selected_executor_connectors(&mut builder); codex_mcp_extension::install(&mut builder); - codex_mcp_extension::install_executor_plugins(&mut builder, environment_manager); + codex_mcp_extension::install_executor_plugins(&mut builder); codex_web_search_extension::install(&mut builder, auth_manager.clone()); codex_image_generation_extension::install(&mut builder, auth_manager, |config: &Config| { Some(config.codex_home.clone()) diff --git a/codex-rs/app-server/src/mcp_refresh.rs b/codex-rs/app-server/src/mcp_refresh.rs index 848d201f2e..1bde5bc41d 100644 --- a/codex-rs/app-server/src/mcp_refresh.rs +++ b/codex-rs/app-server/src/mcp_refresh.rs @@ -237,7 +237,6 @@ mod tests { analytics_events_client: codex_analytics::AnalyticsEventsClient::disabled(), thread_manager: thread_manager.clone(), goal_service: Arc::new(codex_goal_extension::GoalService::new()), - environment_manager: Arc::clone(&environment_manager), executor_skill_provider: Arc::clone(&executor_skill_provider), thread_store: Arc::clone(&thread_store), }, diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 41ecbfe561..5519677de0 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -363,7 +363,6 @@ impl MessageProcessor { analytics_events_client: analytics_events_client.clone(), thread_manager: thread_manager.clone(), goal_service: Arc::clone(&goal_service), - environment_manager: Arc::clone(&environment_manager_for_extensions), executor_skill_provider: Arc::clone(&executor_skill_provider), thread_store: Arc::clone(&thread_store), }, diff --git a/codex-rs/ext/connectors/Cargo.toml b/codex-rs/ext/connectors/Cargo.toml index 48ef48681c..b639d48f29 100644 --- a/codex-rs/ext/connectors/Cargo.toml +++ b/codex-rs/ext/connectors/Cargo.toml @@ -16,7 +16,6 @@ workspace = true [dependencies] codex-connectors = { workspace = true } codex-core-plugins = { workspace = true } -codex-exec-server = { workspace = true } codex-extension-api = { workspace = true } codex-plugin = { workspace = true } codex-protocol = { workspace = true } diff --git a/codex-rs/ext/connectors/src/executor_plugin.rs b/codex-rs/ext/connectors/src/executor_plugin.rs index c27ae3ac62..0ed88a78ae 100644 --- a/codex-rs/ext/connectors/src/executor_plugin.rs +++ b/codex-rs/ext/connectors/src/executor_plugin.rs @@ -1,5 +1,5 @@ use codex_connectors::parse_plugin_app_config; -use codex_core_plugins::ResolvedExecutorPlugin; +use codex_core_plugins::ResolvedSelectedCapabilityRoot; use codex_plugin::AppDeclaration; use codex_plugin::PluginResourceLocator; use codex_utils_path_uri::PathUri; @@ -33,9 +33,11 @@ impl ExecutorPluginConnectorProvider { /// Returns the connector declarations contributed by `plugin`. pub async fn load( &self, - plugin: &ResolvedExecutorPlugin, + selected_root: &ResolvedSelectedCapabilityRoot, ) -> Result, ExecutorPluginConnectorProviderError> { - let resolved_plugin = plugin.plugin(); + let Some(resolved_plugin) = selected_root.plugin() else { + return Ok(Vec::new()); + }; let plugin_id = resolved_plugin.selected_root_id(); let Some(PluginResourceLocator::Environment { path: config_path, .. @@ -43,7 +45,7 @@ impl ExecutorPluginConnectorProvider { else { return Ok(Vec::new()); }; - let contents = plugin + let contents = selected_root .file_system() .read_file_text(config_path, /*sandbox*/ None) .await diff --git a/codex-rs/ext/connectors/src/lib.rs b/codex-rs/ext/connectors/src/lib.rs index d8c384094e..9b1ee41988 100644 --- a/codex-rs/ext/connectors/src/lib.rs +++ b/codex-rs/ext/connectors/src/lib.rs @@ -10,9 +10,8 @@ pub use selected::SelectedExecutorConnectorProvider; /// Installs thread-scoped connector discovery for selected executor plugins. pub fn install_selected_executor_connectors( builder: &mut codex_extension_api::ExtensionRegistryBuilder, - environment_manager: std::sync::Arc, ) { builder.thread_extension_init_contributor(std::sync::Arc::new( - SelectedExecutorConnectorProvider::new(environment_manager), + SelectedExecutorConnectorProvider::new(), )); } diff --git a/codex-rs/ext/connectors/src/selected.rs b/codex-rs/ext/connectors/src/selected.rs index adede0714b..7bd20b13b4 100644 --- a/codex-rs/ext/connectors/src/selected.rs +++ b/codex-rs/ext/connectors/src/selected.rs @@ -1,60 +1,47 @@ use codex_connectors::ConnectorSnapshot; use codex_connectors::PluginConnectorSource; -use codex_core_plugins::ExecutorPluginProvider; -use codex_exec_server::EnvironmentManager; +use codex_core_plugins::SelectedCapabilityBindings; use codex_extension_api::ExtensionDataInit; use codex_extension_api::ExtensionFuture; use codex_extension_api::ThreadExtensionInitContributor; -use codex_protocol::capabilities::SelectedCapabilityRoot; -use std::sync::Arc; use crate::ExecutorPluginConnectorProvider; /// Resolves connector declarations from thread-selected executor plugins. #[derive(Clone, Debug)] pub struct SelectedExecutorConnectorProvider { - plugin_provider: ExecutorPluginProvider, connector_provider: ExecutorPluginConnectorProvider, } impl SelectedExecutorConnectorProvider { - /// Creates a provider backed by the active execution environments. - pub fn new(environment_manager: Arc) -> Self { + /// Creates a provider for already-bound selected capability roots. + pub fn new() -> Self { Self { - plugin_provider: ExecutorPluginProvider::new(environment_manager), connector_provider: ExecutorPluginConnectorProvider, } } /// Resolves one immutable connector snapshot in selected-root order. - pub async fn snapshot_for_roots( + pub async fn snapshot_for_bindings( &self, - selected_roots: &[SelectedCapabilityRoot], + bindings: &SelectedCapabilityBindings, ) -> ConnectorSnapshot { let mut sources = Vec::new(); + let snapshot = bindings.resolve_all().await; - for selected_root in selected_roots { - let plugin = match self.plugin_provider.resolve_bound(selected_root).await { - Ok(Some(plugin)) => plugin, - Ok(None) => continue, - Err(err) => { - tracing::warn!( - selected_root = selected_root.id, - error = %err, - "failed to resolve selected executor plugin for connector discovery" - ); - continue; - } + for selected_root in snapshot.ready() { + let Some(plugin) = selected_root.plugin() else { + continue; }; - match self.connector_provider.load(&plugin).await { + match self.connector_provider.load(selected_root).await { Ok(declarations) => sources.push(PluginConnectorSource::new( - plugin.plugin().selected_root_id(), - plugin.plugin().manifest().display_name(), + plugin.selected_root_id(), + plugin.manifest().display_name(), declarations, )), Err(err) => { tracing::warn!( - selected_root = selected_root.id, + selected_root = selected_root.selected_root().id, error = %err, "failed to load selected executor plugin connectors" ); @@ -72,10 +59,10 @@ impl ThreadExtensionInitContributor for SelectedExecutorConnectorProvider { if thread_init.get::().is_some() { return; } - let Some(selected_roots) = thread_init.get::>() else { + let Some(bindings) = thread_init.get::() else { return; }; - let snapshot = self.snapshot_for_roots(selected_roots.as_ref()).await; + let snapshot = self.snapshot_for_bindings(bindings.as_ref()).await; thread_init.insert(snapshot); }) } diff --git a/codex-rs/ext/mcp/src/executor_plugin.rs b/codex-rs/ext/mcp/src/executor_plugin.rs index 84e4e83c4f..30811d867e 100644 --- a/codex-rs/ext/mcp/src/executor_plugin.rs +++ b/codex-rs/ext/mcp/src/executor_plugin.rs @@ -1,15 +1,12 @@ use codex_core::config::Config; -use codex_core_plugins::ExecutorPluginProvider; -use codex_exec_server::EnvironmentManager; +use codex_core_plugins::SelectedCapabilityBindings; use codex_extension_api::ExtensionDataInit; use codex_extension_api::ExtensionFuture; use codex_extension_api::McpServerContribution; use codex_extension_api::McpServerContributionContext; use codex_extension_api::McpServerContributor; use codex_extension_api::ThreadExtensionInitContributor; -use codex_protocol::capabilities::SelectedCapabilityRoot; use std::collections::HashMap; -use std::sync::Arc; use tokio::sync::OnceCell; use self::provider::ExecutorPluginMcpProvider; @@ -42,7 +39,7 @@ pub(crate) struct SelectedExecutorPluginMcpInitializer; impl ThreadExtensionInitContributor for SelectedExecutorPluginMcpInitializer { fn initialize<'a>(&'a self, thread_init: &'a mut ExtensionDataInit) -> ExtensionFuture<'a, ()> { Box::pin(async move { - if thread_init.get::>().is_some() + if thread_init.get::().is_some() && thread_init .get::() .is_none() @@ -54,47 +51,37 @@ impl ThreadExtensionInitContributor for SelectedExecutorPluginMcpInitializer { } pub(crate) struct SelectedExecutorPluginMcpContributor { - plugin_provider: ExecutorPluginProvider, mcp_provider: ExecutorPluginMcpProvider, } impl SelectedExecutorPluginMcpContributor { - pub(crate) fn new(environment_manager: Arc) -> Self { + pub(crate) fn new() -> Self { Self { - plugin_provider: ExecutorPluginProvider::new(Arc::clone(&environment_manager)), mcp_provider: ExecutorPluginMcpProvider, } } async fn resolve_snapshot( &self, - selected_roots: &[SelectedCapabilityRoot], + bindings: &SelectedCapabilityBindings, ) -> Vec { let mut snapshot = Vec::new(); + let bindings = bindings.resolve_all().await; - for (selection_order, selected_root) in selected_roots.iter().enumerate() { - let plugin = match self.plugin_provider.resolve_bound(selected_root).await { - Ok(Some(plugin)) => plugin, - Ok(None) => continue, - Err(err) => { - tracing::warn!( - selected_root = selected_root.id, - error = %err, - "failed to resolve selected executor plugin for MCP discovery" - ); - continue; - } + for selected_root in bindings.ready() { + let Some(plugin) = selected_root.plugin() else { + continue; }; - match self.mcp_provider.load(&plugin).await { + match self.mcp_provider.load(selected_root).await { Ok(servers) => snapshot.push(SelectedPluginMcpServers { - plugin_id: plugin.plugin().selected_root_id().to_string(), - plugin_display_name: plugin.plugin().manifest().display_name().to_string(), - selection_order, + plugin_id: plugin.selected_root_id().to_string(), + plugin_display_name: plugin.manifest().display_name().to_string(), + selection_order: selected_root.selection_order(), servers, }), Err(err) => { tracing::warn!( - selected_root = selected_root.id, + selected_root = selected_root.selected_root().id, error = %err, "failed to load selected executor plugin MCP servers" ); @@ -119,7 +106,7 @@ impl McpServerContributor for SelectedExecutorPluginMcpContributor { let Some(thread_init) = context.thread_init() else { return Vec::new(); }; - let Some(selected_roots) = thread_init.get::>() else { + let Some(bindings) = thread_init.get::() else { return Vec::new(); }; let Some(state) = thread_init.get::() else { @@ -128,7 +115,7 @@ impl McpServerContributor for SelectedExecutorPluginMcpContributor { }; let snapshot = state .snapshot - .get_or_init(|| self.resolve_snapshot(selected_roots.as_ref())) + .get_or_init(|| self.resolve_snapshot(bindings.as_ref())) .await; let mut contributions = Vec::new(); diff --git a/codex-rs/ext/mcp/src/executor_plugin/provider.rs b/codex-rs/ext/mcp/src/executor_plugin/provider.rs index 9b8a9878e3..c2a6f6d0f1 100644 --- a/codex-rs/ext/mcp/src/executor_plugin/provider.rs +++ b/codex-rs/ext/mcp/src/executor_plugin/provider.rs @@ -1,6 +1,6 @@ use codex_config::McpServerConfig; use codex_config::McpServerTransportConfig; -use codex_core_plugins::ResolvedExecutorPlugin; +use codex_core_plugins::ResolvedSelectedCapabilityRoot; use codex_exec_server::ExecutorFileSystem; use codex_mcp::parse_executor_plugin_mcp_config; use codex_plugin::PluginResourceLocator; @@ -51,11 +51,14 @@ impl ExecutorPluginMcpProvider { /// Returns stdio servers declared by `plugin`, bound to its environment. pub(super) async fn load( &self, - plugin: &ResolvedExecutorPlugin, + selected_root: &ResolvedSelectedCapabilityRoot, ) -> Result, ExecutorPluginMcpProviderError> { - let ResolvedPluginLocation::Environment { root, .. } = plugin.plugin().location(); + let Some(plugin) = selected_root.plugin() else { + return Ok(Vec::new()); + }; + let ResolvedPluginLocation::Environment { root, .. } = plugin.location(); - load_from_file_system(plugin.plugin(), root, plugin.file_system()).await + load_from_file_system(plugin, root, selected_root.file_system().as_ref()).await } } diff --git a/codex-rs/ext/mcp/src/lib.rs b/codex-rs/ext/mcp/src/lib.rs index d76a83bcfa..e3179b4bf5 100644 --- a/codex-rs/ext/mcp/src/lib.rs +++ b/codex-rs/ext/mcp/src/lib.rs @@ -43,14 +43,11 @@ pub fn install(builder: &mut ExtensionRegistryBuilder) { } /// Installs discovery for MCP servers declared by thread-selected executor plugins. -pub fn install_executor_plugins( - builder: &mut ExtensionRegistryBuilder, - environment_manager: std::sync::Arc, -) { +pub fn install_executor_plugins(builder: &mut ExtensionRegistryBuilder) { builder.thread_extension_init_contributor(std::sync::Arc::new( executor_plugin::SelectedExecutorPluginMcpInitializer, )); builder.mcp_server_contributor(std::sync::Arc::new( - executor_plugin::SelectedExecutorPluginMcpContributor::new(environment_manager), + executor_plugin::SelectedExecutorPluginMcpContributor::new(), )); } diff --git a/codex-rs/ext/mcp/tests/executor_plugin_mcp.rs b/codex-rs/ext/mcp/tests/executor_plugin_mcp.rs index 423b5785c5..5d5bd55aaf 100644 --- a/codex-rs/ext/mcp/tests/executor_plugin_mcp.rs +++ b/codex-rs/ext/mcp/tests/executor_plugin_mcp.rs @@ -1,6 +1,7 @@ use codex_config::test_support::CloudConfigBundleFixture; use codex_core::config::Config; use codex_core::config::ConfigBuilder; +use codex_core_plugins::SelectedCapabilityBindings; use codex_exec_server::EnvironmentManager; use codex_exec_server::LOCAL_ENVIRONMENT_ID; use codex_extension_api::ExtensionDataInit; @@ -96,19 +97,21 @@ async fn selected_plugin_contributions( plugin_root: &std::path::Path, ) -> Result, Box> { let mut builder = ExtensionRegistryBuilder::new(); - codex_mcp_extension::install_executor_plugins( - &mut builder, - Arc::new(EnvironmentManager::default_for_tests()), - ); + codex_mcp_extension::install_executor_plugins(&mut builder); let registry = builder.build(); let mut thread_init = ExtensionDataInit::new(); - thread_init.insert(vec![SelectedCapabilityRoot { + let selected_roots = vec![SelectedCapabilityRoot { id: "selected-root".to_string(), location: CapabilityRootLocation::Environment { environment_id: LOCAL_ENVIRONMENT_ID.to_string(), path: PathUri::from_host_native_path(plugin_root)?, }, - }]); + }]; + thread_init.insert(SelectedCapabilityBindings::new( + selected_roots.clone(), + Arc::new(EnvironmentManager::default_for_tests()), + )); + thread_init.insert(selected_roots); registry.initialize_thread_data(&mut thread_init).await; Ok(registry.mcp_server_contributors()[0]