Skip to main content

agent_code_lib/services/mcp/
client.rs

1//! MCP client: high-level interface for MCP server interactions.
2//!
3//! Manages the connection lifecycle, tool discovery, and tool
4//! execution for a single MCP server.
5
6use tracing::{debug, info};
7
8use super::transport::McpTransportConnection;
9use super::types::*;
10
11/// Client for a single MCP server connection.
12pub struct McpClient {
13    /// Server configuration.
14    config: McpServerConfig,
15    /// Transport connection.
16    transport: Option<McpTransportConnection>,
17    /// Discovered tools.
18    tools: Vec<McpTool>,
19    /// Discovered resources.
20    resources: Vec<McpResource>,
21    /// Connection status.
22    status: McpConnectionStatus,
23}
24
25impl McpClient {
26    /// Create a new client for the given server configuration.
27    pub fn new(config: McpServerConfig) -> Self {
28        Self {
29            config,
30            transport: None,
31            tools: Vec::new(),
32            resources: Vec::new(),
33            status: McpConnectionStatus::Disconnected,
34        }
35    }
36
37    /// Get the server name.
38    pub fn name(&self) -> &str {
39        &self.config.name
40    }
41
42    /// Get the connection status.
43    pub fn status(&self) -> &McpConnectionStatus {
44        &self.status
45    }
46
47    /// Get discovered tools.
48    pub fn tools(&self) -> &[McpTool] {
49        &self.tools
50    }
51
52    /// Get discovered resources.
53    pub fn resources(&self) -> &[McpResource] {
54        &self.resources
55    }
56
57    /// Connect to the MCP server and perform initialization.
58    pub async fn connect(&mut self) -> Result<(), String> {
59        self.status = McpConnectionStatus::Connecting;
60
61        let transport = match &self.config.transport {
62            McpTransport::Stdio { command, args } => {
63                McpTransportConnection::connect_stdio(command, args, &self.config.env).await?
64            }
65            McpTransport::Sse { url } => McpTransportConnection::connect_sse(url).await?,
66        };
67
68        // Initialize the connection.
69        let init_result = transport
70            .request(
71                "initialize",
72                Some(serde_json::json!({
73                    "protocolVersion": "2024-11-05",
74                    "capabilities": {
75                        "tools": {},
76                        "resources": {}
77                    },
78                    "clientInfo": {
79                        "name": "agent-code",
80                        "version": env!("CARGO_PKG_VERSION")
81                    }
82                })),
83            )
84            .await?;
85
86        debug!("MCP server initialized: {:?}", init_result);
87
88        // Send initialized notification.
89        transport.notify("notifications/initialized", None).await?;
90
91        self.transport = Some(transport);
92        self.status = McpConnectionStatus::Connected;
93
94        // Discover tools and resources.
95        self.discover_tools().await?;
96        self.discover_resources().await?;
97
98        info!(
99            "MCP server '{}' connected: {} tools, {} resources",
100            self.config.name,
101            self.tools.len(),
102            self.resources.len()
103        );
104
105        Ok(())
106    }
107
108    /// Discover available tools from the server.
109    async fn discover_tools(&mut self) -> Result<(), String> {
110        let transport = self.transport.as_ref().ok_or("Not connected")?;
111
112        let result = transport.request("tools/list", None).await?;
113
114        if let Some(tools) = result.get("tools").and_then(|v| v.as_array()) {
115            self.tools = tools
116                .iter()
117                .filter_map(|t| serde_json::from_value(t.clone()).ok())
118                .collect();
119        }
120
121        Ok(())
122    }
123
124    /// Discover available resources from the server.
125    async fn discover_resources(&mut self) -> Result<(), String> {
126        let transport = self.transport.as_ref().ok_or("Not connected")?;
127
128        match transport.request("resources/list", None).await {
129            Ok(result) => {
130                if let Some(resources) = result.get("resources").and_then(|v| v.as_array()) {
131                    self.resources = resources
132                        .iter()
133                        .filter_map(|r| serde_json::from_value(r.clone()).ok())
134                        .collect();
135                }
136            }
137            Err(e) => {
138                // Resources are optional — server may not support them.
139                debug!(
140                    "MCP server '{}' doesn't support resources: {e}",
141                    self.config.name
142                );
143            }
144        }
145
146        Ok(())
147    }
148
149    /// Call a tool on the MCP server.
150    pub async fn call_tool(
151        &self,
152        tool_name: &str,
153        arguments: serde_json::Value,
154    ) -> Result<McpToolResult, String> {
155        let transport = self.transport.as_ref().ok_or("Not connected")?;
156
157        let result = transport
158            .request(
159                "tools/call",
160                Some(serde_json::json!({
161                    "name": tool_name,
162                    "arguments": arguments,
163                })),
164            )
165            .await?;
166
167        serde_json::from_value(result).map_err(|e| format!("Invalid tool result: {e}"))
168    }
169
170    /// Read a resource from the MCP server.
171    pub async fn read_resource(&self, uri: &str) -> Result<String, String> {
172        let transport = self.transport.as_ref().ok_or("Not connected")?;
173
174        let result = transport
175            .request("resources/read", Some(serde_json::json!({ "uri": uri })))
176            .await?;
177
178        // Extract text content from the response.
179        if let Some(contents) = result.get("contents").and_then(|v| v.as_array()) {
180            let text: Vec<&str> = contents
181                .iter()
182                .filter_map(|c| c.get("text").and_then(|t| t.as_str()))
183                .collect();
184            Ok(text.join("\n"))
185        } else {
186            Ok(result.to_string())
187        }
188    }
189
190    /// Disconnect from the MCP server.
191    pub async fn disconnect(&mut self) {
192        if let Some(transport) = self.transport.take() {
193            transport.shutdown().await;
194        }
195        self.tools.clear();
196        self.resources.clear();
197        self.status = McpConnectionStatus::Disconnected;
198    }
199}