Skip to main content

agentic_core/tool/mcp/
client.rs

1use std::collections::HashMap;
2use std::fmt;
3use std::net::SocketAddr;
4use std::process::Stdio;
5use std::sync::Arc;
6use std::time::Duration;
7
8use http::header::{HeaderName, HeaderValue};
9use rmcp::ClientHandler;
10use rmcp::ServiceExt;
11use rmcp::model::{
12    CallToolRequestParams, CallToolResult, ClientCapabilities, ClientInfo, ClientRequest, Implementation,
13    InitializeRequestParams, ProtocolVersion, ServerResult, Tool,
14};
15use rmcp::service::{ClientInitializeError, PeerRequestOptions, RoleClient, RunningService, ServiceError};
16use rmcp::transport::StreamableHttpClientTransport;
17use rmcp::transport::child_process::TokioChildProcess;
18use rmcp::transport::streamable_http_client::StreamableHttpClientTransportConfig;
19use rmcp_reqwest as http_client;
20use serde_json::Value;
21use tokio::io::{AsyncBufReadExt, BufReader};
22use tokio::process::Command;
23
24const CONNECTION_TIMEOUT: Duration = Duration::from_secs(30);
25const TOOL_TIMEOUT: Duration = Duration::from_secs(60);
26
27#[derive(Debug, Clone, Copy, PartialEq, Eq)]
28pub enum McpOperation {
29    Connect,
30    ListTools,
31    CallTool,
32}
33
34impl fmt::Display for McpOperation {
35    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
36        match self {
37            Self::Connect => f.write_str("connect"),
38            Self::ListTools => f.write_str("tools/list"),
39            Self::CallTool => f.write_str("tools/call"),
40        }
41    }
42}
43
44#[derive(Debug, thiserror::Error)]
45pub enum McpError {
46    #[error("failed to spawn MCP stdio server")]
47    SpawnStdio(#[source] std::io::Error),
48
49    #[error("failed to connect to MCP server")]
50    Connect(#[source] Box<ClientInitializeError>),
51
52    #[error("failed to resolve MCP server host")]
53    ResolveHost(#[source] std::io::Error),
54
55    #[error("MCP server URL has no resolvable host")]
56    UnresolvableHost,
57
58    #[error("failed to build MCP HTTP client")]
59    BuildHttpClient(#[source] http_client::Error),
60
61    #[error("invalid MCP HTTP header name")]
62    InvalidHeaderName(#[source] http::header::InvalidHeaderName),
63
64    #[error("invalid MCP HTTP header value")]
65    InvalidHeaderValue(#[source] http::header::InvalidHeaderValue),
66
67    #[error("MCP operation failed during {operation}")]
68    Operation {
69        operation: McpOperation,
70        #[source]
71        source: ServiceError,
72    },
73
74    #[error("MCP operation timed out during {operation}")]
75    Timeout { operation: McpOperation },
76
77    #[error("MCP tool arguments must be a JSON object")]
78    InvalidArguments,
79
80    #[error("MCP server returned an unexpected response during {operation}")]
81    UnexpectedResponse { operation: McpOperation },
82}
83
84#[derive(Clone)]
85struct AgenticMcpClientHandler;
86
87impl ClientHandler for AgenticMcpClientHandler {
88    fn get_info(&self) -> ClientInfo {
89        InitializeRequestParams::new(
90            ClientCapabilities::default(),
91            Implementation::new("agentic-api", env!("CARGO_PKG_VERSION")),
92        )
93        .with_protocol_version(ProtocolVersion::V_2025_06_18)
94    }
95}
96
97pub struct McpClient {
98    inner: Arc<RunningService<RoleClient, AgenticMcpClientHandler>>,
99    tool_timeout: Duration,
100}
101
102impl McpClient {
103    /// Connects to an MCP server over streamable HTTP.
104    ///
105    /// # Errors
106    ///
107    /// Returns an error if URL resolution, HTTP client or header construction,
108    /// the initialization timeout, or the MCP handshake fails.
109    pub async fn connect(server_url: &str, headers: Option<HashMap<String, String>>) -> Result<Self, McpError> {
110        tokio::time::timeout(CONNECTION_TIMEOUT, Self::connect_streamable_http(server_url, headers))
111            .await
112            .map_err(|_| McpError::Timeout {
113                operation: McpOperation::Connect,
114            })?
115    }
116
117    async fn connect_streamable_http(
118        server_url: &str,
119        headers: Option<HashMap<String, String>>,
120    ) -> Result<Self, McpError> {
121        let http_client = pinned_http_client(server_url).await?;
122        let mut config = StreamableHttpClientTransportConfig::with_uri(server_url.to_owned());
123        if let Some(headers) = headers.filter(|headers| !headers.is_empty()) {
124            let mut custom_headers = HashMap::with_capacity(headers.len());
125            for (name, value) in headers {
126                custom_headers.insert(
127                    HeaderName::try_from(name).map_err(McpError::InvalidHeaderName)?,
128                    HeaderValue::try_from(value).map_err(McpError::InvalidHeaderValue)?,
129                );
130            }
131            config = config.custom_headers(custom_headers);
132        }
133        let transport = StreamableHttpClientTransport::with_client(http_client, config);
134        let service = AgenticMcpClientHandler
135            .serve(transport)
136            .await
137            .map_err(|error| McpError::Connect(Box::new(error)))?;
138
139        Ok(Self {
140            inner: Arc::new(service),
141            tool_timeout: TOOL_TIMEOUT,
142        })
143    }
144
145    /// Spawns a local stdio MCP server and connects over stdin/stdout.
146    ///
147    /// # Errors
148    ///
149    /// Returns [`McpError::SpawnStdio`] if the process cannot be spawned.
150    /// Returns [`McpError::Connect`] if the MCP initialization handshake fails.
151    pub async fn connect_stdio(
152        command: &str,
153        args: &[String],
154        env: Option<&HashMap<String, String>>,
155        cwd: Option<&str>,
156    ) -> Result<Self, McpError> {
157        let mut command_builder = Command::new(command);
158        command_builder
159            .kill_on_drop(true)
160            .stdin(Stdio::piped())
161            .stdout(Stdio::piped())
162            .stderr(Stdio::piped())
163            .args(args);
164
165        if let Some(env) = env {
166            command_builder.envs(env);
167        }
168
169        if let Some(cwd) = cwd {
170            command_builder.current_dir(cwd);
171        }
172
173        let (transport, stderr) = TokioChildProcess::builder(command_builder)
174            .spawn()
175            .map_err(McpError::SpawnStdio)?;
176
177        if let Some(stderr) = stderr {
178            let command = command.to_owned();
179            tokio::spawn(async move {
180                let mut reader = BufReader::new(stderr).lines();
181                loop {
182                    match reader.next_line().await {
183                        Ok(Some(line)) => {
184                            tracing::info!(mcp.command = %command, %line, "MCP server stderr");
185                        }
186                        Ok(None) => break,
187                        Err(error) => {
188                            tracing::warn!(
189                                mcp.command = %command,
190                                error = %error,
191                                "failed to read MCP server stderr"
192                            );
193                            break;
194                        }
195                    }
196                }
197            });
198        }
199
200        let service = tokio::time::timeout(CONNECTION_TIMEOUT, AgenticMcpClientHandler.serve(transport))
201            .await
202            .map_err(|_| McpError::Timeout {
203                operation: McpOperation::Connect,
204            })?
205            .map_err(|error| McpError::Connect(Box::new(error)))?;
206
207        Ok(Self {
208            inner: Arc::new(service),
209            tool_timeout: TOOL_TIMEOUT,
210        })
211    }
212
213    /// Lists tools exposed by the connected MCP server.
214    ///
215    /// # Errors
216    ///
217    /// Returns [`McpError::Timeout`] if `tools/list` exceeds the configured timeout.
218    /// Returns [`McpError::Operation`] if the server rejects or fails the request.
219    pub async fn list_tools(&self) -> Result<Vec<Tool>, McpError> {
220        let result = tokio::time::timeout(self.tool_timeout, self.inner.list_tools(None))
221            .await
222            .map_err(|_| McpError::Timeout {
223                operation: McpOperation::ListTools,
224            })?
225            .map_err(|source| McpError::Operation {
226                operation: McpOperation::ListTools,
227                source,
228            })?;
229
230        Ok(result.tools)
231    }
232
233    /// Calls a tool exposed by the connected MCP server.
234    ///
235    /// # Errors
236    ///
237    /// Returns [`McpError::InvalidArguments`] if `arguments` is not a JSON object.
238    /// Returns [`McpError::Timeout`] if `tools/call` exceeds the configured timeout.
239    /// Returns [`McpError::Operation`] if the server rejects or fails the request.
240    /// Returns [`McpError::UnexpectedResponse`] if the server returns another response kind.
241    pub async fn call_tool(&self, name: &str, arguments: Option<Value>) -> Result<CallToolResult, McpError> {
242        let arguments = match arguments {
243            Some(Value::Object(map)) => Some(map),
244            Some(_) => return Err(McpError::InvalidArguments),
245            None => None,
246        };
247
248        let mut params = CallToolRequestParams::new(name.to_owned());
249        params.arguments = arguments;
250
251        let result = tokio::time::timeout(self.tool_timeout, async {
252            self.inner
253                .peer()
254                .send_request_with_option(
255                    ClientRequest::CallToolRequest(rmcp::model::CallToolRequest::new(params)),
256                    PeerRequestOptions::no_options(),
257                )
258                .await?
259                .await_response()
260                .await
261        })
262        .await
263        .map_err(|_| McpError::Timeout {
264            operation: McpOperation::CallTool,
265        })?
266        .map_err(|source| McpError::Operation {
267            operation: McpOperation::CallTool,
268            source,
269        })?;
270
271        match result {
272            ServerResult::CallToolResult(result) => Ok(result),
273            _ => Err(McpError::UnexpectedResponse {
274                operation: McpOperation::CallTool,
275            }),
276        }
277    }
278}
279
280/// Build the HTTP client used by an MCP connection with DNS pinned to the
281/// addresses resolved during connection setup. This prevents a hostname from
282/// resolving to different addresses between URL validation and later requests.
283async fn pinned_http_client(server_url: &str) -> Result<http_client::Client, McpError> {
284    let url = http_client::Url::parse(server_url).map_err(|_| McpError::UnresolvableHost)?;
285    let port = url.port_or_known_default().ok_or(McpError::UnresolvableHost)?;
286    match url.host().ok_or(McpError::UnresolvableHost)? {
287        url::Host::Domain(host) => {
288            let addresses = tokio::net::lookup_host((host, port))
289                .await
290                .map_err(McpError::ResolveHost)?
291                .collect::<Vec<_>>();
292            if addresses.is_empty() {
293                return Err(McpError::UnresolvableHost);
294            }
295            http_client_for_addresses(host, &addresses)
296        }
297        url::Host::Ipv4(_) | url::Host::Ipv6(_) => http_client_for_literal_address(),
298    }
299}
300
301fn http_client_for_addresses(host: &str, addresses: &[SocketAddr]) -> Result<http_client::Client, McpError> {
302    http_client::Client::builder()
303        .no_proxy()
304        .redirect(http_client::redirect::Policy::none())
305        .resolve_to_addrs(host, addresses)
306        .build()
307        .map_err(McpError::BuildHttpClient)
308}
309
310fn http_client_for_literal_address() -> Result<http_client::Client, McpError> {
311    http_client::Client::builder()
312        .no_proxy()
313        .redirect(http_client::redirect::Policy::none())
314        .build()
315        .map_err(McpError::BuildHttpClient)
316}
317
318#[cfg(test)]
319mod tests {
320    use super::{McpError, http_client_for_addresses, pinned_http_client};
321    use std::io::{Read, Write};
322    use std::net::{Ipv4Addr, SocketAddr, TcpListener as StdTcpListener};
323    use std::process::Command;
324    use std::sync::Arc;
325    use std::sync::atomic::{AtomicBool, Ordering};
326    use std::thread;
327    use std::time::Duration;
328    use tokio::io::{AsyncReadExt, AsyncWriteExt};
329    use tokio::net::TcpListener as TokioTcpListener;
330
331    const PROXY_CHILD_ENV: &str = "AGENTIC_MCP_PROXY_TEST_CHILD";
332    const PROXY_TARGET_ENV: &str = "AGENTIC_MCP_PROXY_TEST_TARGET";
333
334    async fn spawn_http_server(response: &'static str) -> SocketAddr {
335        let listener = TokioTcpListener::bind((Ipv4Addr::LOCALHOST, 0)).await.unwrap();
336        let address = listener.local_addr().unwrap();
337        tokio::spawn(async move {
338            let (mut stream, _) = listener.accept().await.unwrap();
339            let mut request = [0_u8; 1024];
340            let _ = stream.read(&mut request).await.unwrap();
341            stream.write_all(response.as_bytes()).await.unwrap();
342        });
343        address
344    }
345
346    #[tokio::test]
347    async fn pinned_client_tries_every_resolved_address() {
348        let reachable = spawn_http_server("HTTP/1.1 200 OK\r\nContent-Length: 0\r\nConnection: close\r\n\r\n").await;
349        let unreachable = SocketAddr::from(([0, 0, 0, 0, 0, 0, 0, 1], reachable.port()));
350        let client = http_client_for_addresses("mcp.test", &[unreachable, reachable]).unwrap();
351
352        let response = tokio::time::timeout(
353            std::time::Duration::from_secs(5),
354            client.get(format!("http://mcp.test:{}/", reachable.port())).send(),
355        )
356        .await
357        .expect("client must try the reachable address")
358        .unwrap();
359
360        assert_eq!(response.status(), http::StatusCode::OK);
361    }
362
363    #[tokio::test]
364    async fn pinned_client_does_not_follow_redirects() {
365        let address = spawn_http_server(
366            "HTTP/1.1 302 Found\r\nLocation: http://127.0.0.1:9/\r\nContent-Length: 0\r\nConnection: close\r\n\r\n",
367        )
368        .await;
369        let client = http_client_for_addresses("mcp.test", &[address]).unwrap();
370
371        let response = client
372            .get(format!("http://mcp.test:{}/", address.port()))
373            .send()
374            .await
375            .unwrap();
376
377        assert!(response.status().is_redirection());
378    }
379
380    #[tokio::test]
381    async fn pinned_client_rejects_malformed_urls() {
382        assert!(matches!(
383            pinned_http_client("not a URL").await,
384            Err(McpError::UnresolvableHost)
385        ));
386    }
387
388    #[tokio::test]
389    async fn pinned_client_accepts_ipv6_literal_urls_without_dns() {
390        pinned_http_client("http://[::1]:8000/mcp").await.unwrap();
391    }
392
393    fn serve_once(listener: &StdTcpListener, status: &str, accepted: &AtomicBool) {
394        listener.set_nonblocking(true).unwrap();
395        for _ in 0..200 {
396            match listener.accept() {
397                Ok((mut stream, _)) => {
398                    accepted.store(true, Ordering::SeqCst);
399                    let mut request = [0_u8; 1024];
400                    let _ = stream.read(&mut request).unwrap();
401                    write!(
402                        stream,
403                        "HTTP/1.1 {status}\r\nContent-Length: 0\r\nConnection: close\r\n\r\n"
404                    )
405                    .unwrap();
406                    return;
407                }
408                Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => {
409                    thread::sleep(Duration::from_millis(5));
410                }
411                Err(error) => panic!("test server failed: {error}"),
412            }
413        }
414    }
415
416    #[test]
417    fn pinned_client_ignores_system_proxy() {
418        if std::env::var_os(PROXY_CHILD_ENV).is_some() {
419            let target = std::env::var(PROXY_TARGET_ENV).unwrap().parse::<SocketAddr>().unwrap();
420            let runtime = tokio::runtime::Runtime::new().unwrap();
421            runtime.block_on(async move {
422                let client = http_client_for_addresses("mcp.test", &[target]).unwrap();
423                let response = client
424                    .get(format!("http://mcp.test:{}/", target.port()))
425                    .send()
426                    .await
427                    .unwrap();
428                assert_eq!(response.status(), http::StatusCode::OK);
429            });
430            return;
431        }
432
433        let target_listener = StdTcpListener::bind((Ipv4Addr::LOCALHOST, 0)).unwrap();
434        let target_address = target_listener.local_addr().unwrap();
435        let target_accepted = Arc::new(AtomicBool::new(false));
436        let target_thread = {
437            let accepted = Arc::clone(&target_accepted);
438            thread::spawn(move || serve_once(&target_listener, "200 OK", &accepted))
439        };
440
441        let proxy_listener = StdTcpListener::bind((Ipv4Addr::LOCALHOST, 0)).unwrap();
442        let proxy_address = proxy_listener.local_addr().unwrap();
443        let proxy_accepted = Arc::new(AtomicBool::new(false));
444        let proxy_thread = {
445            let accepted = Arc::clone(&proxy_accepted);
446            thread::spawn(move || serve_once(&proxy_listener, "418 I'm a teapot", &accepted))
447        };
448
449        let output = Command::new(std::env::current_exe().unwrap())
450            .arg("--exact")
451            .arg("tool::mcp::client::tests::pinned_client_ignores_system_proxy")
452            .arg("--nocapture")
453            .env(PROXY_CHILD_ENV, "1")
454            .env(PROXY_TARGET_ENV, target_address.to_string())
455            .env("HTTP_PROXY", format!("http://{proxy_address}"))
456            .env("http_proxy", format!("http://{proxy_address}"))
457            .env("ALL_PROXY", format!("http://{proxy_address}"))
458            .env("all_proxy", format!("http://{proxy_address}"))
459            .env_remove("NO_PROXY")
460            .env_remove("no_proxy")
461            .output()
462            .unwrap();
463
464        target_thread.join().unwrap();
465        proxy_thread.join().unwrap();
466        assert!(
467            output.status.success(),
468            "child test failed: {}",
469            String::from_utf8_lossy(&output.stderr)
470        );
471        assert!(target_accepted.load(Ordering::SeqCst));
472        assert!(!proxy_accepted.load(Ordering::SeqCst));
473    }
474}