codex/codex-rs/codex-mcp/src/binding.rs

//! Immutable MCP catalog and execution handles.

#[path = "model_search.rs"]
mod model_search;
pub use model_search::SEARCH_CONNECTOR;
pub use model_search::SearchDiscovery;
pub use model_search::is_model_search_tool;

use std::collections::HashMap;
use std::fmt;
use std::future::Future;
use std::sync::Arc;
use std::time::Duration;

use anyhow::Context;
use anyhow::Result;
use codex_config::AppToolApproval;
use codex_exec_server::FileSystemSandboxContext;
use codex_protocol::mcp::CallToolResult;
use codex_protocol::models::PermissionProfile;
use codex_utils_absolute_path::AbsolutePathBuf;
use rmcp::model::ListResourceTemplatesResult;
use rmcp::model::ListResourcesResult;
use rmcp::model::PaginatedRequestParams;
use rmcp::model::ReadResourceRequestParams;
use rmcp::model::ReadResourceResult;
use rmcp::model::Resource;
use rmcp::model::ResourceTemplate;
use serde_json::Value as JsonValue;

use crate::McpConfig;
use crate::McpServerPermissions;
use crate::binding_clients::McpBindingClients;
use crate::client_tool_catalog::ToolCatalogSnapshot;
use crate::connection_manager::McpConnectionSet;
use crate::rmcp_client::ManagedClient;
use crate::server::McpServerMetadata;
use crate::tools::ToolInfo;

/// The exact tool catalog and execution handles shared by compatible sampling steps.
pub struct McpBinding {
    connections: Arc<McpConnectionSet>,
    clients: Arc<McpBindingClients>,
    config: Arc<McpConfig>,
    plugins_available: bool,
    tools: Arc<[ToolInfo]>,
    calls: HashMap<(String, String), PreparedMcpCall>,
}

impl McpBinding {
    /// Creates an empty binding for tests and callers without a materialized runtime.
    pub fn empty(config: Arc<McpConfig>) -> Self {
        Self::new(
            Arc::new(McpConnectionSet::empty(config.prefix_mcp_tool_names)),
            Arc::new(McpBindingClients::new(HashMap::new())),
            config,
            /*plugins_available*/ false,
            Vec::new(),
            HashMap::new(),
        )
    }

    pub(crate) fn new(
        connections: Arc<McpConnectionSet>,
        clients: Arc<McpBindingClients>,
        config: Arc<McpConfig>,
        plugins_available: bool,
        tools: impl Into<Arc<[ToolInfo]>>,
        calls: HashMap<(String, String), PreparedMcpCall>,
    ) -> Self {
        Self {
            connections,
            clients,
            config,
            plugins_available,
            tools: tools.into(),
            calls,
        }
    }

    pub fn config(&self) -> &Arc<McpConfig> {
        &self.config
    }

    pub fn plugins_available(&self) -> bool {
        self.plugins_available
    }

    /// Returns the frozen model-visible catalog captured for this binding.
    pub fn tools(&self) -> &[ToolInfo] {
        &self.tools
    }

    /// Returns permitted tool metadata, including app-only tools.
    pub fn tool_info(&self, server: &str, tool: &str) -> Option<&ToolInfo> {
        self.calls
            .get(&(server.to_string(), tool.to_string()))
            .map(PreparedMcpCall::tool_info)
    }

    /// Binds a model-visible call to the exact client and metadata in this binding.
    pub fn prepare_call(&self, server: &str, tool: &str) -> Option<PreparedMcpCall> {
        self.calls
            .get(&(server.to_string(), tool.to_string()))
            .filter(|call| crate::tool_is_model_visible(call.tool_info()))
            .cloned()
    }

    pub fn has_servers(&self) -> bool {
        self.connections.has_servers()
    }

    /// The environment owning this server in the step's connection set, even if it exposes no tools.
    pub fn server_environment_id(&self, server: &str) -> Option<&str> {
        self.connections.server_environment_id(server)
    }

    pub async fn list_resources(
        &self,
        server: &str,
        params: Option<PaginatedRequestParams>,
    ) -> Result<ListResourcesResult> {
        if self.clients.client(server).is_some() {
            self.clients.list_resources(server, params).await
        } else {
            self.connections.list_resources(server, params).await
        }
    }

    pub async fn list_all_resources(
        &self,
        include_server: impl Fn(&str) -> bool,
    ) -> HashMap<String, Vec<Resource>> {
        self.clients.list_all_resources(include_server).await
    }

    pub async fn list_resource_templates(
        &self,
        server: &str,
        params: Option<PaginatedRequestParams>,
    ) -> Result<ListResourceTemplatesResult> {
        if self.connections.skips_resource_template_listing(server) {
            return Ok(ListResourceTemplatesResult::with_all_items(Vec::new()));
        }
        if self.clients.client(server).is_some() {
            self.clients.list_resource_templates(server, params).await
        } else {
            self.connections
                .list_resource_templates(server, params)
                .await
        }
    }

    pub async fn list_all_resource_templates(
        &self,
        include_server: impl Fn(&str) -> bool,
    ) -> HashMap<String, Vec<ResourceTemplate>> {
        self.clients
            .list_all_resource_templates(|server| {
                include_server(server) && !self.connections.skips_resource_template_listing(server)
            })
            .await
    }

    pub async fn read_resource(
        &self,
        server: &str,
        params: ReadResourceRequestParams,
    ) -> Result<ReadResourceResult> {
        if self.clients.client(server).is_some() {
            self.clients.read_resource(server, params).await
        } else {
            self.connections.read_resource(server, params).await
        }
    }
}

impl fmt::Debug for McpBinding {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("McpBinding")
            .field("tools", &self.tools)
            .field("prepared_call_count", &self.calls.len())
            .finish_non_exhaustive()
    }
}

/// A call bound to an exact client, catalog revision, timeout, and permission authority.
#[derive(Clone)]
pub struct PreparedMcpCall {
    connections: Arc<McpConnectionSet>,
    client: Arc<ManagedClient>,
    config: Arc<McpConfig>,
    catalog_snapshot: Arc<ToolCatalogSnapshot>,
    tool_info: PreparedToolInfo,
    server_name: String,
    server_metadata: McpServerMetadata,
    plugin_id: Option<String>,
    selected_plugin_server: bool,
}

/// The rules for a prepared call and the folders used to interpret them.
pub enum McpPermissionSource<'a> {
    Environment {
        environment_id: &'a str,
        sandbox: &'a Arc<FileSystemSandboxContext>,
    },
    Thread {
        profile: &'a PermissionProfile,
        host_cwd: &'a AbsolutePathBuf,
    },
}

impl<'a> McpPermissionSource<'a> {
    pub fn profile(&self) -> &'a PermissionProfile {
        match self {
            Self::Environment { sandbox, .. } => &sandbox.permissions,
            Self::Thread { profile, .. } => profile,
        }
    }
}

#[derive(Clone)]
pub(crate) struct PreparedToolInfo {
    catalog: Arc<[ToolInfo]>,
    index: usize,
}

impl PreparedToolInfo {
    pub(crate) fn from_binding_catalog(catalog: Arc<[ToolInfo]>, index: usize) -> Self {
        assert!(
            index < catalog.len(),
            "prepared tool index must refer to its binding catalog"
        );
        Self { catalog, index }
    }

    pub(crate) fn get(&self) -> &ToolInfo {
        &self.catalog[self.index]
    }
}

impl From<ToolInfo> for PreparedToolInfo {
    fn from(tool_info: ToolInfo) -> Self {
        Self {
            catalog: Arc::new([tool_info]),
            index: 0,
        }
    }
}

impl PreparedMcpCall {
    #[expect(
        clippy::too_many_arguments,
        reason = "the exact call authority stays together"
    )]
    pub(crate) fn new(
        connections: Arc<McpConnectionSet>,
        client: Arc<ManagedClient>,
        config: Arc<McpConfig>,
        catalog_snapshot: Arc<ToolCatalogSnapshot>,
        tool_info: impl Into<PreparedToolInfo>,
        server_metadata: McpServerMetadata,
        plugin_id: Option<String>,
        selected_plugin_server: bool,
    ) -> Option<Self> {
        let tool_info = tool_info.into();
        let server_name = tool_info.get().server_name.clone();
        config.permission_profile_for_server(&server_name)?;
        Some(Self {
            connections,
            client,
            config,
            catalog_snapshot,
            tool_info,
            server_name,
            server_metadata,
            plugin_id,
            selected_plugin_server,
        })
    }

    pub fn tool_info(&self) -> &ToolInfo {
        self.tool_info.get()
    }

    /// Returns the configuration and approval authority captured with this client.
    pub fn config(&self) -> &McpConfig {
        &self.config
    }

    /// A hosted app can run through the local server but use thread rules, not local environment rules.
    pub fn permission_source(&self) -> McpPermissionSource<'_> {
        // The constructor only prepares calls with an entry in this immutable config.
        match &self.config.server_permissions[&self.server_name] {
            McpServerPermissions::Environment(sandbox) => McpPermissionSource::Environment {
                environment_id: self.server_environment_id(),
                sandbox,
            },
            McpServerPermissions::Thread => McpPermissionSource::Thread {
                profile: &self.config.permission_profile,
                host_cwd: &self.config.host_cwd,
            },
        }
    }

    /// Returns the owner permissions validated when this immutable call was prepared.
    pub fn permission_profile(&self) -> &PermissionProfile {
        self.permission_source().profile()
    }

    /// Attach policy aliases and the server-host launcher, replacing stale authority.
    /// Callers negotiate the state capability before requesting this metadata.
    pub async fn augment_sandbox_metadata(
        &self,
        meta: &mut Option<JsonValue>,
        state: crate::SandboxState,
    ) -> Result<()> {
        crate::sandbox_metadata::augment_metadata(
            meta,
            state,
            self.sandbox_codex_executable().await,
        )
    }

    /// Full CLI path usable by a local stdio server. An HTTP endpoint (even on
    /// loopback) or remote executor does not establish a shared filesystem.
    /// Capability probes are cached process-wide by absolute executable path.
    pub async fn sandbox_codex_executable(&self) -> Option<AbsolutePathBuf> {
        let server = self
            .config
            .mcp_server_catalog
            .server(&self.server_name)?
            .config();
        if self.server_environment_id() != codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID
            || !server.is_local_environment()
            || !matches!(
                server.transport,
                codex_config::McpServerTransportConfig::Stdio { .. }
            )
        {
            return None;
        }

        // Require a literal absolute path before the constructor can expand `~`.
        let executable = self
            .config
            .codex_self_exe
            .as_deref()
            .filter(|path| path.is_absolute() && path.to_str().is_some())?;
        let executable = AbsolutePathBuf::from_absolute_path_checked(executable).ok()?;
        crate::sandbox_executable::supports_sandbox_state(&executable)
            .await
            .then_some(executable)
    }

    pub fn server_name(&self) -> &str {
        &self.server_name
    }

    /// Returns whether this call is bound to the host-owned Codex Apps server.
    pub fn is_host_owned_apps(&self) -> bool {
        self.config
            .mcp_server_catalog
            .server(&self.server_name)
            .is_some_and(|registration| {
                registration
                    .source()
                    .is_host_owned_apps(&self.server_name, registration.config())
            })
    }

    pub fn server_origin(&self) -> Option<&str> {
        self.server_metadata
            .origin
            .as_ref()
            .map(super::server::McpServerOrigin::as_str)
    }

    pub fn server_environment_id(&self) -> &str {
        &self.server_metadata.environment_id
    }

    pub fn server_pollutes_memory(&self) -> bool {
        self.server_metadata.pollutes_memory
    }

    pub fn tool_approval_mode(&self) -> AppToolApproval {
        self.server_metadata
            .tool_approval_mode(&self.tool_info().tool.name)
    }

    /// Returns the explicit output budget captured with this call's effective server config.
    pub fn output_token_limit(&self) -> Option<usize> {
        self.config
            .mcp_server_catalog
            .server(&self.server_name)?
            .config()
            .tools
            .get(self.tool_info().tool.name.as_ref())?
            .output_token_limit
            .map(std::num::NonZeroUsize::get)
    }

    pub fn plugin_id(&self) -> Option<&str> {
        self.plugin_id.as_deref()
    }

    pub fn is_selected_plugin_server(&self) -> bool {
        self.selected_plugin_server
    }

    pub async fn server_supports_sandbox_state_meta_capability(&self) -> Result<bool> {
        Ok(self.client.server_supports_sandbox_state_meta_capability)
    }

    pub async fn call(
        &self,
        arguments: Option<JsonValue>,
        meta: Option<JsonValue>,
        timeout: Option<Duration>,
    ) -> Result<CallToolResult> {
        self.call_with_preparation(timeout, || async move { Ok((arguments, meta)) })
            .await
    }

    /// Runs irreversible call preparation and execution under the authority of
    /// this call's captured catalog and the extensions owned by the Codex session.
    /// A caller-supplied timeout can further restrict the server's configured timeout.
    pub async fn call_with_preparation<F, Fut>(
        &self,
        requested_timeout: Option<Duration>,
        prepare: F,
    ) -> Result<CallToolResult>
    where
        F: FnOnce() -> Fut,
        Fut: Future<Output = Result<(Option<JsonValue>, Option<JsonValue>)>>,
    {
        let effective_timeout = match (self.client.tool_timeout, requested_timeout) {
            (Some(server_timeout), Some(requested_timeout)) => {
                Some(server_timeout.min(requested_timeout))
            }
            (server_timeout, requested_timeout) => server_timeout.or(requested_timeout),
        };
        let tool_name = self.tool_info().tool.name.to_string();
        let operation =
            self.client
                .tool_catalog
                .run_with_snapshot(&self.catalog_snapshot, || async {
                    let (arguments, meta) = prepare().await?;
                    let timeout_deadline =
                        effective_timeout.map(|timeout| tokio::time::Instant::now() + timeout);
                    let add_trusted_access_context = self.connections.add_trusted_access_context(
                        self.tool_info(),
                        &self.server_metadata,
                        arguments.as_ref(),
                        meta,
                    );
                    let meta = match effective_timeout.zip(timeout_deadline) {
                        Some((timeout, deadline)) => {
                            tokio::time::timeout_at(deadline, add_trusted_access_context)
                                .await
                                .map_err(|_| {
                                    codex_rmcp_client::request_timeout_error("tools/call", timeout)
                                })?
                        }
                        None => add_trusted_access_context.await,
                    };
                    let remaining_timeout = match effective_timeout.zip(timeout_deadline) {
                        Some((timeout, deadline)) => {
                            let remaining =
                                deadline.saturating_duration_since(tokio::time::Instant::now());
                            if remaining.is_zero() {
                                return Err(codex_rmcp_client::request_timeout_error(
                                    "tools/call",
                                    timeout,
                                ));
                            }
                            Some(remaining)
                        }
                        None => None,
                    };
                    self.client
                        .call_tool(
                            tool_name.clone(),
                            arguments,
                            self.connections.network_policy_meta_for_call(
                                &self.client,
                                Some(self.permission_profile()),
                                self.server_environment_id(),
                                meta,
                            ),
                            remaining_timeout,
                        )
                        .await
                        .with_context(|| {
                            format!("tool call failed for `{}/{tool_name}`", self.server_name)
                        })
                });
        self.client
            .with_active_shared_lease(operation)
            .await?
            .ok_or_else(|| anyhow::anyhow!(
                "tool call rejected because the catalog changed after `{}/{tool_name}` was prepared",
                self.server_name
            ))?
            .map(call_tool_result_from_rmcp)
    }
}

pub(crate) fn call_tool_result_from_rmcp(result: rmcp::model::CallToolResult) -> CallToolResult {
    let content = result
        .content
        .into_iter()
        .map(|content| {
            serde_json::to_value(content)
                .unwrap_or_else(|_| JsonValue::String("<content>".to_string()))
        })
        .collect();
    CallToolResult {
        content,
        structured_content: result.structured_content,
        is_error: result.is_error,
        meta: result.meta.and_then(|meta| serde_json::to_value(meta).ok()),
    }
}

#[cfg(test)]
#[path = "binding_tests.rs"]
mod tests;