codex/codex-rs/exec-server/tests/capability_discovery.rs

mod common;
#[cfg(target_os = "linux")]
#[path = "common/fake_bwrap.rs"]
mod fake_bwrap;

#[cfg(target_os = "linux")]
use anyhow::Context as _;
use codex_exec_server::CAPABILITY_ROOTS_DISCOVER_METHOD;
use codex_exec_server::CapabilityRootDiscovery;
use codex_exec_server::CapabilityRootsDiscoverParams;
use codex_exec_server::CapabilityRootsDiscoverResponse;
use codex_exec_server::FileSystemSandboxContext;
use codex_exec_server::InitializeParams;
use codex_exec_server::InitializeResponse;
use codex_exec_server_protocol::CapabilityRootDiscoverRequest;
use codex_exec_server_protocol::JSONRPCMessage;
use codex_exec_server_protocol::JSONRPCResponse;
#[cfg(unix)]
use codex_protocol::models::PermissionProfile;
#[cfg(unix)]
use codex_protocol::permissions::FileSystemAccessMode;
#[cfg(unix)]
use codex_protocol::permissions::FileSystemSandboxEntry;
#[cfg(unix)]
use codex_protocol::permissions::FileSystemSandboxPolicy;
#[cfg(unix)]
use codex_protocol::permissions::NetworkSandboxPolicy;
#[cfg(unix)]
use codex_utils_absolute_path::AbsolutePathBuf;
use codex_utils_path_uri::PathUri;
use common::exec_server::exec_server;
#[cfg(target_os = "linux")]
use common::exec_server::exec_server_with_env;
#[cfg(target_os = "linux")]
use fake_bwrap::write_fake_bwrap;
use pretty_assertions::assert_eq;

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn discovers_a_complete_capability_bundle_in_one_request() -> anyhow::Result<()> {
    let root = tempfile::tempdir()?;
    write_file(
        &root.path().join(".codex-plugin/plugin.json"),
        r#"{
  "name": "demo",
  "interface": {"displayName": "Demo Plugin"},
  "mcpServers": "./config/mcp.json",
  "apps": "./config/apps.json"
}"#,
    )?;
    write_file(
        &root.path().join(".claude-plugin/plugin.json"),
        r#"{"name":"lower-priority-claude"}"#,
    )?;
    write_file(
        &root.path().join(".cursor-plugin/plugin.json"),
        r#"{"name":"lower-priority-cursor"}"#,
    )?;
    write_file(
        &root.path().join("config/mcp.json"),
        r#"{"mcpServers":{"demo":{"command":"demo-server"}}}"#,
    )?;
    write_file(
        &root.path().join("config/apps.json"),
        r#"{"apps":{"demo":{"connector_id":"connector-demo"}}}"#,
    )?;
    write_file(
        &root.path().join("skills/deploy/SKILL.md"),
        "---\nname: deploy\ndescription: Deploy the service.\n---\n\nDeploy instructions.\n",
    )?;
    write_file(
        &root.path().join("skills/deploy/agents/openai.yaml"),
        "policy:\n  allow_implicit_invocation: false\n",
    )?;
    write_file(
        &root.path().join("nested/.claude-plugin/plugin.json"),
        r#"{"name":"nested"}"#,
    )?;
    write_file(
        &root.path().join("nested/skills/audit/SKILL.md"),
        "---\nname: audit\ndescription: Audit the service.\n---\n",
    )?;
    write_file(
        &root.path().join("nested-cursor/.cursor-plugin/plugin.json"),
        r#"{"name":"cursor-nested"}"#,
    )?;
    write_file(
        &root.path().join("nested-cursor/skills/review/SKILL.md"),
        "---\nname: review\ndescription: Review the service.\n---\n",
    )?;

    let mut server = exec_server().await?;
    initialize(&mut server).await?;
    let root_uri = PathUri::from_host_native_path(root.path())?;
    let discovery = discover_root(&mut server, "demo@1", root_uri.clone()).await?;

    assert_eq!(discovery.id, "demo@1");
    assert_eq!(discovery.path, root_uri);
    assert_eq!(discovery.error, None);
    assert_eq!(discovery.warnings, Vec::<String>::new());
    let plugin = discovery.plugin.as_ref().expect("root plugin");
    assert_eq!(
        plugin.manifest.path,
        root_uri.join(".codex-plugin/plugin.json")?
    );
    assert!(plugin.manifest.contents.contains("Demo Plugin"));
    assert_eq!(
        plugin.mcp_config.as_ref().map(|file| &file.path),
        Some(&root_uri.join("config/mcp.json")?)
    );
    assert_eq!(
        plugin.apps_config.as_ref().map(|file| &file.path),
        Some(&root_uri.join("config/apps.json")?)
    );
    assert_eq!(
        discovery
            .namespace_manifests
            .iter()
            .map(|file| file.path.clone())
            .collect::<Vec<_>>(),
        vec![
            root_uri.join(".codex-plugin/plugin.json")?,
            root_uri.join("nested/.claude-plugin/plugin.json")?,
            root_uri.join("nested-cursor/.cursor-plugin/plugin.json")?,
        ]
    );
    assert_eq!(
        discovery
            .skills
            .iter()
            .map(|skill| (
                skill.instructions.path.clone(),
                skill
                    .metadata
                    .as_ref()
                    .map(|metadata| metadata.path.clone()),
            ))
            .collect::<Vec<_>>(),
        vec![
            (root_uri.join("nested-cursor/skills/review/SKILL.md")?, None,),
            (root_uri.join("nested/skills/audit/SKILL.md")?, None,),
            (
                root_uri.join("skills/deploy/SKILL.md")?,
                Some(root_uri.join("skills/deploy/agents/openai.yaml")?),
            ),
        ]
    );

    server.shutdown().await?;
    Ok(())
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn discovers_cursor_plugin_without_reading_default_mcp_for_inline_servers()
-> anyhow::Result<()> {
    let root = tempfile::tempdir()?;
    write_file(
        &root.path().join(".cursor-plugin/plugin.json"),
        r#"{"name":"cursor-demo","mcpServers":{"inline":{"command":"inline"}}}"#,
    )?;
    write_file(
        &root.path().join(".mcp.json"),
        r#"{"mcpServers":{"should-not-load":{"command":"wrong"}}}"#,
    )?;

    let mut server = exec_server().await?;
    initialize(&mut server).await?;
    let root_uri = PathUri::from_host_native_path(root.path())?;
    let discovery = discover_root(&mut server, "cursor@1", root_uri.clone()).await?;

    assert_eq!(discovery.error, None);
    assert_eq!(discovery.warnings, Vec::<String>::new());
    let plugin = discovery.plugin.expect("cursor plugin");
    assert_eq!(
        plugin.manifest.path,
        root_uri.join(".cursor-plugin/plugin.json")?
    );
    assert_eq!(plugin.mcp_config, None);
    assert_eq!(
        discovery
            .namespace_manifests
            .iter()
            .map(|manifest| manifest.path.clone())
            .collect::<Vec<_>>(),
        vec![root_uri.join(".cursor-plugin/plugin.json")?]
    );

    server.shutdown().await?;
    Ok(())
}

#[cfg(unix)]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn sandboxed_discovery_batches_roots_without_combining_different_permissions()
-> anyhow::Result<()> {
    let workspace = tempfile::tempdir()?;
    let first_root = workspace.path().join("first");
    let second_root = workspace.path().join("second");
    write_file(
        &first_root.join("skills/first/SKILL.md"),
        "---\nname: first\ndescription: First skill.\n---\n",
    )?;
    write_file(
        &second_root.join("skills/second/SKILL.md"),
        "---\nname: second\ndescription: Second skill.\n---\n",
    )?;

    let workspace_uri = PathUri::from_host_native_path(workspace.path())?;
    let first_uri = PathUri::from_host_native_path(&first_root)?;
    let second_uri = PathUri::from_host_native_path(&second_root)?;
    let read_workspace = FileSystemSandboxEntry::new(
        AbsolutePathBuf::from_absolute_path(workspace.path())?.into(),
        FileSystemAccessMode::Read,
    );
    let policy = FileSystemSandboxPolicy::restricted(vec![read_workspace]);
    let shared_sandbox = FileSystemSandboxContext::from_permission_profile_with_cwd(
        PermissionProfile::from_runtime_permissions(&policy, NetworkSandboxPolicy::Restricted),
        workspace_uri,
    );

    #[cfg(target_os = "linux")]
    let fake_bwrap_directory = tempfile::tempdir()?;
    #[cfg(target_os = "linux")]
    let (mut server, fake_bwrap) = {
        let fake_bin_dir = fake_bwrap_directory.path().to_path_buf();
        let fake_bwrap = write_fake_bwrap(&fake_bin_dir)?;
        let mut path_entries = vec![fake_bin_dir];
        if let Some(path) = std::env::var_os("PATH") {
            path_entries.extend(std::env::split_paths(&path));
        }
        let helper_path = std::env::join_paths(path_entries)?;
        (
            exec_server_with_env([("PATH", helper_path.as_os_str())], &[]).await?,
            fake_bwrap,
        )
    };
    #[cfg(not(target_os = "linux"))]
    let mut server = exec_server().await?;
    initialize(&mut server).await?;
    let response = discover_roots(
        &mut server,
        vec![
            CapabilityRootDiscoverRequest {
                id: "first".to_string(),
                path: first_uri.clone(),
                sandbox: Some(shared_sandbox.clone()),
            },
            CapabilityRootDiscoverRequest {
                id: "second".to_string(),
                path: second_uri.clone(),
                sandbox: Some(shared_sandbox.clone()),
            },
        ],
    )
    .await?;
    assert_eq!(
        response
            .roots
            .into_iter()
            .map(|root| (
                root.id,
                root.path,
                root.skills
                    .into_iter()
                    .map(|skill| skill.instructions.path)
                    .collect::<Vec<_>>(),
                root.error,
            ))
            .collect::<Vec<_>>(),
        vec![
            (
                "first".to_string(),
                first_uri.clone(),
                vec![first_uri.join("skills/first/SKILL.md")?],
                None,
            ),
            (
                "second".to_string(),
                second_uri.clone(),
                vec![second_uri.join("skills/second/SKILL.md")?],
                None,
            ),
        ]
    );

    #[cfg(target_os = "linux")]
    {
        let launch_log = fake_bwrap.with_file_name("bwrap.log");
        let launch_count = std::fs::read_to_string(&launch_log)
            .with_context(|| format!("expected fake bwrap launch log at {}", launch_log.display()))?
            .lines()
            .count();
        assert_eq!(launch_count, 1);

        std::fs::write(fake_bwrap.with_file_name("bwrap.fail-once"), "")?;
        let fallback = discover_roots(
            &mut server,
            vec![
                CapabilityRootDiscoverRequest {
                    id: "fallback-first".to_string(),
                    path: first_uri.clone(),
                    sandbox: Some(shared_sandbox.clone()),
                },
                CapabilityRootDiscoverRequest {
                    id: "fallback-second".to_string(),
                    path: second_uri.clone(),
                    sandbox: Some(shared_sandbox.clone()),
                },
            ],
        )
        .await?;
        assert_eq!(
            fallback
                .roots
                .into_iter()
                .map(|root| (root.id, root.skills.len(), root.error))
                .collect::<Vec<_>>(),
            vec![
                ("fallback-first".to_string(), 1, None),
                ("fallback-second".to_string(), 1, None),
            ]
        );
        assert!(std::fs::read_to_string(&launch_log)?.lines().count() > 2);

        server.shutdown().await?;
        server = exec_server().await?;
        initialize(&mut server).await?;
    }

    let read_first_root = FileSystemSandboxEntry::new(
        AbsolutePathBuf::from_absolute_path(&first_root)?.into(),
        FileSystemAccessMode::Read,
    );
    let first_policy = FileSystemSandboxPolicy::restricted(vec![read_first_root]);
    let first_only_sandbox = FileSystemSandboxContext::from_permission_profile_with_cwd(
        PermissionProfile::from_runtime_permissions(
            &first_policy,
            NetworkSandboxPolicy::Restricted,
        ),
        first_uri.clone(),
    );
    let read_second_root = FileSystemSandboxEntry::new(
        AbsolutePathBuf::from_absolute_path(&second_root)?.into(),
        FileSystemAccessMode::Read,
    );
    let second_policy = FileSystemSandboxPolicy::restricted(vec![read_second_root]);
    let second_only_sandbox = FileSystemSandboxContext::from_permission_profile_with_cwd(
        PermissionProfile::from_runtime_permissions(
            &second_policy,
            NetworkSandboxPolicy::Restricted,
        ),
        second_uri.clone(),
    );
    let response = discover_roots(
        &mut server,
        vec![
            CapabilityRootDiscoverRequest {
                id: "first-isolated".to_string(),
                path: first_uri,
                sandbox: Some(first_only_sandbox),
            },
            CapabilityRootDiscoverRequest {
                id: "second-isolated".to_string(),
                path: second_uri,
                sandbox: Some(second_only_sandbox),
            },
        ],
    )
    .await?;
    assert_eq!(
        response
            .roots
            .into_iter()
            .map(|root| (root.id, root.skills.len(), root.error))
            .collect::<Vec<_>>(),
        vec![
            ("first-isolated".to_string(), 1, None),
            ("second-isolated".to_string(), 1, None),
        ]
    );

    server.shutdown().await?;
    Ok(())
}

#[cfg(unix)]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn sandboxed_discovery_follows_only_permitted_external_symlinks() -> anyhow::Result<()> {
    let root = tempfile::tempdir()?;
    let external = tempfile::tempdir()?;
    write_file(
        &root.path().join(".codex-plugin/plugin.json"),
        r#"{"name":"linked-plugin","mcpServers":"./external-mcp.json"}"#,
    )?;
    write_file(
        &external.path().join("mcp.json"),
        r#"{"mcpServers":{"linked":{"command":"linked-server"}}}"#,
    )?;
    write_file(
        &external.path().join("skill/SKILL.md"),
        "---\nname: linked\ndescription: Linked external skill.\n---\n",
    )?;
    std::fs::create_dir_all(root.path().join("skills"))?;
    std::os::unix::fs::symlink(
        external.path().join("skill"),
        root.path().join("skills/linked"),
    )?;
    std::os::unix::fs::symlink(
        external.path().join("mcp.json"),
        root.path().join("external-mcp.json"),
    )?;

    let mut server = exec_server().await?;
    initialize(&mut server).await?;
    let root_uri = PathUri::from_host_native_path(root.path())?;
    let root_path = AbsolutePathBuf::from_absolute_path(root.path())?;
    let external_root = AbsolutePathBuf::from_absolute_path(external.path())?;
    let path_entry =
        |path: AbsolutePathBuf, access| FileSystemSandboxEntry::new(path.into(), access);
    let read_root = path_entry(root_path, FileSystemAccessMode::Read);
    let read_external = path_entry(external_root.clone(), FileSystemAccessMode::Read);
    let deny_external_skill = path_entry(external_root.join("skill"), FileSystemAccessMode::Deny);
    let cases = [
        (
            "permitted symlinks",
            vec![read_root.clone(), read_external.clone()],
            true,
            true,
        ),
        (
            "denied external root",
            vec![read_root.clone()],
            false,
            false,
        ),
        (
            "denied external skill",
            vec![read_root, read_external, deny_external_skill],
            false,
            true,
        ),
    ];

    for (scenario, entries, has_skill, has_mcp) in cases {
        let policy = FileSystemSandboxPolicy::restricted(entries);
        let sandbox = FileSystemSandboxContext::from_permission_profile_with_cwd(
            PermissionProfile::from_runtime_permissions(&policy, NetworkSandboxPolicy::Restricted),
            root_uri.clone(),
        );
        let discovery =
            discover_root_with_sandbox(&mut server, "linked", root_uri.clone(), Some(sandbox))
                .await?;

        assert_eq!(discovery.error, None, "{scenario}");
        assert_eq!(discovery.skills.len(), usize::from(has_skill), "{scenario}");
        assert_eq!(
            discovery
                .plugin
                .and_then(|plugin| plugin.mcp_config)
                .is_some_and(|config| config.contents.contains("linked-server")),
            has_mcp,
            "{scenario}"
        );
    }

    server.shutdown().await?;
    Ok(())
}

async fn discover_root(
    server: &mut common::exec_server::ExecServerHarness,
    id: &str,
    path: PathUri,
) -> anyhow::Result<CapabilityRootDiscovery> {
    discover_root_with_sandbox(server, id, path, /*sandbox*/ None).await
}

async fn discover_root_with_sandbox(
    server: &mut common::exec_server::ExecServerHarness,
    id: &str,
    path: PathUri,
    sandbox: Option<FileSystemSandboxContext>,
) -> anyhow::Result<CapabilityRootDiscovery> {
    let response = discover_roots(
        server,
        vec![CapabilityRootDiscoverRequest {
            id: id.to_string(),
            path,
            sandbox,
        }],
    )
    .await?;
    let [discovery] = response.roots.as_slice() else {
        anyhow::bail!("expected exactly one discovered root");
    };
    Ok(discovery.clone())
}

async fn discover_roots(
    server: &mut common::exec_server::ExecServerHarness,
    roots: Vec<CapabilityRootDiscoverRequest>,
) -> anyhow::Result<CapabilityRootsDiscoverResponse> {
    let request_id = server
        .send_request(
            CAPABILITY_ROOTS_DISCOVER_METHOD,
            serde_json::to_value(CapabilityRootsDiscoverParams { roots })?,
        )
        .await?;
    let response = server.next_event().await?;
    let JSONRPCMessage::Response(JSONRPCResponse { id, result }) = response else {
        anyhow::bail!("expected discovery response, received {response:?}");
    };
    assert_eq!(id, request_id);
    Ok(serde_json::from_value(result)?)
}

async fn initialize(server: &mut common::exec_server::ExecServerHarness) -> anyhow::Result<()> {
    let initialize_id = server
        .send_request(
            "initialize",
            serde_json::to_value(InitializeParams {
                client_name: "capability-discovery-test".to_string(),
                resume_session_id: None,
            })?,
        )
        .await?;
    let response = server
        .wait_for_event(|event| {
            matches!(event, JSONRPCMessage::Response(response) if response.id == initialize_id)
        })
        .await?;
    let JSONRPCMessage::Response(JSONRPCResponse { result, .. }) = response else {
        unreachable!("wait predicate only accepts a response");
    };
    let _: InitializeResponse = serde_json::from_value(result)?;
    server
        .send_notification("initialized", serde_json::json!({}))
        .await?;
    Ok(())
}

fn write_file(path: &std::path::Path, contents: &str) -> anyhow::Result<()> {
    let parent = path
        .parent()
        .ok_or_else(|| anyhow::anyhow!("test file should have a parent"))?;
    std::fs::create_dir_all(parent)?;
    std::fs::write(path, contents)?;
    Ok(())
}