Codex Handbook
ext/mcp/src/executor_plugin.rs 139 lines
use codex_core::config::Config;use codex_core_plugins::ExecutorPluginProvider;use codex_exec_server::EnvironmentManager;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_protocol::capabilities::SelectedCapabilityRoot;use std::collections::HashMap;use std::sync::Arc;use tokio::sync::OnceCell;use self::provider::ExecutorPluginMcpProvider;mod provider;/// Frozen MCP declarations for one selected package.////// Each server config retains the stable logical environment ID. Reconnection may replace the/// concrete environment instance without changing that authority.#[derive(Clone)]struct SelectedPluginMcpServers {    plugin_id: String,    plugin_display_name: String,    selection_order: usize,    servers: Vec<(String, codex_config::McpServerConfig)>,}#[derive(Default)]pub(crate) struct SelectedExecutorPluginMcpState {    snapshot: OnceCell<Vec<SelectedPluginMcpServers>>,}pub(crate) fn seed_thread_state(thread_init: &mut ExtensionDataInit) {    thread_init.insert(SelectedExecutorPluginMcpState::default());}pub(crate) struct SelectedExecutorPluginMcpContributor {    plugin_provider: ExecutorPluginProvider,    mcp_provider: ExecutorPluginMcpProvider,}impl SelectedExecutorPluginMcpContributor {    pub(crate) fn new(environment_manager: Arc<EnvironmentManager>) -> Self {        Self {            plugin_provider: ExecutorPluginProvider::new(Arc::clone(&environment_manager)),            mcp_provider: ExecutorPluginMcpProvider,        }    }    async fn resolve_snapshot(        &self,        selected_roots: &[SelectedCapabilityRoot],    ) -> Vec<SelectedPluginMcpServers> {        let mut snapshot = Vec::new();        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;                }            };            match self.mcp_provider.load(&plugin).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,                    servers,                }),                Err(err) => {                    tracing::warn!(                        selected_root = selected_root.id,                        error = %err,                        "failed to load selected executor plugin MCP servers"                    );                }            }        }        snapshot    }}impl McpServerContributor<Config> for SelectedExecutorPluginMcpContributor {    fn id(&self) -> &'static str {        "selected_executor_plugin_mcp"    }    fn contribute<'a>(        &'a self,        context: McpServerContributionContext<'a, Config>,    ) -> ExtensionFuture<'a, Vec<McpServerContribution>> {        Box::pin(async move {            let Some(thread_init) = context.thread_init() else {                return Vec::new();            };            let Some(selected_roots) = thread_init.get::<Vec<SelectedCapabilityRoot>>() else {                return Vec::new();            };            let Some(state) = thread_init.get::<SelectedExecutorPluginMcpState>() else {                tracing::warn!("selected executor plugin MCP state was not initialized");                return Vec::new();            };            let snapshot = state                .snapshot                .get_or_init(|| self.resolve_snapshot(selected_roots.as_ref()))                .await;            let mut contributions = Vec::new();            for plugin in snapshot {                let mut servers = plugin.servers.iter().cloned().collect::<HashMap<_, _>>();                context                    .config()                    .apply_plugin_mcp_server_requirements(&plugin.plugin_id, &mut servers);                let mut servers = servers.into_iter().collect::<Vec<_>>();                servers.sort_unstable_by(|left, right| left.0.cmp(&right.0));                contributions.extend(servers.into_iter().map(|(name, config)| {                    McpServerContribution::SelectedPlugin {                        name,                        plugin_id: plugin.plugin_id.clone(),                        plugin_display_name: plugin.plugin_display_name.clone(),                        selection_order: plugin.selection_order,                        config: Box::new(config),                    }                }));            }            contributions        })    }}