codex-rs/codex-mcp/src/catalog_tests.rs

use std::collections::BTreeMap;
use std::collections::HashMap;
use std::time::Duration;

use codex_config::AppToolApproval;
use codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID;
use codex_config::McpServerAuth;
use codex_config::McpServerConfig;
use codex_config::McpServerDisabledReason;
use codex_config::McpServerIdpOAuthConfig;
use codex_config::McpServerToolConfig;
use codex_config::McpServerTransportConfig;
use codex_config::types::PluginMcpServerEmaAuthConfig;
use codex_protocol::mcp_policy::EnvironmentMcpPolicy;
use codex_protocol::mcp_policy::PluginMcpRequirements;
use codex_utils_path_uri::PathUri;
use pretty_assertions::assert_eq;

use crate::CODEX_APPS_MCP_SERVER_NAME;
use crate::McpProtocolMode;

use super::McpEnvironmentAuthority;
use super::McpPluginAttribution;
use super::McpServerConflict;
use super::McpServerConflictAction;
use super::McpServerRegistration;
use super::McpServerSource;
use super::ResolvedMcpCatalog;
use super::ResolvedMcpServer;

fn server(url: &str) -> McpServerConfig {
    McpServerConfig {
        auth: Default::default(),
        transport: McpServerTransportConfig::StreamableHttp {
            url: url.to_string(),
            bearer_token_env_var: None,
            http_headers: None,
            env_http_headers: None,
            http_headers_helper: None,
        },
        environment_id: DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(),
        enabled: true,
        required: true,
        supports_parallel_tool_calls: true,
        omit_tools_from: None,
        disabled_reason: None,
        startup_timeout_sec: Some(Duration::from_secs(7)),
        tool_timeout_sec: Some(Duration::from_secs(11)),
        default_tools_approval_mode: Some(AppToolApproval::Prompt),
        enabled_tools: Some(vec!["read".to_string()]),
        disabled_tools: Some(vec!["write".to_string()]),
        scopes: None,
        oauth: None,
        oauth_resource: None,
        tools: HashMap::from([(
            "read".to_string(),
            McpServerToolConfig {
                approval_mode: Some(AppToolApproval::Approve),
                ..Default::default()
            },
        )]),
    }
}

fn plugin(plugin_id: &str) -> McpPluginAttribution {
    McpPluginAttribution::new(plugin_id.to_string(), plugin_id.to_string())
}

fn plugin_source(plugin_id: &str) -> McpServerSource {
    McpServerSource::Plugin(plugin(plugin_id))
}

fn selected_plugin_source(plugin_id: &str) -> McpServerSource {
    McpServerSource::SelectedPlugin(plugin(plugin_id))
}

fn compatibility_source(id: &str) -> McpServerSource {
    McpServerSource::Compatibility { id: id.to_string() }
}

fn extension_source(id: &str) -> McpServerSource {
    McpServerSource::Extension {
        id: id.to_string(),
        host_owned_apps: false,
    }
}

fn register(source: McpServerSource) -> McpServerConflictAction {
    McpServerConflictAction::Register(source)
}

fn remove(source: McpServerSource) -> McpServerConflictAction {
    McpServerConflictAction::Remove(source)
}

#[test]
fn plugin_host_root_is_retained_in_catalog_identity() {
    let original_root = PathUri::parse("file:///plugins/original").expect("valid plugin root URI");
    let replacement_root =
        PathUri::parse("file:///plugins/replacement").expect("valid plugin root URI");
    let catalog_for_root = |root| {
        let mut builder = ResolvedMcpCatalog::builder();
        builder.register(McpServerRegistration::from_plugin(
            "docs".to_string(),
            plugin("plugin@test").with_host_root(root),
            /*plugin_order*/ 0,
            server("https://plugin.example/mcp"),
        ));
        builder.build()
    };
    let original = catalog_for_root(original_root.clone());
    let replacement = catalog_for_root(replacement_root);

    let Some(McpServerSource::Plugin(attribution)) =
        original.server("docs").map(ResolvedMcpServer::source)
    else {
        panic!("expected host-discovered plugin registration");
    };
    assert_eq!(attribution.host_root(), Some(&original_root));
    assert!(!original.has_same_servers(&replacement));
}

#[test]
fn ema_policy_survives_rebuilds_and_rebinds_materialized_servers() {
    let idp = McpServerIdpOAuthConfig {
        issuer: "https://idp.example".to_string(),
        client_id: "enterprise-client".to_string(),
    };
    let mut builder = ResolvedMcpCatalog::builder();
    builder.enable_ema(idp.clone());
    let enabled_empty = builder.build();
    assert!(!enabled_empty.has_same_servers(&ResolvedMcpCatalog::default()));

    let mut builder = enabled_empty.to_builder();
    let mut config = server("https://resource.example/mcp");
    config.auth = McpServerAuth::EmaAuth;
    builder.register(McpServerRegistration::from_config(
        "enterprise".to_string(),
        config,
    ));
    let original = builder.build();
    assert!(original.has_same_servers(&original.to_builder().build()));

    let mut servers = original.configured_servers();
    let config = servers.get_mut("enterprise").expect("EMA server");
    let McpServerTransportConfig::StreamableHttp { url, .. } = &mut config.transport else {
        panic!("expected HTTP server");
    };
    *url = "https://resource.example/revised".to_string();
    let materialized = original.with_materialized_servers(servers);
    let config = materialized
        .server("enterprise")
        .expect("materialized EMA server")
        .config();
    let registration = config
        .ema_registration()
        .expect("finalized EMA registration");
    assert_eq!(
        (
            config.enabled,
            registration.idp(),
            registration.server_url()
        ),
        (true, &idp, "https://resource.example/revised")
    );

    let mut builder = ResolvedMcpCatalog::builder();
    builder.register(McpServerRegistration::from_config(
        "enterprise".to_string(),
        config.clone(),
    ));
    let denied = builder.build();
    let config = denied
        .server("enterprise")
        .expect("denied EMA server")
        .config();
    assert_eq!((config.enabled, config.ema_registration()), (false, None));
}

#[test]
fn rejected_plugin_ema_registration_does_not_veto_hosted_apps() {
    let idp = McpServerIdpOAuthConfig {
        issuer: "https://idp.example".to_string(),
        client_id: "enterprise-client".to_string(),
    };
    for (url, resource) in [
        ("https://other.example/mcp", "https://resource.example"),
        ("https://resource.example/mcp", " "),
    ] {
        for initially_enabled in [true, false] {
            let mut rejected = server(url);
            rejected.enabled = initially_enabled;
            let policy = PluginMcpServerEmaAuthConfig {
                url: "https://resource.example/mcp".to_string(),
                client_id: "resource-client".to_string(),
                authorization_server_issuer: "https://as.example".to_string(),
                scopes: Vec::new(),
                resource: resource.to_string(),
            };
            policy.apply(&mut rejected);
            assert!(rejected.resolve_ema_registration(&idp).is_err());
            assert_eq!(
                (
                    rejected.enabled,
                    rejected.auth.clone(),
                    rejected.ema_registration()
                ),
                (false, McpServerAuth::EmaAuth, None),
            );
            assert_eq!(
                rejected.disabled_reason,
                initially_enabled.then_some(McpServerDisabledReason::EmaRegistration),
            );

            let mut builder = ResolvedMcpCatalog::builder();
            builder.enable_ema(idp.clone());
            builder.register(McpServerRegistration::from_plugin(
                CODEX_APPS_MCP_SERVER_NAME.to_string(),
                plugin("plugin@test"),
                /*plugin_order*/ 0,
                rejected.clone(),
            ));
            let catalog = builder.build();
            assert_eq!(
                catalog.server(CODEX_APPS_MCP_SERVER_NAME).unwrap().config(),
                &rejected,
            );

            let materialized = catalog.with_materialized_servers(catalog.configured_servers());
            for catalog in [catalog, materialized] {
                let mut builder = catalog.to_builder();
                let mut expected = server("https://chatgpt.com/mcp");
                builder.register(McpServerRegistration::from_hosted_apps(
                    "apps",
                    /*contribution_order*/ 0,
                    expected.clone(),
                ));
                expected.enabled = initially_enabled;
                assert_eq!(
                    builder.build().server(CODEX_APPS_MCP_SERVER_NAME),
                    Some(&ResolvedMcpServer {
                        source: McpServerSource::Extension {
                            id: "apps".to_string(),
                            host_owned_apps: true,
                        },
                        config: expected,
                        protocol_mode: None,
                    }),
                );
            }
        }
    }
}

#[test]
fn source_precedence_preserves_the_winning_registration() {
    let extension = server("https://extension.example/mcp");
    let mut plugin_server = server("https://plugin.example/mcp");
    plugin_server.enabled = false;
    let mut builder = ResolvedMcpCatalog::builder();
    builder.register(McpServerRegistration::from_extension(
        "docs".to_string(),
        "hosted",
        /*contribution_order*/ 0,
        extension.clone(),
    ));
    builder.register(McpServerRegistration::from_plugin(
        "docs".to_string(),
        plugin("plugin@test"),
        /*plugin_order*/ 0,
        plugin_server,
    ));
    builder.register(McpServerRegistration::from_plugin(
        "docs".to_string(),
        plugin("other-plugin@test"),
        /*plugin_order*/ 1,
        server("https://other-plugin.example/mcp"),
    ));
    builder.register(McpServerRegistration::from_compatibility(
        "docs".to_string(),
        "legacy",
        server("https://compatibility.example/mcp"),
    ));
    builder.register(McpServerRegistration::from_config(
        "docs".to_string(),
        server("https://config.example/mcp"),
    ));

    let catalog = builder.build();
    let resolved = catalog.server("docs").expect("resolved server");

    assert_eq!(
        resolved.source(),
        &McpServerSource::Extension {
            id: "hosted".to_string(),
            host_owned_apps: false,
        }
    );
    assert_eq!(resolved.config(), &extension);
    assert!(catalog.plugin_attributions_by_server_name().is_empty());
    assert_eq!(
        catalog.conflicts(),
        &[McpServerConflict {
            name: "docs".to_string(),
            outcome: register(extension_source("hosted")),
            contenders: vec![
                register(plugin_source("other-plugin@test")),
                register(plugin_source("plugin@test")),
            ],
        }]
    );
}

#[test]
fn disabled_veto_only_disables_the_winning_registration() {
    let extension = server("https://extension.example/mcp");
    let mut expected = extension.clone();
    expected.enabled = false;
    let mut builder = ResolvedMcpCatalog::builder();
    builder.register(McpServerRegistration::from_extension(
        "docs".to_string(),
        "hosted",
        /*contribution_order*/ 0,
        extension,
    ));
    builder.disable("docs".to_string());

    let actual = builder
        .build()
        .server("docs")
        .expect("resolved server")
        .config()
        .clone();

    assert_eq!(actual, expected);
}

#[test]
fn disabled_winner_remains_a_veto_when_the_catalog_is_extended() {
    let mut disabled = server("https://config.example/mcp");
    disabled.enabled = false;
    let mut expected = server("https://extension.example/mcp");
    expected.enabled = false;
    let mut builder = ResolvedMcpCatalog::builder();
    builder.register(McpServerRegistration::from_config(
        "docs".to_string(),
        disabled,
    ));
    let mut builder = builder.build().to_builder();
    builder.register(McpServerRegistration::from_extension(
        "docs".to_string(),
        "hosted",
        /*contribution_order*/ 0,
        server("https://extension.example/mcp"),
    ));

    let resolved = builder.build();

    assert_eq!(
        resolved.server("docs"),
        Some(&super::ResolvedMcpServer {
            source: extension_source("hosted"),
            config: expected,
            protocol_mode: None,
        })
    );
}

#[test]
fn disabled_discovered_plugin_remains_a_veto_for_runtime_overlays() {
    let mut disabled = server("https://plugin.example/mcp");
    disabled.enabled = false;
    let mut expected = server("https://extension.example/mcp");
    expected.enabled = false;
    let mut builder = ResolvedMcpCatalog::builder();
    builder.register(McpServerRegistration::from_plugin(
        "docs".to_string(),
        plugin("plugin@test"),
        /*plugin_order*/ 0,
        disabled,
    ));
    let mut builder = builder.build().to_builder();
    builder.register(McpServerRegistration::from_extension(
        "docs".to_string(),
        "hosted",
        /*contribution_order*/ 0,
        server("https://extension.example/mcp"),
    ));

    let resolved = builder.build();

    assert_eq!(
        resolved.server("docs"),
        Some(&super::ResolvedMcpServer {
            source: extension_source("hosted"),
            config: expected,
            protocol_mode: None,
        })
    );
}

#[test]
fn earlier_plugin_wins_with_an_explicit_conflict() {
    let mut builder = ResolvedMcpCatalog::builder();
    builder.register(McpServerRegistration::from_plugin(
        "docs".to_string(),
        plugin("alpha@test"),
        /*plugin_order*/ 0,
        server("https://alpha.example/mcp"),
    ));
    builder.register(McpServerRegistration::from_plugin(
        "docs".to_string(),
        plugin("beta@test"),
        /*plugin_order*/ 1,
        server("https://beta.example/mcp"),
    ));

    let catalog = builder.build();

    assert_eq!(
        catalog.plugin_attributions_by_server_name(),
        HashMap::from([("docs".to_string(), plugin("alpha@test"))])
    );
    assert_eq!(
        catalog.conflicts(),
        &[McpServerConflict {
            name: "docs".to_string(),
            outcome: register(plugin_source("alpha@test")),
            contenders: vec![
                register(plugin_source("beta@test")),
                register(plugin_source("alpha@test")),
            ],
        }]
    );
}

#[test]
fn selected_plugins_override_discovered_plugins_but_not_config() {
    let selected = server("https://selected-alpha.example/mcp");
    let mut discovered = server("https://local.example/mcp");
    discovered.enabled = false;
    discovered.default_tools_approval_mode = Some(AppToolApproval::Auto);
    let mut builder = ResolvedMcpCatalog::builder();
    builder.register(McpServerRegistration::from_plugin(
        "docs".to_string(),
        plugin("local@test"),
        /*plugin_order*/ 0,
        discovered,
    ));
    builder.register(McpServerRegistration::from_selected_plugin(
        "docs".to_string(),
        plugin("selected-beta"),
        /*selection_order*/ 1,
        server("https://selected-beta.example/mcp"),
    ));
    builder.register(McpServerRegistration::from_selected_plugin(
        "docs".to_string(),
        plugin("selected-alpha"),
        /*selection_order*/ 0,
        selected.clone(),
    ));

    let catalog = builder.build();

    assert_eq!(
        catalog.server("docs"),
        Some(&super::ResolvedMcpServer {
            source: selected_plugin_source("selected-alpha"),
            config: selected,
            protocol_mode: None,
        })
    );
    assert_eq!(
        catalog.plugin_attributions_by_server_name(),
        HashMap::from([("docs".to_string(), plugin("selected-alpha"))])
    );
    assert_eq!(
        catalog.conflicts(),
        &[McpServerConflict {
            name: "docs".to_string(),
            outcome: register(selected_plugin_source("selected-alpha")),
            contenders: vec![
                register(selected_plugin_source("selected-beta")),
                register(selected_plugin_source("selected-alpha")),
            ],
        }]
    );

    let refreshed = server("https://refreshed.example/mcp");
    let catalog =
        catalog.with_materialized_servers(HashMap::from([("docs".to_string(), refreshed.clone())]));
    assert_eq!(
        catalog.server("docs"),
        Some(&super::ResolvedMcpServer {
            source: selected_plugin_source("selected-alpha"),
            config: refreshed,
            protocol_mode: None,
        })
    );

    let mut builder = catalog.to_builder();
    let configured = server("https://config.example/mcp");
    builder.register(McpServerRegistration::from_config(
        "docs".to_string(),
        configured.clone(),
    ));
    let catalog = builder.build();

    assert_eq!(
        catalog.server("docs"),
        Some(&super::ResolvedMcpServer {
            source: McpServerSource::Config,
            config: configured,
            protocol_mode: None,
        })
    );
}

#[test]
fn disabled_selected_plugin_does_not_veto_runtime_overlays() {
    let mut disabled = server("https://selected.example/mcp");
    disabled.enabled = false;
    let extension = server("https://extension.example/mcp");
    let mut builder = ResolvedMcpCatalog::builder();
    builder.register(McpServerRegistration::from_selected_plugin(
        "docs".to_string(),
        plugin("selected"),
        /*selection_order*/ 0,
        disabled,
    ));
    let mut builder = builder.build().to_builder();
    builder.register(McpServerRegistration::from_extension(
        "docs".to_string(),
        "hosted",
        /*contribution_order*/ 0,
        extension.clone(),
    ));

    let resolved = builder.build();

    assert_eq!(
        resolved.server("docs"),
        Some(&super::ResolvedMcpServer {
            source: extension_source("hosted"),
            config: extension,
            protocol_mode: None,
        })
    );
}

#[test]
fn equal_precedence_uses_insertion_order_not_source_identity() {
    let mut builder = ResolvedMcpCatalog::builder();
    builder.register(McpServerRegistration::from_compatibility(
        "docs".to_string(),
        "z-first",
        server("https://first.example/mcp"),
    ));
    builder.register(McpServerRegistration::from_compatibility(
        "docs".to_string(),
        "a-second",
        server("https://second.example/mcp"),
    ));

    let catalog = builder.build();

    assert_eq!(
        catalog.server("docs"),
        Some(&super::ResolvedMcpServer {
            source: compatibility_source("a-second"),
            config: server("https://second.example/mcp"),
            protocol_mode: None,
        })
    );
    let mut builder = catalog.to_builder();
    builder.remove_compatibility("docs".to_string(), "remove-last");

    let catalog = builder.build();

    assert_eq!(catalog.server("docs"), None);
    assert_eq!(
        catalog.conflicts(),
        &[McpServerConflict {
            name: "docs".to_string(),
            outcome: remove(compatibility_source("remove-last")),
            contenders: vec![
                register(compatibility_source("z-first")),
                register(compatibility_source("a-second")),
                remove(compatibility_source("remove-last")),
            ],
        }]
    );
}

#[test]
fn extension_protocol_mode_follows_the_winner_through_materialization() {
    let config = server("https://apps.example/mcp");
    let mut builder = ResolvedMcpCatalog::builder();
    builder.register(
        McpServerRegistration::from_extension(
            "apps".to_string(),
            "loser",
            /*contribution_order*/ 0,
            config.clone(),
        )
        .with_protocol_mode(McpProtocolMode::V20260728),
    );
    builder.register(McpServerRegistration::from_extension(
        "apps".to_string(),
        "winner",
        /*contribution_order*/ 1,
        config.clone(),
    ));
    let without_override = builder.build();
    assert_eq!(
        without_override
            .server("apps")
            .and_then(ResolvedMcpServer::protocol_mode),
        None
    );

    let mut builder = ResolvedMcpCatalog::builder();
    builder.register(
        McpServerRegistration::from_extension(
            "apps".to_string(),
            "winner",
            /*contribution_order*/ 1,
            config,
        )
        .with_protocol_mode(McpProtocolMode::Legacy),
    );
    let with_override = builder.build();
    assert!(!without_override.has_same_servers(&with_override));

    let refreshed = server("https://refreshed.example/mcp");
    let materialized = with_override
        .with_materialized_servers(HashMap::from([("apps".to_string(), refreshed.clone())]));
    assert_eq!(
        materialized.server("apps"),
        Some(&ResolvedMcpServer {
            source: extension_source("winner"),
            config: refreshed,
            protocol_mode: Some(McpProtocolMode::Legacy),
        })
    );
}

#[test]
fn environment_policy_exempts_only_explicitly_host_owned_apps() {
    let policy = EnvironmentMcpPolicy {
        servers: Some(BTreeMap::new()),
        plugins: None,
    };
    for (registration, expected) in [
        (
            McpServerRegistration::from_extension(
                CODEX_APPS_MCP_SERVER_NAME.to_string(),
                "apps",
                /*contribution_order*/ 0,
                server("https://apps.example/mcp"),
            ),
            false,
        ),
        (
            McpServerRegistration::from_hosted_apps(
                "apps",
                /*contribution_order*/ 0,
                server("https://apps.example/mcp"),
            ),
            true,
        ),
    ] {
        let mut builder = ResolvedMcpCatalog::builder();
        builder.register(registration);
        let catalog = builder
            .build_with_environment_authority(|_| McpEnvironmentAuthority::Restricted(&policy));
        assert_eq!(
            catalog
                .server(CODEX_APPS_MCP_SERVER_NAME)
                .expect("Apps registration")
                .config()
                .enabled,
            expected
        );
    }
}

#[test]
fn environment_policy_preserves_selected_plugin_and_empty_server_allowlist_semantics() {
    let mut builder = ResolvedMcpCatalog::builder();
    builder.register(McpServerRegistration::from_selected_plugin(
        "selected".to_string(),
        plugin("selected-plugin"),
        /*selection_order*/ 0,
        server("https://plugin.example/mcp"),
    ));
    let metadata_only_policy = EnvironmentMcpPolicy {
        servers: None,
        plugins: Some(BTreeMap::from([(
            "metadata-only-plugin".to_string(),
            PluginMcpRequirements { mcp_servers: None },
        )])),
    };
    let deny_all_policy = EnvironmentMcpPolicy {
        servers: Some(BTreeMap::new()),
        plugins: None,
    };

    for (policy, expected) in [(&metadata_only_policy, true), (&deny_all_policy, false)] {
        let resolved = builder
            .clone()
            .build_with_environment_authority(|_| McpEnvironmentAuthority::Restricted(policy));
        assert_eq!(
            resolved
                .server("selected")
                .expect("selected plugin")
                .config()
                .enabled,
            expected
        );
    }
}