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

use codex_config::AppToolApproval;
use codex_config::McpServerToolConfig;
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::collections::HashMap;
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"}
  }
}"#,
    )?;
    std::fs::write(
        codex_home.path().join("config.toml"),
        "[plugins.\"selected-root\".mcp_servers.mismatched]\nenabled = true\n[plugins.\"selected-root\".mcp_servers.unlisted]\nenabled = true",
    )?;
    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 managed_plugins_requirement_disables_selected_executor_plugin_capabilities() -> 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":"selected-root","interface":{"displayName":"Selected Root"}}"#,
    )?;
    std::fs::write(
        plugin_root.path().join(".mcp.json"),
        r#"{"mcpServers":{"probe":{"command":"probe-command"}}}"#,
    )?;
    let mut 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#"
[features]
plugins = false
"#,
            ),
        )
        .build()
        .await?;
    assert!(!config.features.enabled(Feature::Plugins));

    let direct = raw_selected_plugin_contributions(&config, plugin_root.path()).await?;
    assert!(
        matches!(
            direct.as_slice(),
            [McpServerContribution::SelectedPluginPackage { selected_root_id, .. }]
                if selected_root_id == "selected-root"
        ),
        "managed Plugins disable should preserve only the direct selected-root identity"
    );

    config
        .features
        .enable(Feature::ExecutorCapabilityDiscovery)
        .expect("test config should allow feature update");
    let discovered = raw_selected_plugin_contributions(&config, plugin_root.path()).await?;
    assert!(
        matches!(
            discovered.as_slice(),
            [McpServerContribution::SelectedPluginPackage { selected_root_id, .. }]
                if selected_root_id == "selected-root"
        ),
        "managed Plugins disable should preserve only the discovered selected-root identity"
    );
    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",
      "default_tools_approval_mode": "writes",
      "enabled_tools": ["read", "deploy", "trusted", "package-only"],
      "disabled_tools": ["package-denied"],
      "tools": {
        "read": {"approval_mode": "prompt", "output_token_limit": 12000},
        "deploy": {"approval_mode": "approve", "output_token_limit": 4000},
        "trusted": {"approval_mode": "approve"}
      }
    },
    "second": {
      "command": "second",
      "enabled": false,
      "default_tools_approval_mode": "prompt"
    }
  }
}"#,
    )?;
    std::fs::write(
        codex_home.path().join("config.toml"),
        r#"
[plugins."selected-root".mcp_servers.first]
enabled = false
default_tools_approval_mode = "prompt"
enabled_tools = ["read", "deploy", "trusted", "host-only"]
disabled_tools = ["write"]

[plugins."selected-root".mcp_servers.first.tools.read]
approval_mode = "approve"
output_token_limit = 8000

[plugins."selected-root".mcp_servers.first.tools.deploy]
output_token_limit = 9000

[plugins."selected-root".mcp_servers.first.tools.trusted]
approval_mode = "approve"

[plugins."selected-root".mcp_servers.second]
enabled = true
default_tools_approval_mode = "auto"
"#,
    )?;
    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?;
    let mut servers = raw_selected_plugin_contributions(&config, plugin_root.path())
        .await?
        .into_iter()
        .filter_map(|contribution| match contribution {
            McpServerContribution::SelectedPlugin { name, config, .. } => Some((name, config)),
            _ => None,
        })
        .collect::<HashMap<_, _>>();
    let server = servers
        .remove("first")
        .expect("disabled selected server remains registered");
    let declared_disabled_server = servers
        .remove("second")
        .expect("package-disabled server remains registered");
    assert_eq!(
        (
            server.enabled,
            server.default_tools_approval_mode,
            server.enabled_tools,
            server.disabled_tools,
            server.tools,
            declared_disabled_server.enabled,
            declared_disabled_server.default_tools_approval_mode,
        ),
        (
            false,
            Some(AppToolApproval::Prompt),
            Some(vec![
                "read".to_string(),
                "deploy".to_string(),
                "trusted".to_string(),
            ]),
            Some(vec!["package-denied".to_string(), "write".to_string()]),
            HashMap::from([
                (
                    "read".to_string(),
                    McpServerToolConfig {
                        approval_mode: Some(AppToolApproval::Prompt),
                        output_token_limit: std::num::NonZeroUsize::new(8_000),
                    },
                ),
                (
                    "deploy".to_string(),
                    McpServerToolConfig {
                        approval_mode: Some(AppToolApproval::Prompt),
                        output_token_limit: std::num::NonZeroUsize::new(4_000),
                    },
                ),
                (
                    "trusted".to_string(),
                    McpServerToolConfig {
                        approval_mode: Some(AppToolApproval::Approve),
                        ..Default::default()
                    },
                ),
            ]),
            false,
            Some(AppToolApproval::Prompt),
        )
    );
    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::HostedApps { .. }
            | 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, &Default::default())
                .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)
}