PicoBot/src/mcp/mod.rs

pub mod tool_wrapper;

use std::collections::HashMap;
use std::sync::{Arc, Mutex};

use anyhow::Context;
use http::{HeaderName, HeaderValue};
use rmcp::model::{CallToolRequestParams, RawContent};
use rmcp::transport::streamable_http_client::StreamableHttpClientTransportConfig;
use rmcp::transport::{StreamableHttpClientTransport, TokioChildProcess};
use rmcp::{Peer, RoleClient, ServiceExt};
use tokio::process::Command;

use crate::config::{McpConfig, McpServerConfig, McpTransport};
use crate::tools::ToolResult;

pub use tool_wrapper::McpToolWrapper;

/// Status of a single MCP tool.
#[derive(Debug, Clone)]
pub struct McpToolStatus {
    pub name: String,
    pub description: String,
}

/// Status of a single MCP server.
#[derive(Debug, Clone)]
pub struct McpServerStatus {
    pub name: String,
    pub transport: String,
    pub connected: bool,
    pub error: Option<String>,
    pub tools: Vec<McpToolStatus>,
}

static MCP_SERVER_STATUS: Mutex<Vec<McpServerStatus>> = Mutex::new(Vec::new());

pub fn get_mcp_status() -> Vec<McpServerStatus> {
    MCP_SERVER_STATUS.lock().unwrap().clone()
}

fn update_mcp_status(servers: Vec<McpServerStatus>) {
    let mut status = MCP_SERVER_STATUS.lock().unwrap();
    *status = servers;
}

/// A connected MCP server. Holds a clonable Peer handle for tool calls,
/// and keeps the underlying service alive via a background task.
pub struct McpConnection {
    peer: Peer<RoleClient>,
    /// Keep the service alive. When dropped, the MCP connection is closed.
    _service: Option<Box<dyn std::any::Any + Send + Sync>>,
}

impl McpConnection {
    pub async fn call_tool(
        &self,
        tool_name: &str,
        arguments: serde_json::Value,
    ) -> anyhow::Result<ToolResult> {
        let result = self
            .peer
            .call_tool(
                CallToolRequestParams::new(tool_name.to_string())
                    .with_arguments(arguments.as_object().cloned().unwrap_or_default()),
            )
            .await
            .context("MCP tool call failed")?;

        let is_error = result.is_error.unwrap_or(false);
        let output = extract_text(&result);

        Ok(ToolResult {
            success: !is_error,
            output,
            error: if is_error {
                Some("MCP server returned an error".to_string())
            } else {
                None
            },
        })
    }
}

fn extract_text(result: &rmcp::model::CallToolResult) -> String {
    let mut parts = Vec::new();
    for content in &result.content {
        match &**content {
            RawContent::Text(text) => {
                parts.push(text.text.clone());
            }
            RawContent::Image(image) => {
                parts.push(format!("[image: {}]", image.mime_type,));
            }
            RawContent::Resource(resource) => match &resource.resource {
                rmcp::model::ResourceContents::TextResourceContents { text, .. } => {
                    parts.push(format!(
                        "[resource text: {}]",
                        text.chars().take(200).collect::<String>(),
                    ));
                }
                rmcp::model::ResourceContents::BlobResourceContents { uri, .. } => {
                    parts.push(format!("[resource blob: {}]", uri));
                }
            },
            _ => {
                parts.push("[unsupported content]".to_string());
            }
        }
    }
    if parts.is_empty() {
        String::new()
    } else {
        parts.join("\n")
    }
}

pub struct ToolInfo {
    pub server_name: String,
    pub tool_name: String,
    pub description: String,
    pub schema: serde_json::Value,
    pub connection: Arc<McpConnection>,
}

pub async fn connect_all(config: &McpConfig) -> Vec<ToolInfo> {
    let mut tools = Vec::new();
    let mut server_statuses = Vec::new();

    for server_config in &config.servers {
        let transport_str = match server_config.transport {
            McpTransport::Stdio => "stdio",
            McpTransport::Sse => "sse",
            McpTransport::StreamableHttp => "streamable-http",
        };

        match connect_server(server_config).await {
            Ok(connection) => {
                let connection = Arc::new(connection);
                match list_tools(&connection).await {
                    Ok(server_tools) => {
                        tracing::info!(
                            server = %server_config.name,
                            count = server_tools.len(),
                            "MCP server connected"
                        );
                        let tool_statuses: Vec<McpToolStatus> = server_tools
                            .iter()
                            .map(|(name, desc, _)| McpToolStatus {
                                name: name.clone(),
                                description: desc.clone(),
                            })
                            .collect();
                        server_statuses.push(McpServerStatus {
                            name: server_config.name.clone(),
                            transport: transport_str.to_string(),
                            connected: true,
                            error: None,
                            tools: tool_statuses,
                        });
                        for (orig_name, desc, schema) in server_tools {
                            tools.push(ToolInfo {
                                server_name: server_config.name.clone(),
                                tool_name: orig_name,
                                description: desc,
                                schema,
                                connection: connection.clone(),
                            });
                        }
                    }
                    Err(e) => {
                        tracing::error!(
                            server = %server_config.name,
                            error = %e,
                            "Failed to list MCP tools"
                        );
                        server_statuses.push(McpServerStatus {
                            name: server_config.name.clone(),
                            transport: transport_str.to_string(),
                            connected: false,
                            error: Some(e.to_string()),
                            tools: Vec::new(),
                        });
                    }
                }
            }
            Err(e) => {
                tracing::error!(
                    server = %server_config.name,
                    error = %e,
                    "Failed to connect to MCP server"
                );
                server_statuses.push(McpServerStatus {
                    name: server_config.name.clone(),
                    transport: transport_str.to_string(),
                    connected: false,
                    error: Some(e.to_string()),
                    tools: Vec::new(),
                });
            }
        }
    }

    update_mcp_status(server_statuses);
    tools
}

async fn connect_server(config: &McpServerConfig) -> anyhow::Result<McpConnection> {
    match config.transport {
        McpTransport::Stdio => {
            let command = config
                .command
                .as_ref()
                .context("stdio transport requires 'command'")?;
            let mut cmd = Command::new(command);
            cmd.args(&config.args);
            for (k, v) in &config.env {
                cmd.env(k, v);
            }

            let service =
                ().serve(
                    TokioChildProcess::new(cmd).context("failed to create stdio MCP transport")?,
                )
                .await
                .context("failed to connect to stdio MCP server")?;

            let peer = service.peer().clone();

            Ok(McpConnection {
                peer,
                _service: Some(Box::new(service)),
            })
        }
        McpTransport::Sse | McpTransport::StreamableHttp => {
            let url = config
                .url
                .as_ref()
                .context("sse/streamable-http transport requires 'url'")?;

            let mut headers_map = HashMap::new();
            for (k, v) in &config.headers {
                if let (Ok(name), Ok(value)) = (
                    HeaderName::from_bytes(k.as_bytes()),
                    HeaderValue::from_str(v),
                ) {
                    headers_map.insert(name, value);
                }
            }

            let transport = if headers_map.is_empty() {
                StreamableHttpClientTransport::from_uri(url.to_string())
            } else {
                StreamableHttpClientTransport::from_config(
                    StreamableHttpClientTransportConfig::with_uri(url.to_string())
                        .custom_headers(headers_map),
                )
            };

            let service =
                ().serve(transport)
                    .await
                    .context("failed to connect to HTTP/SSE MCP server")?;

            let peer = service.peer().clone();

            Ok(McpConnection {
                peer,
                _service: Some(Box::new(service)),
            })
        }
    }
}

async fn list_tools(
    connection: &McpConnection,
) -> anyhow::Result<Vec<(String, String, serde_json::Value)>> {
    let tools = connection
        .peer
        .list_all_tools()
        .await
        .context("failed to list MCP tools")?;

    Ok(tools
        .into_iter()
        .map(|tool| {
            (
                tool.name.to_string(),
                tool.description.map(|d| d.to_string()).unwrap_or_default(),
                serde_json::Value::Object((*tool.input_schema).clone()),
            )
        })
        .collect())
}