Codex Handbook
codex-mcp/src/catalog.rs 404 lines
use std::cmp::Reverse;use std::collections::BTreeMap;use std::collections::BTreeSet;use std::collections::HashMap;use codex_config::McpServerConfig;/// Plugin identity retained with an MCP registration for tool attribution.#[derive(Clone, Debug, PartialEq, Eq)]pub struct McpPluginAttribution {    plugin_id: String,    display_name: String,}impl McpPluginAttribution {    pub fn new(plugin_id: String, display_name: String) -> Self {        Self {            plugin_id,            display_name,        }    }    pub fn plugin_id(&self) -> &str {        &self.plugin_id    }    pub fn display_name(&self) -> &str {        &self.display_name    }}/// The component that declared an MCP server registration.#[derive(Clone, Debug, PartialEq, Eq)]pub enum McpServerSource {    /// A plugin discovered through the process-wide legacy plugin manager.    Plugin(McpPluginAttribution),    /// A plugin explicitly selected for this thread through a capability root.    SelectedPlugin(McpPluginAttribution),    Config,    Compatibility {        id: String,    },    Extension {        id: String,    },}impl McpServerSource {    fn disabled_registration_is_name_veto(&self) -> bool {        // A selected package's policy applies to its registration, not to a higher runtime source        // that happens to use the same logical server name.        !matches!(self, Self::SelectedPlugin(_))    }}#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]enum RegistrationPrecedence {    Plugin(Reverse<usize>),    SelectedPlugin(Reverse<usize>),    Config,    Compatibility,    Extension(usize),}impl RegistrationPrecedence {    fn tier(self) -> u8 {        match self {            Self::Plugin(_) => 0,            Self::SelectedPlugin(_) => 1,            Self::Config => 2,            Self::Compatibility => 3,            Self::Extension(_) => 4,        }    }}/// One named MCP server declaration before source resolution.#[derive(Clone, Debug, PartialEq)]pub struct McpServerRegistration {    name: String,    source: McpServerSource,    config: McpServerConfig,    precedence: RegistrationPrecedence,}impl McpServerRegistration {    pub fn from_config(name: String, config: McpServerConfig) -> Self {        Self::new(            name,            McpServerSource::Config,            config,            RegistrationPrecedence::Config,        )    }    pub fn from_plugin(        name: String,        attribution: McpPluginAttribution,        plugin_order: usize,        config: McpServerConfig,    ) -> Self {        Self::new(            name,            McpServerSource::Plugin(attribution),            config,            RegistrationPrecedence::Plugin(Reverse(plugin_order)),        )    }    /// Registers a thread-selected plugin above discovered plugins and below config.    pub fn from_selected_plugin(        name: String,        attribution: McpPluginAttribution,        selection_order: usize,        config: McpServerConfig,    ) -> Self {        Self::new(            name,            McpServerSource::SelectedPlugin(attribution),            config,            RegistrationPrecedence::SelectedPlugin(Reverse(selection_order)),        )    }    pub fn from_compatibility(        name: String,        id: impl Into<String>,        config: McpServerConfig,    ) -> Self {        Self::new(            name,            McpServerSource::Compatibility { id: id.into() },            config,            RegistrationPrecedence::Compatibility,        )    }    pub fn from_extension(        name: String,        id: impl Into<String>,        contribution_order: usize,        config: McpServerConfig,    ) -> Self {        Self::new(            name,            McpServerSource::Extension { id: id.into() },            config,            RegistrationPrecedence::Extension(contribution_order),        )    }    fn new(        name: String,        source: McpServerSource,        config: McpServerConfig,        precedence: RegistrationPrecedence,    ) -> Self {        Self {            name,            source,            config,            precedence,        }    }}/// One side of an MCP server conflict, including whether it registers or/// removes the server.#[derive(Clone, Debug, PartialEq, Eq)]pub enum McpServerConflictAction {    Register(McpServerSource),    Remove(McpServerSource),}/// A same-tier name collision and the final outcome after all precedence is applied.#[derive(Clone, Debug, PartialEq, Eq)]pub struct McpServerConflict {    pub name: String,    pub outcome: McpServerConflictAction,    pub contenders: Vec<McpServerConflictAction>,}#[derive(Clone, Debug)]enum CatalogAction {    Register(Box<McpServerRegistration>),    Remove {        name: String,        source: McpServerSource,        precedence: RegistrationPrecedence,    },}impl CatalogAction {    fn name(&self) -> &str {        match self {            Self::Register(registration) => &registration.name,            Self::Remove { name, .. } => name,        }    }    fn precedence(&self) -> RegistrationPrecedence {        match self {            Self::Register(registration) => registration.precedence,            Self::Remove { precedence, .. } => *precedence,        }    }    fn conflict_action(&self) -> McpServerConflictAction {        match self {            Self::Register(registration) => {                McpServerConflictAction::Register(registration.source.clone())            }            Self::Remove { source, .. } => McpServerConflictAction::Remove(source.clone()),        }    }}/// Mutable inputs used to produce an immutable resolved catalog.#[derive(Clone, Debug, Default)]pub struct McpCatalogBuilder {    actions: Vec<CatalogAction>,    disabled_server_names: BTreeSet<String>,}impl McpCatalogBuilder {    pub fn register(&mut self, registration: McpServerRegistration) {        self.actions            .push(CatalogAction::Register(Box::new(registration)));    }    /// Applies the legacy name-scoped disabled veto after source resolution.    pub fn disable(&mut self, name: String) {        self.disabled_server_names.insert(name);    }    pub fn remove_compatibility(&mut self, name: String, id: impl Into<String>) {        self.actions.push(CatalogAction::Remove {            name,            source: McpServerSource::Compatibility { id: id.into() },            precedence: RegistrationPrecedence::Compatibility,        });    }    pub fn remove_extension(        &mut self,        name: String,        id: impl Into<String>,        contribution_order: usize,    ) {        self.actions.push(CatalogAction::Remove {            name,            source: McpServerSource::Extension { id: id.into() },            precedence: RegistrationPrecedence::Extension(contribution_order),        });    }    pub fn build(mut self) -> ResolvedMcpCatalog {        // Stable sorting makes action order the tie-breaker when precedence is equal.        self.actions.sort_by_key(CatalogAction::precedence);        let mut winners = BTreeMap::<String, CatalogAction>::new();        let mut actions_by_name_and_tier = BTreeMap::<(String, u8), Vec<&CatalogAction>>::new();        for action in &self.actions {            winners.insert(action.name().to_string(), action.clone());            actions_by_name_and_tier                .entry((action.name().to_string(), action.precedence().tier()))                .or_default()                .push(action);        }        let mut conflicts = Vec::new();        for ((name, _), actions) in actions_by_name_and_tier {            if actions.len() < 2 {                continue;            }            let Some(outcome) = winners.get(&name).map(CatalogAction::conflict_action) else {                continue;            };            conflicts.push(McpServerConflict {                name,                outcome,                contenders: actions                    .into_iter()                    .map(CatalogAction::conflict_action)                    .collect(),            });        }        let mut disabled_server_names = self.disabled_server_names;        let servers = winners            .into_iter()            .filter_map(|(name, action)| match action {                CatalogAction::Register(registration) => {                    let mut registration = *registration;                    let persist_disabled_name =                        registration.source.disabled_registration_is_name_veto();                    if !registration.config.enabled || disabled_server_names.contains(&name) {                        registration.config.enabled = false;                        if persist_disabled_name {                            // Preserve legacy disabled winners across later runtime overlays.                            disabled_server_names.insert(name.clone());                        }                    }                    Some((                        name,                        ResolvedMcpServer {                            source: registration.source,                            config: registration.config,                        },                    ))                }                CatalogAction::Remove { .. } => None,            })            .collect();        ResolvedMcpCatalog {            actions: self.actions,            disabled_server_names,            servers,            conflicts,        }    }}/// A single winning MCP registration.#[derive(Clone, Debug, PartialEq)]pub struct ResolvedMcpServer {    source: McpServerSource,    config: McpServerConfig,}impl ResolvedMcpServer {    pub fn source(&self) -> &McpServerSource {        &self.source    }    pub fn config(&self) -> &McpServerConfig {        &self.config    }}/// Immutable result of MCP registration resolution.#[derive(Clone, Debug, Default)]pub struct ResolvedMcpCatalog {    actions: Vec<CatalogAction>,    disabled_server_names: BTreeSet<String>,    servers: BTreeMap<String, ResolvedMcpServer>,    conflicts: Vec<McpServerConflict>,}impl ResolvedMcpCatalog {    pub fn builder() -> McpCatalogBuilder {        McpCatalogBuilder::default()    }    pub fn to_builder(&self) -> McpCatalogBuilder {        McpCatalogBuilder {            actions: self.actions.clone(),            disabled_server_names: self.disabled_server_names.clone(),        }    }    pub fn server(&self, name: &str) -> Option<&ResolvedMcpServer> {        self.servers.get(name)    }    pub fn configured_servers(&self) -> HashMap<String, McpServerConfig> {        self.servers            .iter()            .map(|(name, server)| (name.clone(), server.config.clone()))            .collect()    }    /// Returns package attribution for each winning plugin-owned server.    pub fn plugin_attributions_by_server_name(&self) -> HashMap<String, McpPluginAttribution> {        self.servers            .iter()            .filter_map(|(name, server)| match server.source() {                McpServerSource::Plugin(attribution)                | McpServerSource::SelectedPlugin(attribution) => {                    Some((name.clone(), attribution.clone()))                }                McpServerSource::Config                | McpServerSource::Compatibility { .. }                | McpServerSource::Extension { .. } => None,            })            .collect()    }    /// Returns the names of winning servers supplied by thread-selected plugins.    pub(crate) fn selected_plugin_server_names(&self) -> impl Iterator<Item = &str> {        self.servers.iter().filter_map(|(name, server)| {            matches!(server.source(), McpServerSource::SelectedPlugin(_)).then_some(name.as_str())        })    }    pub fn conflicts(&self) -> &[McpServerConflict] {        &self.conflicts    }}#[cfg(test)]#[path = "catalog_tests.rs"]mod tests;