Skip to main content

mermaid_cli/mcp/
client.rs

1//! MCP protocol client — higher-level API built on StdioTransport.
2//!
3//! Implements the three protocol methods we need:
4//! - `initialize` — handshake and capability negotiation
5//! - `tools/list` — discover available tools
6//! - `tools/call` — invoke a tool and get results
7
8use anyhow::{Result, anyhow};
9use serde::{Deserialize, Serialize};
10use serde_json::{Value, json};
11
12use super::transport::StdioTransport;
13
14/// MCP protocol client for a single server connection.
15pub struct McpClient {
16    transport: StdioTransport,
17    /// Server info from initialization
18    pub server_info: Option<ServerInfo>,
19}
20
21/// Info returned by the server during initialization
22#[derive(Debug, Clone, Serialize, Deserialize)]
23pub struct ServerInfo {
24    pub name: String,
25    pub version: Option<String>,
26}
27
28/// A tool definition discovered from an MCP server
29#[derive(Debug, Clone)]
30pub struct McpToolDef {
31    pub name: String,
32    pub description: String,
33    pub input_schema: Value,
34}
35
36/// Result of calling an MCP tool
37#[derive(Debug, Clone)]
38pub struct McpToolResult {
39    pub content: Vec<ContentBlock>,
40    pub is_error: bool,
41}
42
43/// A content block in an MCP tool result
44#[derive(Debug, Clone)]
45pub enum ContentBlock {
46    Text(String),
47    Image { data: String, mime_type: String },
48}
49
50impl McpClient {
51    /// Create a new MCP client wrapping a transport.
52    pub fn new(transport: StdioTransport) -> Self {
53        Self {
54            transport,
55            server_info: None,
56        }
57    }
58
59    /// Perform the MCP initialization handshake.
60    ///
61    /// Sends `initialize` request with our client info and protocol version,
62    /// then sends `notifications/initialized` to signal readiness.
63    pub async fn initialize(&mut self) -> Result<ServerInfo> {
64        let result = self
65            .transport
66            .send_request(
67                "initialize",
68                json!({
69                    "protocolVersion": "2025-03-26",
70                    "capabilities": {},
71                    "clientInfo": {
72                        "name": "mermaid",
73                        "version": env!("CARGO_PKG_VERSION"),
74                    }
75                }),
76            )
77            .await?;
78
79        // Parse server info
80        let server_info = ServerInfo {
81            name: result
82                .pointer("/serverInfo/name")
83                .and_then(|v| v.as_str())
84                .unwrap_or("unknown")
85                .to_string(),
86            version: result
87                .pointer("/serverInfo/version")
88                .and_then(|v| v.as_str())
89                .map(|s| s.to_string()),
90        };
91
92        // Send initialized notification
93        self.transport
94            .send_notification("notifications/initialized", json!({}))
95            .await?;
96
97        self.server_info = Some(server_info.clone());
98        Ok(server_info)
99    }
100
101    /// Discover all tools available from this server.
102    pub async fn list_tools(&self) -> Result<Vec<McpToolDef>> {
103        let result = self
104            .transport
105            .send_request("tools/list", json!({}))
106            .await?;
107
108        let tools_array = result
109            .get("tools")
110            .and_then(|v| v.as_array())
111            .ok_or_else(|| anyhow!("MCP tools/list response missing 'tools' array"))?;
112
113        let mut tools = Vec::new();
114        for tool in tools_array {
115            let name = tool
116                .get("name")
117                .and_then(|v| v.as_str())
118                .unwrap_or("")
119                .to_string();
120            let description = tool
121                .get("description")
122                .and_then(|v| v.as_str())
123                .unwrap_or("")
124                .to_string();
125            let input_schema = tool
126                .get("inputSchema")
127                .cloned()
128                .unwrap_or_else(|| json!({"type": "object", "properties": {}}));
129
130            if !name.is_empty() {
131                tools.push(McpToolDef {
132                    name,
133                    description,
134                    input_schema,
135                });
136            }
137        }
138
139        Ok(tools)
140    }
141
142    /// Call a tool on this server and return the result.
143    pub async fn call_tool(&self, name: &str, arguments: &Value) -> Result<McpToolResult> {
144        let params = json!({
145            "name": name,
146            "arguments": arguments,
147        });
148
149        let result = self.transport.send_request("tools/call", params).await?;
150
151        let is_error = result
152            .get("isError")
153            .and_then(|v| v.as_bool())
154            .unwrap_or(false);
155
156        let content_array = result
157            .get("content")
158            .and_then(|v| v.as_array())
159            .cloned()
160            .unwrap_or_default();
161
162        let mut content = Vec::new();
163        for block in content_array {
164            let block_type = block.get("type").and_then(|v| v.as_str()).unwrap_or("");
165            match block_type {
166                "text" => {
167                    if let Some(text) = block.get("text").and_then(|v| v.as_str()) {
168                        content.push(ContentBlock::Text(text.to_string()));
169                    }
170                },
171                "image" => {
172                    let data = block
173                        .get("data")
174                        .and_then(|v| v.as_str())
175                        .unwrap_or("")
176                        .to_string();
177                    let mime_type = block
178                        .get("mimeType")
179                        .and_then(|v| v.as_str())
180                        .unwrap_or("image/png")
181                        .to_string();
182                    content.push(ContentBlock::Image { data, mime_type });
183                },
184                _ => {
185                    // Unknown content type — treat as text if it has a text field
186                    if let Some(text) = block.get("text").and_then(|v| v.as_str()) {
187                        content.push(ContentBlock::Text(text.to_string()));
188                    }
189                },
190            }
191        }
192
193        Ok(McpToolResult { content, is_error })
194    }
195
196    /// Shut down the transport (kills the server process).
197    pub async fn shutdown(&self) {
198        self.transport.shutdown().await;
199    }
200}