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//! [`open`] performs the MCP `initialize` handshake; sessions are then kept alive
5//! and reused by the [`super::pool`] (#1078), so list/call operations no longer
6//! pay spawn+handshake latency on every invocation. The pool only ever hands out
7//! a *live* session (it sweeps closed ones at acquire), so requests never reach a
8//! dead pipe. Listing is idempotent, so a mid-flight transport failure is evicted
9//! and reopened once; a *call* is never auto-retried (a downstream tool may be
10//! non-idempotent — re-issuing could double-execute a side effect), only evicted.
11//! The catalog-listing cost is additionally amortized by the TTL cache in
12//! [`super::catalog`].
13
14use std::time::Duration;
15
16use rmcp::ServiceExt;
17use rmcp::model::{CallToolRequestParams, CallToolResult, Tool};
18use rmcp::service::{RoleClient, RunningService};
19use rmcp::transport::streamable_http_client::StreamableHttpClientTransportConfig;
20use rmcp::transport::{StreamableHttpClientTransport, TokioChildProcess};
21use serde_json::{Map, Value};
22
23use super::config::ResolvedTransport;
24
25/// A connected downstream MCP client session. Transport-erased: stdio, HTTP,
26/// and (in tests) in-process duplex all collapse to this one type.
27pub type ClientService = RunningService<RoleClient, ()>;
28
29/// Open a connection to a downstream MCP server (runs the MCP `initialize`
30/// handshake). The whole connect is bounded by `timeout`.
31pub async fn open(
32    transport: &ResolvedTransport,
33    timeout: Duration,
34) -> Result<ClientService, String> {
35    let connect = async {
36        match transport {
37            ResolvedTransport::Stdio {
38                command,
39                args,
40                env,
41                binary_sha256,
42                capabilities,
43            } => {
44                // Binary-hash pin (#403, P3): if the addon pinned its binary's
45                // sha256, verify the file on PATH before doing anything else, so
46                // a swapped executable is refused (fail-closed). No-op when
47                // unpinned.
48                crate::core::addons::binhash::verify_binary(command, binary_sha256)?;
49                // Per-addon OS sandbox (#865, P1): declared capabilities drive
50                // the profile (network/filesystem); absent caps fall back to the
51                // legacy `addons.sandbox` mode. May wrap with sandbox-exec /
52                // bwrap, or refuse to spawn (strict / enforce_capabilities).
53                let (spawn_cmd, spawn_args) =
54                    crate::core::addons::sandbox::apply_for(command, args, capabilities.as_ref())?;
55                let mut cmd = tokio::process::Command::new(&spawn_cmd);
56                cmd.args(&spawn_args);
57                // Secure-by-default environment (P1): a capability-declaring
58                // addon gets a scrubbed env (base allowlist + declared names);
59                // legacy addons inherit the host env unchanged.
60                crate::core::addons::env_scrub::apply_env(&mut cmd, env, capabilities.as_ref());
61                let child = TokioChildProcess::new(cmd)
62                    .map_err(|e| format!("spawn `{command}` failed: {e}"))?;
63                ().serve(child)
64                    .await
65                    .map_err(|e| format!("MCP handshake failed (stdio): {e}"))
66            }
67            ResolvedTransport::Http { url, headers } => {
68                let mut cfg = StreamableHttpClientTransportConfig::with_uri(url.clone());
69                if !headers.is_empty() {
70                    let mut custom = std::collections::HashMap::new();
71                    for (k, v) in headers {
72                        let name = http::HeaderName::from_bytes(k.as_bytes())
73                            .map_err(|e| format!("invalid header name `{k}`: {e}"))?;
74                        let val = http::HeaderValue::from_str(v)
75                            .map_err(|e| format!("invalid header value for `{k}`: {e}"))?;
76                        custom.insert(name, val);
77                    }
78                    cfg = cfg.custom_headers(custom);
79                }
80                let t = StreamableHttpClientTransport::from_config(cfg);
81                ().serve(t)
82                    .await
83                    .map_err(|e| format!("MCP handshake failed (http): {e}"))
84            }
85        }
86    };
87    tokio::time::timeout(timeout, connect)
88        .await
89        .map_err(|_| "downstream connect timed out".to_string())?
90}
91
92/// List tools on an already-connected session (bounded by `timeout`).
93pub async fn list_tools_on(
94    service: &ClientService,
95    timeout: Duration,
96) -> Result<Vec<Tool>, String> {
97    tokio::time::timeout(timeout, service.list_all_tools())
98        .await
99        .map_err(|_| "downstream tools/list timed out".to_string())
100        .and_then(|r| r.map_err(|e| format!("downstream tools/list failed: {e}")))
101}
102
103/// Call a tool on an already-connected session (bounded by `timeout`).
104pub async fn call_tool_on(
105    service: &ClientService,
106    tool: &str,
107    arguments: Map<String, Value>,
108    timeout: Duration,
109) -> Result<CallToolResult, String> {
110    let param = CallToolRequestParams::new(tool.to_string()).with_arguments(arguments);
111    tokio::time::timeout(timeout, service.call_tool(param))
112        .await
113        .map_err(|_| "downstream tools/call timed out".to_string())
114        .and_then(|r| r.map_err(|e| format!("downstream tools/call failed: {e}")))
115}
116
117/// List a downstream server's tools over a pooled session (`tools/list`).
118/// [`super::pool::acquire`] guarantees a live session, so the only failure left
119/// is a rare mid-flight transport death; because listing is idempotent we evict
120/// the suspect session and reopen once. A timeout is surfaced (not retried).
121pub async fn fetch_tools(
122    transport: &ResolvedTransport,
123    timeout: Duration,
124) -> Result<Vec<Tool>, String> {
125    let key = super::pool::key(transport);
126    let service = super::pool::acquire(transport, timeout).await?;
127    match list_tools_on(&service, timeout).await {
128        Ok(tools) => Ok(tools),
129        Err(e) => {
130            super::pool::evict(key);
131            if is_broken_connection(&e) {
132                let service = super::pool::acquire(transport, timeout).await?;
133                list_tools_on(&service, timeout).await
134            } else {
135                Err(e)
136            }
137        }
138    }
139}
140
141/// Proxy a single tool call to a downstream server over a pooled session
142/// (`tools/call`). [`super::pool::acquire`] only returns a live session, so the
143/// request is never sent into a dead pipe. A failed call is **never** retried:
144/// a downstream tool may be non-idempotent, so re-issuing could double-execute a
145/// side effect. On any failure we evict the (now-suspect) session — the next
146/// call reopens cleanly — and surface the error for the caller to decide.
147pub async fn proxy_call(
148    transport: &ResolvedTransport,
149    tool: &str,
150    arguments: Map<String, Value>,
151    timeout: Duration,
152) -> Result<CallToolResult, String> {
153    let key = super::pool::key(transport);
154    let service = super::pool::acquire(transport, timeout).await?;
155    let result = call_tool_on(&service, tool, arguments, timeout).await;
156    if result.is_err() {
157        super::pool::evict(key);
158    }
159    result
160}
161
162/// Whether `err` indicates the pooled connection is broken (so a fresh reopen +
163/// retry of an *idempotent* op is safe), as opposed to a timeout (the request
164/// may still be running) or a higher-level failure. Errors here are produced by
165/// this module, so the match is on our own stable strings.
166fn is_broken_connection(err: &str) -> bool {
167    !err.contains("timed out")
168}
169
170/// Flatten a downstream [`CallToolResult`] into plain text. Text blocks are
171/// concatenated; non-text blocks (images/resources) are summarized so the proxy
172/// never returns binary blobs into the model context.
173pub fn result_to_text(result: &CallToolResult) -> String {
174    let mut parts: Vec<String> = Vec::new();
175    for c in &result.content {
176        if let Some(t) = c.as_text() {
177            parts.push(t.text.clone());
178        } else if c.as_image().is_some() {
179            parts.push("[image content omitted by gateway]".to_string());
180        } else {
181            parts.push("[non-text content omitted by gateway]".to_string());
182        }
183    }
184    parts.join("\n")
185}
lean_ctx/core/gateway/client.rs

lean_ctx/core/gateway/
client.rs
