mermaid_cli/mcp/
server_manager.rs

1//! MCP server lifecycle management.
2//!
3//! Manages multiple MCP server processes, handles tool discovery,
4//! and routes tool calls to the correct server.
5
6use anyhow::{Result, anyhow};
7use std::collections::HashMap;
8use tracing::{info, warn};
9
10use super::client::{ContentBlock, McpClient, McpToolDef, McpToolResult};
11use super::transport::StdioTransport;
12use crate::app::McpServerConfig;
13
14/// Manages multiple MCP server connections.
15pub struct McpServerManager {
16    /// Active server connections: server_name → client
17    servers: HashMap<String, McpClient>,
18    /// Cached tool definitions: (server_name, tool_def)
19    tools: Vec<(String, McpToolDef)>,
20}
21
22impl McpServerManager {
23    /// Start all configured MCP servers, initialize them, and discover tools.
24    ///
25    /// Servers that fail to start are logged and skipped (non-fatal).
26    pub async fn start(configs: &HashMap<String, McpServerConfig>) -> Self {
27        let mut servers = HashMap::new();
28        let mut all_tools = Vec::new();
29
30        for (name, config) in configs {
31            info!("Starting MCP server: {} ({} {})", name, config.command, config.args.join(" "));
32
33            match Self::start_one(name, config).await {
34                Ok((client, tools)) => {
35                    let tool_count = tools.len();
36                    for tool in &tools {
37                        all_tools.push((name.clone(), tool.clone()));
38                    }
39                    info!(
40                        "MCP server '{}' ready: {} tools ({})",
41                        name,
42                        tool_count,
43                        client
44                            .server_info
45                            .as_ref()
46                            .map(|s| s.name.as_str())
47                            .unwrap_or("?")
48                    );
49                    servers.insert(name.clone(), client);
50                },
51                Err(e) => {
52                    warn!("Failed to start MCP server '{}': {}", name, e);
53                },
54            }
55        }
56
57        Self {
58            servers,
59            tools: all_tools,
60        }
61    }
62
63    /// Start a single MCP server, initialize, and list tools.
64    async fn start_one(
65        name: &str,
66        config: &McpServerConfig,
67    ) -> Result<(McpClient, Vec<McpToolDef>)> {
68        let transport =
69            StdioTransport::spawn(&config.command, &config.args, &config.env).await?;
70        let mut client = McpClient::new(transport);
71
72        client.initialize().await.map_err(|e| {
73            anyhow!("MCP server '{}' initialization failed: {}", name, e)
74        })?;
75
76        let tools = client.list_tools().await.map_err(|e| {
77            anyhow!("MCP server '{}' tool discovery failed: {}", name, e)
78        })?;
79
80        Ok((client, tools))
81    }
82
83    /// Get all discovered tools with their server names.
84    pub fn get_all_tools(&self) -> &[(String, McpToolDef)] {
85        &self.tools
86    }
87
88    /// Check if any MCP servers are active.
89    pub fn has_servers(&self) -> bool {
90        !self.servers.is_empty()
91    }
92
93    /// Call a tool on a specific server.
94    ///
95    /// # Concurrency
96    ///
97    /// Multiple concurrent calls to the same server will serialize at the
98    /// transport layer (`StdioTransport` holds a mutex over stdin writes and
99    /// uses a shared pending-response map for JSON-RPC correlation). This is
100    /// intentional: JSON-RPC over stdio is a byte stream, and interleaved
101    /// writes would corrupt messages. Calls to *different* servers run fully
102    /// in parallel since each has its own transport.
103    pub async fn call_tool(
104        &self,
105        server_name: &str,
106        tool_name: &str,
107        arguments: &serde_json::Value,
108    ) -> Result<McpToolResult> {
109        let client = self
110            .servers
111            .get(server_name)
112            .ok_or_else(|| anyhow!("MCP server '{}' not found or not running", server_name))?;
113
114        client.call_tool(tool_name, arguments).await
115    }
116
117    /// Convert an MCP tool result into text suitable for a tool result message.
118    /// Images are returned separately for multimodal attachment.
119    pub fn format_tool_result(result: &McpToolResult) -> (String, Option<Vec<String>>) {
120        let mut text_parts = Vec::new();
121        let mut images = Vec::new();
122
123        for block in &result.content {
124            match block {
125                ContentBlock::Text(text) => text_parts.push(text.clone()),
126                ContentBlock::Image { data, .. } => images.push(data.clone()),
127            }
128        }
129
130        let text = if text_parts.is_empty() {
131            if result.is_error {
132                "MCP tool returned an error with no message".to_string()
133            } else {
134                "MCP tool returned no text content".to_string()
135            }
136        } else {
137            text_parts.join("\n")
138        };
139
140        let images = if images.is_empty() {
141            None
142        } else {
143            Some(images)
144        };
145
146        (text, images)
147    }
148
149    /// Gracefully shut down all MCP servers.
150    pub async fn shutdown(&self) {
151        for (name, client) in &self.servers {
152            info!("Shutting down MCP server: {}", name);
153            client.shutdown().await;
154        }
155    }
156}
mermaid_cli/mcp/server_manager.rs

mermaid_cli/mcp/
server_manager.rs
