harn_vm/mcp/
mod.rs

1//! MCP (Model Context Protocol) client for connecting to external tool servers.
2//!
3//! Supports stdio transport and streamable HTTP-style request/response transport.
4
5pub(crate) use std::collections::BTreeMap;
6pub(crate) use std::future::Future;
7pub(crate) use std::path::{Path, PathBuf};
8pub(crate) use std::sync::Arc;
9
10pub(crate) use base64::Engine;
11pub(crate) use futures::StreamExt;
12pub(crate) use reqwest_eventsource::{Event as SseEvent, EventSource};
13pub(crate) use serde::Deserialize;
14pub(crate) use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
15pub(crate) use tokio::process::{Child, ChildStdin, ChildStdout};
16pub(crate) use tokio::sync::Mutex;
17
18pub(crate) use crate::stdlib::json_to_vm_value;
19pub(crate) use crate::value::{VmError, VmValue};
20pub(crate) use crate::vm::Vm;
21
22pub(crate) use crate::mcp_protocol::{
23    cache_hints_to_json, rc_name_header_value, McpCacheHint, McpProtocolMode,
24    DRAFT_PROTOCOL_VERSION, LEGACY_2025_06_18_PROTOCOL_VERSION, MCP_SESSION_HEADER_LEGACY,
25    PROTOCOL_VERSION, RC_HEADER_METHOD, RC_HEADER_NAME, RC_HEADER_PROTOCOL_VERSION,
26    RC_META_KEY_CLIENT_CAPABILITIES, RC_META_KEY_CLIENT_INFO, RC_META_KEY_PROTOCOL_VERSION,
27    RESULT_TYPE_INPUT_REQUIRED, UNSUPPORTED_PROTOCOL_VERSION_CODE,
28};
29
30mod builtins;
31mod connect;
32mod notifications;
33mod protocol;
34mod roots;
35mod transport;
36
37pub use builtins::*;
38pub use connect::*;
39pub(crate) use notifications::*;
40pub(crate) use protocol::*;
41pub(crate) use roots::*;
42pub(crate) use transport::*;
43
44const X_MCP_HEADER: &str = "x-mcp-header";
45
46const MCP_INPUT_REQUIRED_MAX_ROUNDS: usize = 8;
47
48/// Default timeout for MCP requests (60 seconds).
49const MCP_TIMEOUT: std::time::Duration = std::time::Duration::from_mins(1);
50
51#[derive(Clone, Debug, Deserialize)]
52#[serde(rename_all = "lowercase")]
53pub(crate) enum McpTransport {
54    Stdio,
55    Http,
56}
57
58#[derive(Clone, Debug, Deserialize)]
59pub struct McpServerSpec {
60    pub name: String,
61    #[serde(default = "default_transport")]
62    transport: McpTransport,
63    #[serde(default)]
64    pub command: String,
65    #[serde(default)]
66    pub args: Vec<String>,
67    #[serde(default)]
68    pub env: BTreeMap<String, String>,
69    #[serde(default)]
70    pub url: String,
71    #[serde(default)]
72    pub auth_token: Option<String>,
73    #[serde(default)]
74    pub token_exchange: Option<crate::mcp_oauth::McpTokenExchangeConfig>,
75    #[serde(default)]
76    pub protocol_version: Option<String>,
77    #[serde(default)]
78    pub protocol_mode: Option<String>,
79    #[serde(default)]
80    pub proxy_server_name: Option<String>,
81}
82
83fn default_transport() -> McpTransport {
84    McpTransport::Stdio
85}
86
87/// Internal state for an MCP client connection.
88pub(crate) enum McpClientInner {
89    Stdio(StdioMcpClientInner),
90    Http(HttpMcpClientInner),
91}
92
93pub(crate) struct StdioMcpClientInner {
94    child: Child,
95    stdin: ChildStdin,
96    reader: BufReader<ChildStdout>,
97    next_id: u64,
98    protocol_mode: McpProtocolMode,
99    protocol_version: String,
100}
101
102pub(crate) struct HttpMcpClientInner {
103    client: reqwest::Client,
104    url: String,
105    auth_token: Option<String>,
106    auth_token_source: HttpAuthTokenSource,
107    token_exchange: Option<Arc<crate::mcp_oauth::McpTokenExchangeConfig>>,
108    protocol_mode: McpProtocolMode,
109    protocol_version: String,
110    session_id: Option<String>,
111    next_id: u64,
112    proxy_server_name: Option<String>,
113    get_stream_task: Option<tokio::task::JoinHandle<()>>,
114    tool_headers: BTreeMap<String, Vec<McpToolHeader>>,
115}
116
117#[derive(Clone, Copy, Debug, PartialEq, Eq)]
118pub(crate) enum HttpAuthTokenSource {
119    None,
120    Config,
121    OAuthStore,
122}
123
124#[derive(Clone, Debug, PartialEq, Eq)]
125pub(crate) struct ResolvedHttpAuthToken {
126    pub token: Option<String>,
127    pub source: HttpAuthTokenSource,
128}
129
130#[derive(Clone, Debug, PartialEq, Eq)]
131pub(crate) struct McpToolHeader {
132    parameter: String,
133    header_name: String,
134}
135
136#[derive(Clone, Debug, PartialEq, Eq)]
137pub(crate) struct McpRoot {
138    path: String,
139    uri: String,
140    name: String,
141}
142
143impl McpRoot {
144    fn protocol_json(&self) -> serde_json::Value {
145        serde_json::json!({
146            "uri": self.uri,
147            "name": self.name,
148        })
149    }
150
151    fn script_json(&self) -> serde_json::Value {
152        serde_json::json!({
153            "uri": self.uri,
154            "name": self.name,
155            "path": self.path,
156        })
157    }
158}
159
160impl HttpMcpClientInner {
161    fn abort_get_stream(&mut self) {
162        if let Some(task) = self.get_stream_task.take() {
163            task.abort();
164        }
165    }
166}
167
168impl Drop for StdioMcpClientInner {
169    fn drop(&mut self) {
170        let _ = self.child.start_kill();
171    }
172}
173
174impl Drop for HttpMcpClientInner {
175    fn drop(&mut self) {
176        self.abort_get_stream();
177    }
178}
179
180/// Handle to an MCP client connection, stored in VmValue.
181#[derive(Clone)]
182pub struct VmMcpClientHandle {
183    pub name: String,
184    inner: Arc<Mutex<Option<McpClientInner>>>,
185    last_roots: Arc<Mutex<Vec<McpRoot>>>,
186    pub(crate) initialize_result: Arc<Mutex<Option<serde_json::Value>>>,
187    cache_hints: Arc<Mutex<BTreeMap<String, McpCacheHint>>>,
188}
189
190impl std::fmt::Debug for VmMcpClientHandle {
191    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
192        write!(f, "McpClient({})", self.name)
193    }
194}
195
196impl VmMcpClientHandle {
197    async fn protocol_mode(&self) -> Result<McpProtocolMode, VmError> {
198        let guard = self.inner.lock().await;
199        let inner = guard
200            .as_ref()
201            .ok_or_else(|| VmError::Runtime("MCP client is disconnected".into()))?;
202        Ok(match inner {
203            McpClientInner::Stdio(inner) => inner.protocol_mode,
204            McpClientInner::Http(inner) => inner.protocol_mode,
205        })
206    }
207
208    async fn protocol_version(&self) -> Result<String, VmError> {
209        let guard = self.inner.lock().await;
210        let inner = guard
211            .as_ref()
212            .ok_or_else(|| VmError::Runtime("MCP client is disconnected".into()))?;
213        Ok(match inner {
214            McpClientInner::Stdio(inner) => inner.protocol_version.clone(),
215            McpClientInner::Http(inner) => inner.protocol_version.clone(),
216        })
217    }
218
219    async fn switch_to_legacy_protocol(&self) -> Result<(), VmError> {
220        let mut guard = self.inner.lock().await;
221        let inner = guard
222            .as_mut()
223            .ok_or_else(|| VmError::Runtime("MCP client is disconnected".into()))?;
224        match inner {
225            McpClientInner::Stdio(inner) => {
226                inner.protocol_mode = McpProtocolMode::Legacy;
227                inner.protocol_version = PROTOCOL_VERSION.to_string();
228            }
229            McpClientInner::Http(inner) => {
230                inner.protocol_mode = McpProtocolMode::Legacy;
231                inner.protocol_version = PROTOCOL_VERSION.to_string();
232            }
233        }
234        Ok(())
235    }
236
237    pub(crate) async fn call(
238        &self,
239        method: &str,
240        params: serde_json::Value,
241    ) -> Result<serde_json::Value, VmError> {
242        let msg = self.call_raw(method, params).await?;
243        parse_jsonrpc_result(msg)
244    }
245
246    async fn call_raw(
247        &self,
248        method: &str,
249        params: serde_json::Value,
250    ) -> Result<serde_json::Value, VmError> {
251        if method != "initialize" && method != "server/discover" {
252            self.notify_roots_list_changed_if_needed().await?;
253        }
254        let record_request = serde_json::json!({
255            "jsonrpc": "2.0",
256            "method": method,
257            "params": params.clone(),
258        });
259        let started_at = crate::clock_mock::instant_now();
260        let mut guard = self.inner.lock().await;
261        let inner = guard
262            .as_mut()
263            .ok_or_else(|| VmError::Runtime("MCP client is disconnected".into()))?;
264
265        let result = match inner {
266            McpClientInner::Stdio(inner) => stdio_call_raw(inner, &self.name, method, params).await,
267            McpClientInner::Http(inner) => http_call_raw(inner, &self.name, method, params).await,
268        };
269        let latency_ms = crate::clock_mock::instant_now()
270            .duration_since(started_at)
271            .as_millis()
272            .min(u64::MAX as u128) as u64;
273        let record_response = match &result {
274            Ok(response) => response.clone(),
275            Err(error) => serde_json::json!({
276                "jsonrpc": "2.0",
277                "error": {
278                    "message": error.to_string(),
279                }
280            }),
281        };
282        crate::testbench::tape::record_mcp_json_rpc(
283            &self.name,
284            method,
285            &record_request,
286            &record_response,
287            latency_ms,
288        );
289        result
290    }
291
292    async fn notify(&self, method: &str, params: serde_json::Value) -> Result<(), VmError> {
293        let mut guard = self.inner.lock().await;
294        let inner = guard
295            .as_mut()
296            .ok_or_else(|| VmError::Runtime("MCP client is disconnected".into()))?;
297
298        match inner {
299            McpClientInner::Stdio(inner) => stdio_notify(inner, method, params).await,
300            McpClientInner::Http(inner) => http_notify(inner, &self.name, method, params).await,
301        }
302    }
303
304    pub(crate) async fn disconnect(&self) -> Result<(), VmError> {
305        let mut guard = self.inner.lock().await;
306        if let Some(inner) = guard.take() {
307            match inner {
308                McpClientInner::Stdio(mut inner) => {
309                    let _ = inner.child.kill().await;
310                }
311                McpClientInner::Http(mut inner) => {
312                    inner.abort_get_stream();
313                }
314            }
315        }
316        Ok(())
317    }
318
319    async fn notify_roots_list_changed_if_needed(&self) -> Result<(), VmError> {
320        if self.protocol_mode().await? == McpProtocolMode::Modern {
321            return Ok(());
322        }
323        let roots = current_mcp_roots();
324        let mut last_roots = self.last_roots.lock().await;
325        if *last_roots == roots {
326            return Ok(());
327        }
328
329        self.notify(
330            crate::mcp_protocol::METHOD_ROOTS_LIST_CHANGED_NOTIFICATION,
331            serde_json::json!({}),
332        )
333        .await?;
334        *last_roots = roots;
335        Ok(())
336    }
337
338    async fn record_cache_hint(&self, method: &str, result: &serde_json::Value) {
339        let Some(hint) = McpCacheHint::from_result(result) else {
340            return;
341        };
342        self.cache_hints
343            .lock()
344            .await
345            .insert(method.to_string(), hint);
346    }
347
348    async fn store_http_tool_headers(&self, tools: &[serde_json::Value]) {
349        let mut valid_headers = BTreeMap::new();
350        let mut valid_tools = std::collections::BTreeSet::new();
351        for tool in tools {
352            let Some(name) = tool.get("name").and_then(|value| value.as_str()) else {
353                continue;
354            };
355            match extract_tool_headers(tool) {
356                Ok(headers) => {
357                    valid_tools.insert(name.to_string());
358                    if !headers.is_empty() {
359                        valid_headers.insert(name.to_string(), headers);
360                    }
361                }
362                Err(reason) => {
363                    tracing::warn!(tool = name, %reason, "rejecting MCP tool with invalid x-mcp-header annotation");
364                }
365            }
366        }
367
368        let mut guard = self.inner.lock().await;
369        if let Some(McpClientInner::Http(inner)) = guard.as_mut() {
370            inner
371                .tool_headers
372                .retain(|tool, _| valid_tools.contains(tool));
373            inner.tool_headers.extend(valid_headers);
374        }
375    }
376}
377
378#[derive(Clone)]
379pub(crate) struct HttpStreamConfig {
380    client: reqwest::Client,
381    url: String,
382    auth_token: Option<String>,
383    protocol_mode: McpProtocolMode,
384    protocol_version: String,
385    session_id: Option<String>,
386    proxy_server_name: Option<String>,
387    server_name: String,
388}
389
390pub(crate) struct McpConnectOptions {
391    protocol_mode: McpProtocolMode,
392    protocol_version: String,
393}
394
395/// Process-wide registry mapping MCP progress tokens to the agent
396/// session that issued the originating `tools/call`. The MCP transport
397/// reader (line-by-line stdio loop, SSE parser, multiplexed HTTP
398/// stream) consults this registry when it sees an inbound
399/// `notifications/progress` so the right session's `agent_inbox`
400/// receives the update — even when the reader runs on a task that
401/// doesn't share the issuer's thread-local current-session.
402pub(crate) mod client_progress {
403    use std::collections::HashMap;
404    use std::sync::{Mutex, OnceLock, PoisonError};
405    use uuid::Uuid;
406
407    /// Snapshot of an active progress-token registration.
408    #[derive(Clone, Debug)]
409    pub struct ProgressTokenContext {
410        pub token: String,
411        pub session_id: Option<String>,
412        #[allow(dead_code)] // recorded for telemetry / future debug surface
413        pub server: String,
414        pub tool: String,
415    }
416
417    fn registry() -> &'static Mutex<HashMap<String, ProgressTokenContext>> {
418        static REGISTRY: OnceLock<Mutex<HashMap<String, ProgressTokenContext>>> = OnceLock::new();
419        REGISTRY.get_or_init(|| Mutex::new(HashMap::new()))
420    }
421
422    fn lock<T>(m: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
423        m.lock().unwrap_or_else(PoisonError::into_inner)
424    }
425
426    /// Issue a fresh progress token for an outgoing `tools/call`. The
427    /// token is registered against the active session id when one is
428    /// available; calls made outside a session context still get a
429    /// token (so the server emits notifications), but inbound updates
430    /// for it have no inbox to land in and are dropped.
431    pub fn issue_token(server: &str, tool: &str) -> Option<ProgressTokenContext> {
432        let session_id = crate::llm::current_agent_session_id();
433        let ctx = ProgressTokenContext {
434            token: format!("hpt_{}", Uuid::now_v7()),
435            session_id,
436            server: server.to_string(),
437            tool: tool.to_string(),
438        };
439        lock(registry()).insert(ctx.token.clone(), ctx.clone());
440        Some(ctx)
441    }
442
443    /// Look up the registration for `token`. Returns `None` for unknown
444    /// or already-released tokens (e.g. a late progress notification
445    /// that arrived after the call completed).
446    pub fn lookup(token: &str) -> Option<ProgressTokenContext> {
447        lock(registry()).get(token).cloned()
448    }
449
450    /// Drop the registration for `token`. Idempotent.
451    pub fn release(token: &str) {
452        lock(registry()).remove(token);
453    }
454
455    /// RAII guard returned by [`issue_token`]; releases the
456    /// registration on drop so a long-lived MCP connection doesn't
457    /// accumulate stale tokens.
458    pub struct ProgressTokenGuard {
459        pub token: String,
460    }
461
462    impl Drop for ProgressTokenGuard {
463        fn drop(&mut self) {
464            release(&self.token);
465        }
466    }
467
468    #[cfg(any(test, feature = "vm-bench-internals"))]
469    #[allow(dead_code)]
470    pub fn reset() {
471        lock(registry()).clear();
472    }
473}
474
475#[derive(Clone, Debug)]
476pub(crate) struct McpInputRound {
477    input_responses: serde_json::Value,
478    request_state: Option<serde_json::Value>,
479}
480
481#[cfg(test)]
482mod tests;
harn_vm/mcp/mod.rs

harn_vm/mcp/
mod.rs
