mentra 0.14.0

//! Bridge that wraps MCP server tools as Mentra `ExecutableTool` instances.

use std::sync::Arc;

use async_trait::async_trait;
use serde_json::{Value, json};

use crate::tool::{
    ParallelToolContext, RuntimeToolDescriptor, ToolApprovalCategory, ToolCapability,
    ToolDefinition, ToolDurability, ToolExecutionCategory, ToolExecutor, ToolResult,
    ToolSideEffectLevel,
};

use super::client::McpStdioClient;
use super::protocol::{McpToolCallResult, McpToolDefinition};
use super::sse::client::McpSseClient;

/// The transport-independent surface [`McpBridgedTool`] needs from a client.
///
/// Each transport reports failures with its own error type, so this trait
/// flattens them to a message rather than forcing a shared error enum on the
/// public clients.
///
/// This is a sealed trait: it is public only so that
/// [`McpBridgedTool::new`] can be generic over the transport, and it is not
/// implementable outside this crate.
#[async_trait]
pub trait McpToolClient: sealed::Sealed + Send + Sync {
    /// Calls one tool, rendering any transport failure as a message.
    async fn call_tool(
        &self,
        tool_name: &str,
        arguments: Option<Value>,
    ) -> Result<McpToolCallResult, String>;
}

mod sealed {
    /// Prevents outside implementations of [`super::McpToolClient`].
    pub trait Sealed {}

    impl Sealed for super::McpStdioClient {}
    impl Sealed for super::McpSseClient {}

    #[cfg(test)]
    impl Sealed for crate::mcp::tests::SuccessfulMcpClient {}
}

#[async_trait]
impl McpToolClient for McpStdioClient {
    async fn call_tool(
        &self,
        tool_name: &str,
        arguments: Option<Value>,
    ) -> Result<McpToolCallResult, String> {
        McpStdioClient::call_tool(self, tool_name, arguments)
            .await
            .map_err(|error| error.to_string())
    }
}

#[async_trait]
impl McpToolClient for McpSseClient {
    async fn call_tool(
        &self,
        tool_name: &str,
        arguments: Option<Value>,
    ) -> Result<McpToolCallResult, String> {
        McpSseClient::call_tool(self, tool_name, arguments)
            .await
            .map_err(|error| error.to_string())
    }
}

/// Prefix applied to MCP tool names to namespace them.
const MCP_TOOL_PREFIX: &str = "mcp__";

/// Construct the namespaced tool name for an MCP tool.
pub fn mcp_tool_name(server_name: &str, tool_name: &str) -> String {
    format!("{MCP_TOOL_PREFIX}{server_name}__{tool_name}")
}

/// Parse a namespaced MCP tool name back into `(server_name, tool_name)`.
pub fn parse_mcp_tool_name(name: &str) -> Option<(&str, &str)> {
    let rest = name.strip_prefix(MCP_TOOL_PREFIX)?;
    let (server, tool) = rest.split_once("__")?;
    Some((server, tool))
}

/// A Mentra tool backed by an MCP server tool.
pub struct McpBridgedTool {
    server_name: String,
    tool_def: McpToolDefinition,
    client: Arc<dyn McpToolClient>,
}

impl McpBridgedTool {
    /// Wraps one tool from a connected MCP server.
    ///
    /// The client is generic over the transport, so this accepts an
    /// `Arc<McpStdioClient>` and an `Arc<McpSseClient>` alike.
    pub fn new<C>(server_name: String, tool_def: McpToolDefinition, client: Arc<C>) -> Self
    where
        C: McpToolClient + 'static,
    {
        Self::from_client(server_name, tool_def, client)
    }

    fn from_client(
        server_name: String,
        tool_def: McpToolDefinition,
        client: Arc<dyn McpToolClient>,
    ) -> Self {
        Self {
            server_name,
            tool_def,
            client,
        }
    }

    #[cfg(test)]
    pub(crate) fn new_for_test(
        server_name: String,
        tool_def: McpToolDefinition,
        client: Arc<dyn McpToolClient>,
    ) -> Self {
        Self::from_client(server_name, tool_def, client)
    }

    fn full_name(&self) -> String {
        mcp_tool_name(&self.server_name, &self.tool_def.name)
    }
}

impl std::fmt::Debug for McpBridgedTool {
    /// Renders the bridged identity without reaching into the client, which
    /// holds transport credentials.
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("McpBridgedTool")
            .field("name", &self.full_name())
            .finish_non_exhaustive()
    }
}

impl ToolDefinition for McpBridgedTool {
    fn descriptor(&self) -> RuntimeToolDescriptor {
        let description = self.tool_def.description.clone().unwrap_or_default();

        let input_schema = self
            .tool_def
            .input_schema
            .clone()
            .unwrap_or_else(|| json!({"type": "object", "properties": {}}));

        RuntimeToolDescriptor::builder(self.full_name())
            .description(description)
            .input_schema(input_schema)
            .capability(ToolCapability::Custom(format!("mcp:{}", self.server_name)))
            .side_effect_level(ToolSideEffectLevel::External)
            .durability(ToolDurability::Ephemeral)
            .execution_category(ToolExecutionCategory::ExclusiveLocalMutation)
            .approval_category(ToolApprovalCategory::Process)
            .build()
    }
}

#[async_trait]
impl ToolExecutor for McpBridgedTool {
    async fn execute(&self, _ctx: ParallelToolContext, input: Value) -> ToolResult {
        let arguments = if input.is_null()
            || (input.is_object() && input.as_object().is_none_or(|o| o.is_empty()))
        {
            None
        } else {
            Some(input)
        };

        let result = self
            .client
            .call_tool(&self.tool_def.name, arguments)
            .await
            .map_err(|error| format!("MCP tool call failed: {error}"))?;

        // Concatenate text content blocks into the result string.
        let mut output = String::new();
        for block in &result.content {
            if let Some(text) = &block.text {
                if !output.is_empty() {
                    output.push('\n');
                }
                output.push_str(text);
            }
        }

        if result.is_error {
            Err(output)
        } else {
            Ok(output)
        }
    }
}