context_harness/agent_runtime/
mcp_client.rs

1//! Bounded stdio MCP sessions. Authorization must precede `connect`: spawning
2//! a server is process execution even when its advertised tools claim read-only.
3use crate::{
4    agent_resource::Capability,
5    config::McpServerConfig,
6    tool_binding::{self, LoadedToolResource, ResolvedToolBinding},
7    traits::{Tool, ToolContext, ToolRegistry},
8};
9use anyhow::{anyhow, ensure, Result};
10use async_trait::async_trait;
11use futures_util::{SinkExt, StreamExt};
12use rmcp::{
13    model::{
14        CallToolRequestParams, ClientInfo, ClientResult, PaginatedRequestParams,
15        ServerNotification, ServerRequest, TaskSupport,
16    },
17    service::{NotificationContext, RequestContext, RunningService},
18    transport::Transport,
19    ErrorData, Peer, RoleClient, Service, ServiceExt,
20};
21use serde_json::Value;
22use sha2::{Digest, Sha256};
23use std::{collections::HashSet, path::Path, process::Stdio, sync::Arc, time::Duration};
24use tokio::{
25    process::{Child, ChildStdin, ChildStdout, Command},
26    sync::Mutex,
27};
28use tokio_util::codec::{FramedRead, FramedWrite, LinesCodec};
29
30const MAX_FRAME: usize = 1024 * 1024;
31const MAX_ARGUMENTS: usize = 64 * 1024;
32type Incoming = rmcp::service::RxJsonRpcMessage<RoleClient>;
33type Outgoing = rmcp::service::TxJsonRpcMessage<RoleClient>;
34struct BoundedTransport {
35    reader: FramedRead<ChildStdout, LinesCodec>,
36    writer: Arc<Mutex<Option<FramedWrite<ChildStdin, LinesCodec>>>>,
37    ended: bool,
38}
39impl Transport<RoleClient> for BoundedTransport {
40    type Error = std::io::Error;
41    fn send(
42        &mut self,
43        message: Outgoing,
44    ) -> impl std::future::Future<Output = std::io::Result<()>> + Send + 'static {
45        let writer = self.writer.clone();
46        async move {
47            let bytes = serde_json::to_vec(&message)
48                .map_err(|_| std::io::Error::other("MCP request encoding failed"))?;
49            if bytes.len() > MAX_FRAME {
50                return Err(std::io::Error::other("MCP request exceeds size limit"));
51            }
52            let mut writer = writer.lock().await;
53            writer
54                .as_mut()
55                .ok_or_else(|| std::io::Error::other("MCP transport closed"))?
56                .send(
57                    String::from_utf8(bytes)
58                        .map_err(|_| std::io::Error::other("MCP request encoding failed"))?,
59                )
60                .await
61                .map_err(|_| std::io::Error::other("MCP write failed"))
62        }
63    }
64    async fn receive(&mut self) -> Option<Incoming> {
65        if self.ended {
66            return None;
67        }
68        match self.reader.next().await {
69            Some(Ok(line)) => match serde_json::from_str(&line) {
70                Ok(message) => Some(message),
71                Err(_) => {
72                    self.ended = true;
73                    None
74                }
75            },
76            _ => {
77                self.ended = true;
78                None
79            }
80        }
81    }
82    async fn close(&mut self) -> std::io::Result<()> {
83        self.ended = true;
84        self.writer.lock().await.take();
85        Ok(())
86    }
87}
88/// No sampling, roots, elicitation, or server-initiated execution is supported.
89struct EmptyClient;
90impl Service<RoleClient> for EmptyClient {
91    async fn handle_request(
92        &self,
93        _: ServerRequest,
94        _: RequestContext<RoleClient>,
95    ) -> Result<ClientResult, ErrorData> {
96        Err(ErrorData::invalid_request(
97            "client requests are unsupported",
98            None,
99        ))
100    }
101    async fn handle_notification(
102        &self,
103        _: ServerNotification,
104        _: NotificationContext<RoleClient>,
105    ) -> Result<(), ErrorData> {
106        Ok(())
107    }
108    fn get_info(&self) -> ClientInfo {
109        ClientInfo::default()
110    }
111}
112
113pub(super) struct Session {
114    // Child stays with this guard, never in a detached SDK service task.
115    child: Child,
116    service: RunningService<RoleClient, EmptyClient>,
117    tools: Vec<RemoteTool>,
118}
119impl Drop for Session {
120    fn drop(&mut self) {
121        self.service.cancellation_token().cancel();
122        let _ = self.child.start_kill();
123    }
124}
125impl Session {
126    pub(super) fn register(&self, registry: &mut ToolRegistry) -> Result<()> {
127        ensure!(
128            self.tools
129                .iter()
130                .all(|tool| registry.find(&tool.name).is_none()),
131            "duplicate MCP tool namespace"
132        );
133        for tool in &self.tools {
134            registry.register(Box::new(tool.clone()));
135        }
136        Ok(())
137    }
138    pub(super) async fn close(&mut self) {
139        let _ = self.child.start_kill();
140        let _ = self
141            .service
142            .close_with_timeout(Duration::from_secs(1))
143            .await;
144        let _ = tokio::time::timeout(Duration::from_secs(1), self.child.wait()).await;
145    }
146
147    pub(super) fn register_alias(
148        &self,
149        resource: &LoadedToolResource,
150        server: &McpServerConfig,
151        registry: &mut ToolRegistry,
152    ) -> Result<()> {
153        let (_, remote_name) =
154            tool_binding::mcp_reference(&resource.definition.tool.implementation)
155                .ok_or_else(|| anyhow!("invalid MCP alias implementation"))?;
156        let remote = self
157            .tools
158            .iter()
159            .find(|tool| tool.remote == remote_name)
160            .ok_or_else(|| anyhow!("MCP alias remote tool was not discovered"))?;
161        let binding = tool_binding::resolve_mcp_binding(
162            &resource.definition,
163            &remote.schema,
164            &remote.description,
165            &serde_json::to_value(server)?,
166        )?;
167        ensure!(
168            registry.find(&binding.name).is_none(),
169            "duplicate MCP alias"
170        );
171        registry.register(Box::new(RemoteAlias {
172            remote: remote.clone(),
173            binding,
174        }));
175        Ok(())
176    }
177}
178
179pub(super) async fn connect(name: &str, config: &McpServerConfig, root: &Path) -> Result<Session> {
180    config.validate(name)?;
181    let timeout = Duration::from_secs(config.timeout_seconds);
182    let mut child = Command::new(&config.command)
183        .args(&config.args)
184        .current_dir(root)
185        .stdin(Stdio::piped())
186        .stdout(Stdio::piped())
187        .stderr(Stdio::null())
188        .kill_on_drop(true)
189        .spawn()
190        .map_err(|_| anyhow!("MCP process could not start"))?;
191    let reader = child
192        .stdout
193        .take()
194        .ok_or_else(|| anyhow!("MCP stdout unavailable"))?;
195    let writer = child
196        .stdin
197        .take()
198        .ok_or_else(|| anyhow!("MCP stdin unavailable"))?;
199    let transport = BoundedTransport {
200        reader: FramedRead::new(reader, LinesCodec::new_with_max_length(MAX_FRAME)),
201        writer: Arc::new(Mutex::new(Some(FramedWrite::new(
202            writer,
203            LinesCodec::new_with_max_length(MAX_FRAME),
204        )))),
205        ended: false,
206    };
207    let service = tokio::time::timeout(timeout, EmptyClient.serve(transport))
208        .await
209        .map_err(|_| anyhow!("MCP initialization timed out"))?
210        .map_err(|_| anyhow!("MCP initialization failed"))?;
211    let mut session = Session {
212        child,
213        service,
214        tools: vec![],
215    };
216    let mut cursor = None;
217    let mut cursors = HashSet::new();
218    let mut names = HashSet::new();
219    for page in 0..16 {
220        let result = tokio::time::timeout(
221            timeout,
222            session
223                .service
224                .list_tools(Some(PaginatedRequestParams { meta: None, cursor })),
225        )
226        .await
227        .map_err(|_| anyhow!("MCP discovery timed out"))?
228        .map_err(|_| anyhow!("MCP discovery failed"))?;
229        ensure!(
230            session.tools.len() + result.tools.len() <= 128,
231            "MCP discovery exceeds tool limit"
232        );
233        for tool in result.tools {
234            ensure!(
235                !tool.name.is_empty()
236                    && tool.name.len() <= 128
237                    && tool
238                        .name
239                        .bytes()
240                        .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'_' | b'-' | b'.')),
241                "invalid MCP tool name"
242            );
243            ensure!(
244                names.insert(tool.name.to_string()),
245                "duplicate MCP tool name"
246            );
247            ensure!(
248                tool.input_schema.get("type").and_then(Value::as_str) == Some("object"),
249                "MCP tool schema must declare object type"
250            );
251            ensure!(
252                tool.execution
253                    .as_ref()
254                    .and_then(|execution| execution.task_support)
255                    != Some(TaskSupport::Required),
256                "MCP task-only tools are unsupported"
257            );
258            let public_name = format!("mcp.{name}.{}", tool.name);
259            let remote_name = tool.name.into_owned();
260            let schema = Value::Object((*tool.input_schema).clone());
261            let capabilities = vec![Capability::ProcessExecute, Capability::ExternalSideEffect];
262            let implementation_version = format!(
263                "sha256:{:x}",
264                Sha256::digest(serde_json::to_vec(&serde_json::json!({
265                    "server": config,
266                    "remote": remote_name,
267                    "schema": schema,
268                }))?)
269            );
270            let metadata = tool_binding::compatibility_binding_metadata_versioned(
271                &public_name,
272                &public_name,
273                &implementation_version,
274                schema.clone(),
275                capabilities,
276                crate::tool_binding::ToolTrustClass::McpExternal,
277            );
278            session.tools.push(RemoteTool {
279                name: public_name,
280                remote: remote_name,
281                description: tool
282                    .description
283                    .map(|description| description.into_owned())
284                    .unwrap_or_default(),
285                schema,
286                metadata,
287                peer: session.service.peer().clone(),
288                timeout,
289            });
290        }
291        cursor = result.next_cursor;
292        if let Some(next) = &cursor {
293            ensure!(
294                page < 15 && next.len() <= 4096 && cursors.insert(next.clone()),
295                "MCP discovery pagination limit"
296            );
297        } else {
298            return Ok(session);
299        }
300    }
301    unreachable!()
302}
303
304pub(super) fn validate_arguments(arguments: &Value) -> Result<()> {
305    ensure!(arguments.is_object(), "MCP arguments must be an object");
306    ensure!(
307        serde_json::to_vec(arguments)?.len() <= MAX_ARGUMENTS,
308        "MCP arguments exceed size limit"
309    );
310    Ok(())
311}
312#[derive(Clone)]
313struct RemoteTool {
314    name: String,
315    remote: String,
316    description: String,
317    schema: Value,
318    metadata: Value,
319    peer: Peer<RoleClient>,
320    timeout: Duration,
321}
322
323#[derive(Clone)]
324struct RemoteAlias {
325    remote: RemoteTool,
326    binding: ResolvedToolBinding,
327}
328
329impl RemoteAlias {
330    fn effective_arguments(&self, arguments: Value) -> Result<Value> {
331        tool_binding::validate_binding_arguments(&self.binding.public_schema, &arguments)?;
332        let mut arguments = arguments
333            .as_object()
334            .ok_or_else(|| anyhow!("MCP arguments must be an object"))?
335            .clone();
336        for (name, value) in self
337            .binding
338            .fixed
339            .as_object()
340            .ok_or_else(|| anyhow!("fixed arguments must be an object"))?
341        {
342            ensure!(
343                !arguments.contains_key(name),
344                "fixed argument cannot be overridden"
345            );
346            arguments.insert(name.clone(), value.clone());
347        }
348        let arguments = Value::Object(arguments);
349        tool_binding::validate_binding_arguments(&self.remote.schema, &arguments)?;
350        Ok(arguments)
351    }
352}
353
354#[async_trait]
355impl Tool for RemoteAlias {
356    fn name(&self) -> &str {
357        &self.binding.name
358    }
359    fn description(&self) -> &str {
360        &self.binding.description
361    }
362    fn parameters_schema(&self) -> Value {
363        self.binding.public_schema.clone()
364    }
365    fn capabilities(&self) -> Option<Vec<Capability>> {
366        Some(self.binding.capabilities.clone())
367    }
368    fn binding_metadata(&self) -> Option<Value> {
369        let mut metadata = tool_binding::binding_metadata(&self.binding);
370        metadata["remote_metadata"] = serde_json::json!(self.binding.remote_metadata);
371        Some(metadata)
372    }
373    fn validate_arguments(&self, arguments: &Value) -> Result<()> {
374        tool_binding::validate_binding_arguments(&self.binding.public_schema, arguments)
375    }
376    fn approval_arguments(&self, arguments: &Value) -> Result<Value> {
377        self.effective_arguments(arguments.clone())
378    }
379    async fn execute(&self, arguments: Value, context: &ToolContext) -> Result<Value> {
380        let result = self
381            .remote
382            .execute(self.effective_arguments(arguments)?, context)
383            .await?;
384        let limit = self
385            .binding
386            .restrictions
387            .max_output_bytes
388            .unwrap_or(MAX_FRAME as u64);
389        ensure!(
390            serde_json::to_vec(&result)?.len() as u64 <= limit,
391            "MCP alias result exceeds output limit"
392        );
393        Ok(result)
394    }
395}
396#[async_trait]
397impl Tool for RemoteTool {
398    fn name(&self) -> &str {
399        &self.name
400    }
401    fn description(&self) -> &str {
402        &self.description
403    }
404    fn parameters_schema(&self) -> Value {
405        self.schema.clone()
406    }
407    fn capabilities(&self) -> Option<Vec<Capability>> {
408        Some(vec![
409            Capability::ProcessExecute,
410            Capability::ExternalSideEffect,
411        ])
412    }
413    fn binding_metadata(&self) -> Option<Value> {
414        Some(self.metadata.clone())
415    }
416    fn validate_arguments(&self, arguments: &Value) -> Result<()> {
417        validate_arguments(arguments)
418    }
419    async fn execute(&self, arguments: Value, _: &ToolContext) -> Result<Value> {
420        validate_arguments(&arguments)?;
421        let result = tokio::time::timeout(
422            self.timeout,
423            self.peer.call_tool(CallToolRequestParams {
424                meta: None,
425                name: self.remote.clone().into(),
426                arguments: arguments.as_object().cloned(),
427                task: None,
428            }),
429        )
430        .await
431        .map_err(|_| anyhow!("MCP tool call timed out"))?
432        .map_err(|_| anyhow!("MCP tool call failed"))?;
433        ensure!(result.is_error != Some(true), "MCP tool returned an error");
434        let bytes =
435            serde_json::to_vec(&result).map_err(|_| anyhow!("MCP result encoding failed"))?;
436        ensure!(bytes.len() <= MAX_FRAME, "MCP result exceeds size limit");
437        serde_json::from_slice(&bytes).map_err(|_| anyhow!("MCP result encoding failed"))
438    }
439}
context_harness/agent_runtime/mcp_client.rs

context_harness/agent_runtime/
mcp_client.rs
