Resolve MCP server registrations through a catalog

codex-rs/codex-mcp/src/catalog.rs

@@ -0,0 +1,333 @@
+use std::cmp::Reverse;
+use std::collections::BTreeMap;
+use std::collections::BTreeSet;
+use std::collections::HashMap;
+
+use codex_config::McpServerConfig;
+
+/// The component that declared an MCP server registration.
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub enum McpServerSource {
+    Plugin { plugin_id: String },
+    Config,
+    Compatibility { id: String },
+    Extension { id: String },
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
+enum RegistrationPrecedence {
+    Plugin(Reverse<usize>),
+    Config,
+    Compatibility,
+    Extension(usize),
+}
+
+impl RegistrationPrecedence {
+    fn tier(self) -> u8 {
+        match self {
+            Self::Plugin(_) => 0,
+            Self::Config => 1,
+            Self::Compatibility => 2,
+            Self::Extension(_) => 3,
+        }
+    }
+}
+
+/// 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,
+        plugin_id: String,
+        plugin_order: usize,
+        config: McpServerConfig,
+    ) -> Self {
+        Self::new(
+            name,
+            McpServerSource::Plugin { plugin_id },
+            config,
+            RegistrationPrecedence::Plugin(Reverse(plugin_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;
+                    // Effective disabled winners remain name-scoped vetoes for later overlays.
+                    if !registration.config.enabled || disabled_server_names.contains(&name) {
+                        registration.config.enabled = false;
+                        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()
+    }
+
+    pub fn plugin_ids_by_server_name(&self) -> HashMap<String, String> {
+        self.servers
+            .iter()
+            .filter_map(|(name, server)| match server.source() {
+                McpServerSource::Plugin { plugin_id } => Some((name.clone(), plugin_id.clone())),
+                McpServerSource::Config
+                | McpServerSource::Compatibility { .. }
+                | McpServerSource::Extension { .. } => None,
+            })
+            .collect()
+    }
+
+    pub fn conflicts(&self) -> &[McpServerConflict] {
+        &self.conflicts
+    }
+}
+
+#[cfg(test)]
+#[path = "catalog_tests.rs"]
+mod tests;