codex/codex-rs/ext/mcp/tests/executor_plugin_mcp.rs

use codex_config::test_support::CloudConfigBundleFixture;
use codex_core::config::Config;
use codex_core::config::ConfigBuilder;
use codex_exec_server::EnvironmentManager;
use codex_exec_server::ExecutorCapabilityDiscoveryCache;
use codex_exec_server::LOCAL_ENVIRONMENT_ID;
use codex_extension_api::ExtensionData;
use codex_extension_api::ExtensionDataInit;
use codex_extension_api::ExtensionRegistryBuilder;
use codex_extension_api::McpServerContribution;
use codex_extension_api::McpServerContributionContext;
use codex_features::Feature;
use codex_protocol::capabilities::CapabilityRootLocation;
use codex_protocol::capabilities::SelectedCapabilityRoot;
use codex_utils_path_uri::PathUri;
use pretty_assertions::assert_eq;
use std::sync::Arc;

type TestResult = Result<(), Box<dyn std::error::Error>>;

#[derive(Debug, PartialEq, Eq)]
struct ContributionSummary {
    name: String,
    plugin_id: String,
    plugin_display_name: String,
    selection_order: usize,
    enabled: bool,
}

#[derive(Debug, PartialEq, Eq)]
struct PackageSummary {
    plugin_id: String,
    plugin_display_name: String,
    connector_ids: Vec<String>,
}

#[tokio::test]
async fn selected_plugin_servers_use_managed_requirements_for_the_selected_root_id() -> TestResult {
    let codex_home = tempfile::tempdir()?;
    let plugin_root = tempfile::tempdir()?;
    std::fs::create_dir_all(plugin_root.path().join(".codex-plugin"))?;
    std::fs::write(
        plugin_root.path().join(".codex-plugin/plugin.json"),
        r#"{"name":"different-manifest-name","interface":{"displayName":"Selected Demo"}}"#,
    )?;
    std::fs::write(
        plugin_root.path().join(".mcp.json"),
        r#"{
  "mcpServers": {
    "allowed": {"command":"allowed-command"},
    "mismatched": {"command":"wrong-command"},
    "unlisted": {"command":"unlisted-command"}
  }
}"#,
    )?;
    let config = ConfigBuilder::default()
        .codex_home(codex_home.path().to_path_buf())
        .fallback_cwd(Some(codex_home.path().to_path_buf()))
        .cloud_config_bundle(
            CloudConfigBundleFixture::loader_with_enterprise_requirement(
                r#"
[plugins."selected-root".mcp_servers.allowed.identity]
command = "allowed-command"

[plugins."selected-root".mcp_servers.mismatched.identity]
command = "expected-command"
"#,
            ),
        )
        .build()
        .await?;

    let contributions = selected_plugin_contributions(&config, plugin_root.path()).await?;

    assert_eq!(
        contributions,
        vec![
            ContributionSummary {
                name: "allowed".to_string(),
                plugin_id: "selected-root".to_string(),
                plugin_display_name: "Selected Demo".to_string(),
                selection_order: 0,
                enabled: true,
            },
            ContributionSummary {
                name: "mismatched".to_string(),
                plugin_id: "selected-root".to_string(),
                plugin_display_name: "Selected Demo".to_string(),
                selection_order: 0,
                enabled: false,
            },
            ContributionSummary {
                name: "unlisted".to_string(),
                plugin_id: "selected-root".to_string(),
                plugin_display_name: "Selected Demo".to_string(),
                selection_order: 0,
                enabled: false,
            },
        ]
    );
    Ok(())
}

#[tokio::test]
async fn selected_plugin_package_is_contributed_without_servers_or_connectors() -> TestResult {
    let codex_home = tempfile::tempdir()?;
    let plugin_root = tempfile::tempdir()?;
    std::fs::create_dir_all(plugin_root.path().join(".codex-plugin"))?;
    std::fs::create_dir_all(plugin_root.path().join("skills/deploy"))?;
    std::fs::write(
        plugin_root.path().join(".codex-plugin/plugin.json"),
        r#"{"name":"skill-only","interface":{"displayName":"Skill Only"}}"#,
    )?;
    std::fs::write(
        plugin_root.path().join("skills/deploy/SKILL.md"),
        "---\nname: deploy\ndescription: Deploy the project.\n---\n",
    )?;
    let config = ConfigBuilder::default()
        .codex_home(codex_home.path().to_path_buf())
        .fallback_cwd(Some(codex_home.path().to_path_buf()))
        .build()
        .await?;

    let contributions = raw_selected_plugin_contributions(&config, plugin_root.path()).await?;
    let package = contributions.into_iter().find_map(|contribution| {
        let McpServerContribution::SelectedPluginPackage {
            plugin_id,
            plugin_display_name,
            connector_ids,
        } = contribution
        else {
            return None;
        };
        Some(PackageSummary {
            plugin_id,
            plugin_display_name,
            connector_ids,
        })
    });

    assert_eq!(
        package,
        Some(PackageSummary {
            plugin_id: "selected-root".to_string(),
            plugin_display_name: "Skill Only".to_string(),
            connector_ids: Vec::new(),
        })
    );
    Ok(())
}

#[tokio::test]
async fn high_level_discovery_matches_the_existing_plugin_provider() -> TestResult {
    let codex_home = tempfile::tempdir()?;
    let plugin_root = tempfile::tempdir()?;
    std::fs::create_dir_all(plugin_root.path().join(".codex-plugin"))?;
    std::fs::write(
        plugin_root.path().join(".codex-plugin/plugin.json"),
        r#"{"name":"demo","interface":{"displayName":"Demo"},"mcpServers":"./servers.json"}"#,
    )?;
    std::fs::write(
        plugin_root.path().join("servers.json"),
        r#"{"mcpServers":{"first":{"command":"first"},"second":{"command":"second"}}}"#,
    )?;
    let mut config = ConfigBuilder::default()
        .codex_home(codex_home.path().to_path_buf())
        .fallback_cwd(Some(codex_home.path().to_path_buf()))
        .build()
        .await?;
    let existing = selected_plugin_contributions(&config, plugin_root.path()).await?;
    config
        .features
        .enable(Feature::ExecutorCapabilityDiscovery)
        .expect("test config should allow feature update");
    let high_level = selected_plugin_contributions(&config, plugin_root.path()).await?;

    assert_eq!(high_level, existing);
    Ok(())
}

async fn selected_plugin_contributions(
    config: &Config,
    plugin_root: &std::path::Path,
) -> Result<Vec<ContributionSummary>, Box<dyn std::error::Error>> {
    Ok(raw_selected_plugin_contributions(config, plugin_root)
        .await?
        .into_iter()
        .filter_map(|contribution| match contribution {
            McpServerContribution::SelectedPlugin {
                name,
                plugin_id,
                plugin_display_name,
                selection_order,
                config,
            } => Some(ContributionSummary {
                name,
                plugin_id,
                plugin_display_name,
                selection_order,
                enabled: config.enabled,
            }),
            McpServerContribution::SelectedPluginPackage { .. } => None,
            McpServerContribution::Set { .. } | McpServerContribution::Remove { .. } => {
                panic!("expected selected plugin contribution")
            }
        })
        .collect())
}

async fn raw_selected_plugin_contributions(
    config: &Config,
    plugin_root: &std::path::Path,
) -> Result<Vec<McpServerContribution>, Box<dyn std::error::Error>> {
    let mut builder = ExtensionRegistryBuilder::new();
    let environment_manager = Arc::new(EnvironmentManager::default_for_tests());
    codex_mcp_extension::install_executor_plugins(&mut builder, Arc::clone(&environment_manager));
    let registry = builder.build();
    let thread_init = ExtensionDataInit::new();
    let selected_capability_roots = vec![SelectedCapabilityRoot {
        id: "selected-root".to_string(),
        location: CapabilityRootLocation::Environment {
            environment_id: LOCAL_ENVIRONMENT_ID.to_string(),
            path: PathUri::from_host_native_path(plugin_root)?,
        },
    }];
    let thread_store = ExtensionData::new_with_init("test-thread", thread_init.clone());
    let executor_capability_discovery = if config
        .features
        .enabled(Feature::ExecutorCapabilityDiscovery)
    {
        Some(
            ExecutorCapabilityDiscoveryCache::new(environment_manager)
                .snapshot(&selected_capability_roots)
                .await,
        )
    } else {
        None
    };

    Ok(registry.mcp_server_contributors()[0]
        .contribute(McpServerContributionContext::for_step(
            config,
            &thread_init,
            &thread_store,
            "test_originator",
            &selected_capability_roots,
            executor_capability_discovery.as_ref(),
        ))
        .await)
}