lean_ctx/core/gateway/
client.rs

1//! Downstream MCP client (#210).
2//!
3//! A *real* MCP client built on the official `rmcp` SDK — no bespoke JSON-RPC.
4//! Each operation opens a fresh connection (performs the `initialize`
5//! handshake), does its work, and shuts the connection down. This keeps the
6//! gateway stateless and robust (no stale child processes / sessions); the
7//! expensive part — listing the full catalog — is amortized by the TTL cache in
8//! [`super::catalog`].
9
10use std::time::Duration;
11
12use rmcp::ServiceExt;
13use rmcp::model::{CallToolRequestParams, CallToolResult, Tool};
14use rmcp::service::{RoleClient, RunningService};
15use rmcp::transport::streamable_http_client::StreamableHttpClientTransportConfig;
16use rmcp::transport::{StreamableHttpClientTransport, TokioChildProcess};
17use serde_json::{Map, Value};
18
19use super::config::ResolvedTransport;
20
21/// A connected downstream MCP client session. Transport-erased: stdio, HTTP,
22/// and (in tests) in-process duplex all collapse to this one type.
23pub type ClientService = RunningService<RoleClient, ()>;
24
25/// Open a connection to a downstream MCP server (runs the MCP `initialize`
26/// handshake). The whole connect is bounded by `timeout`.
27pub async fn open(
28    transport: &ResolvedTransport,
29    timeout: Duration,
30) -> Result<ClientService, String> {
31    let connect = async {
32        match transport {
33            ResolvedTransport::Stdio {
34                command,
35                args,
36                env,
37                binary_sha256,
38                capabilities,
39            } => {
40                // Binary-hash pin (#403, P3): if the addon pinned its binary's
41                // sha256, verify the file on PATH before doing anything else, so
42                // a swapped executable is refused (fail-closed). No-op when
43                // unpinned.
44                crate::core::addons::binhash::verify_binary(command, binary_sha256)?;
45                // Per-addon OS sandbox (#865, P1): declared capabilities drive
46                // the profile (network/filesystem); absent caps fall back to the
47                // legacy `addons.sandbox` mode. May wrap with sandbox-exec /
48                // bwrap, or refuse to spawn (strict / enforce_capabilities).
49                let (spawn_cmd, spawn_args) =
50                    crate::core::addons::sandbox::apply_for(command, args, capabilities.as_ref())?;
51                let mut cmd = tokio::process::Command::new(&spawn_cmd);
52                cmd.args(&spawn_args);
53                // Secure-by-default environment (P1): a capability-declaring
54                // addon gets a scrubbed env (base allowlist + declared names);
55                // legacy addons inherit the host env unchanged.
56                crate::core::addons::env_scrub::apply_env(&mut cmd, env, capabilities.as_ref());
57                let child = TokioChildProcess::new(cmd)
58                    .map_err(|e| format!("spawn `{command}` failed: {e}"))?;
59                ().serve(child)
60                    .await
61                    .map_err(|e| format!("MCP handshake failed (stdio): {e}"))
62            }
63            ResolvedTransport::Http { url, headers } => {
64                let mut cfg = StreamableHttpClientTransportConfig::with_uri(url.clone());
65                if !headers.is_empty() {
66                    let mut custom = std::collections::HashMap::new();
67                    for (k, v) in headers {
68                        let name = http::HeaderName::from_bytes(k.as_bytes())
69                            .map_err(|e| format!("invalid header name `{k}`: {e}"))?;
70                        let val = http::HeaderValue::from_str(v)
71                            .map_err(|e| format!("invalid header value for `{k}`: {e}"))?;
72                        custom.insert(name, val);
73                    }
74                    cfg = cfg.custom_headers(custom);
75                }
76                let t = StreamableHttpClientTransport::from_config(cfg);
77                ().serve(t)
78                    .await
79                    .map_err(|e| format!("MCP handshake failed (http): {e}"))
80            }
81        }
82    };
83    tokio::time::timeout(timeout, connect)
84        .await
85        .map_err(|_| "downstream connect timed out".to_string())?
86}
87
88/// List tools on an already-connected session (bounded by `timeout`).
89pub async fn list_tools_on(
90    service: &ClientService,
91    timeout: Duration,
92) -> Result<Vec<Tool>, String> {
93    tokio::time::timeout(timeout, service.list_all_tools())
94        .await
95        .map_err(|_| "downstream tools/list timed out".to_string())
96        .and_then(|r| r.map_err(|e| format!("downstream tools/list failed: {e}")))
97}
98
99/// Call a tool on an already-connected session (bounded by `timeout`).
100pub async fn call_tool_on(
101    service: &ClientService,
102    tool: &str,
103    arguments: Map<String, Value>,
104    timeout: Duration,
105) -> Result<CallToolResult, String> {
106    let param = CallToolRequestParams::new(tool.to_string()).with_arguments(arguments);
107    tokio::time::timeout(timeout, service.call_tool(param))
108        .await
109        .map_err(|_| "downstream tools/call timed out".to_string())
110        .and_then(|r| r.map_err(|e| format!("downstream tools/call failed: {e}")))
111}
112
113/// List a downstream server's tools (connect → `tools/list` → disconnect).
114pub async fn fetch_tools(
115    transport: &ResolvedTransport,
116    timeout: Duration,
117) -> Result<Vec<Tool>, String> {
118    let service = open(transport, timeout).await?;
119    let listed = list_tools_on(&service, timeout).await;
120    let _ = service.cancel().await;
121    listed
122}
123
124/// Proxy a single tool call to a downstream server (connect → `tools/call` →
125/// disconnect).
126pub async fn proxy_call(
127    transport: &ResolvedTransport,
128    tool: &str,
129    arguments: Map<String, Value>,
130    timeout: Duration,
131) -> Result<CallToolResult, String> {
132    let service = open(transport, timeout).await?;
133    let called = call_tool_on(&service, tool, arguments, timeout).await;
134    let _ = service.cancel().await;
135    called
136}
137
138/// Flatten a downstream [`CallToolResult`] into plain text. Text blocks are
139/// concatenated; non-text blocks (images/resources) are summarized so the proxy
140/// never returns binary blobs into the model context.
141pub fn result_to_text(result: &CallToolResult) -> String {
142    let mut parts: Vec<String> = Vec::new();
143    for c in &result.content {
144        if let Some(t) = c.as_text() {
145            parts.push(t.text.clone());
146        } else if c.as_image().is_some() {
147            parts.push("[image content omitted by gateway]".to_string());
148        } else {
149            parts.push("[non-text content omitted by gateway]".to_string());
150        }
151    }
152    parts.join("\n")
153}
lean_ctx/core/gateway/client.rs

lean_ctx/core/gateway/
client.rs
