codex/codex-rs/ext/extension-api/src/contributors/mcp.rs

use codex_config::McpServerConfig;
use codex_exec_server_protocol::ExecutorCapabilityDiscoverySnapshot;
use codex_protocol::capabilities::SelectedCapabilityRoot;
use codex_protocol::protocol::SessionSource;

use crate::ExtensionData;
use crate::ExtensionDataInit;

/// Input supplied while resolving MCP server contributions.
///
/// Thread-scoped implementations can read stable host inputs through [`Self::thread_init`] and
/// keep their cache in [`Self::thread_store`]. Implementations should not retain borrowed context
/// after contribution completes.
pub struct McpServerContributionContext<'a, C> {
    /// Host configuration visible during MCP resolution.
    config: &'a C,
    /// Extension-owned data for the active thread, when resolution is thread-scoped.
    thread_store: Option<&'a ExtensionData>,
    /// Stable host inputs for the active thread, when resolution is thread-scoped.
    thread_init: Option<&'a ExtensionDataInit>,
    /// Source of the active thread, when supplied by the host runtime.
    session_source: Option<&'a SessionSource>,
    /// Effective request originator for the active thread, when resolution is thread-scoped.
    originator: Option<&'a str>,
    /// Selected roots resolved against ready environments for this exact step.
    ready_selected_capability_roots: Option<&'a [SelectedCapabilityRoot]>,
    /// Executor-materialized capability files shared by all consumers in this exact step.
    executor_capability_discovery: Option<&'a ExecutorCapabilityDiscoverySnapshot>,
}

impl<C> Clone for McpServerContributionContext<'_, C> {
    fn clone(&self) -> Self {
        *self
    }
}

impl<C> Copy for McpServerContributionContext<'_, C> {}

impl<'a, C> McpServerContributionContext<'a, C> {
    /// Creates context for resolution that is not associated with a running thread.
    pub fn global(config: &'a C) -> Self {
        Self {
            config,
            thread_store: None,
            thread_init: None,
            session_source: None,
            originator: None,
            ready_selected_capability_roots: None,
            executor_capability_discovery: None,
        }
    }

    /// Creates context for one model step using only currently available environments.
    pub fn for_step(
        config: &'a C,
        thread_init: &'a ExtensionDataInit,
        thread_store: &'a ExtensionData,
        originator: &'a str,
        ready_selected_capability_roots: &'a [SelectedCapabilityRoot],
        executor_capability_discovery: Option<&'a ExecutorCapabilityDiscoverySnapshot>,
    ) -> Self {
        Self {
            config,
            thread_store: Some(thread_store),
            thread_init: Some(thread_init),
            session_source: None,
            originator: Some(originator),
            ready_selected_capability_roots: Some(ready_selected_capability_roots),
            executor_capability_discovery,
        }
    }

    /// Attaches the stable source of the active thread to this contribution.
    pub fn with_session_source(mut self, session_source: &'a SessionSource) -> Self {
        self.session_source = Some(session_source);
        self
    }

    /// Returns the host configuration visible during resolution.
    pub fn config(&self) -> &'a C {
        self.config
    }

    /// Returns extension-owned state when resolving for a running thread.
    pub fn thread_store(&self) -> Option<&'a ExtensionData> {
        self.thread_store
    }

    /// Returns stable host inputs when resolving for a running thread.
    pub fn thread_init(&self) -> Option<&'a ExtensionDataInit> {
        self.thread_init
    }

    /// Returns the active thread's source when supplied by the host runtime.
    pub fn session_source(&self) -> Option<&'a SessionSource> {
        self.session_source
    }

    /// Returns the effective request originator when resolving for a running thread.
    pub fn originator(&self) -> Option<&'a str> {
        self.originator
    }

    /// 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
    }

    /// Returns the executor-materialized capability files for this model step, when enabled.
    pub fn executor_capability_discovery(&self) -> Option<&'a ExecutorCapabilityDiscoverySnapshot> {
        self.executor_capability_discovery
    }
}

/// One extension-owned overlay for the runtime MCP server configuration.
#[derive(Clone, Debug)]
pub enum McpServerContribution {
    /// Adds or replaces a named MCP server.
    Set {
        name: String,
        config: Box<McpServerConfig>,
    },
    /// Registers a server declared by a plugin selected for this thread.
    SelectedPlugin {
        name: String,
        plugin_id: String,
        plugin_display_name: String,
        selection_order: usize,
        config: Box<McpServerConfig>,
    },
    /// Records a plugin selected for this thread and any connector IDs it declares.
    SelectedPluginPackage {
        plugin_id: String,
        plugin_display_name: String,
        connector_ids: Vec<String>,
    },
    /// Removes a named MCP server.
    Remove { name: String },
}