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::model::{CallToolRequestParams, CallToolResult, Tool};
13use rmcp::service::{RoleClient, RunningService};
14use rmcp::transport::streamable_http_client::StreamableHttpClientTransportConfig;
15use rmcp::transport::{StreamableHttpClientTransport, TokioChildProcess};
16use rmcp::ServiceExt;
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 { command, args, env } => {
34                let mut cmd = tokio::process::Command::new(command);
35                cmd.args(args);
36                for (k, v) in env {
37                    cmd.env(k, v);
38                }
39                let child = TokioChildProcess::new(cmd)
40                    .map_err(|e| format!("spawn `{command}` failed: {e}"))?;
41                ().serve(child)
42                    .await
43                    .map_err(|e| format!("MCP handshake failed (stdio): {e}"))
44            }
45            ResolvedTransport::Http { url, headers } => {
46                let mut cfg = StreamableHttpClientTransportConfig::with_uri(url.clone());
47                if !headers.is_empty() {
48                    let mut custom = std::collections::HashMap::new();
49                    for (k, v) in headers {
50                        let name = http::HeaderName::from_bytes(k.as_bytes())
51                            .map_err(|e| format!("invalid header name `{k}`: {e}"))?;
52                        let val = http::HeaderValue::from_str(v)
53                            .map_err(|e| format!("invalid header value for `{k}`: {e}"))?;
54                        custom.insert(name, val);
55                    }
56                    cfg = cfg.custom_headers(custom);
57                }
58                let t = StreamableHttpClientTransport::from_config(cfg);
59                ().serve(t)
60                    .await
61                    .map_err(|e| format!("MCP handshake failed (http): {e}"))
62            }
63        }
64    };
65    tokio::time::timeout(timeout, connect)
66        .await
67        .map_err(|_| "downstream connect timed out".to_string())?
68}
69
70/// List tools on an already-connected session (bounded by `timeout`).
71pub async fn list_tools_on(
72    service: &ClientService,
73    timeout: Duration,
74) -> Result<Vec<Tool>, String> {
75    tokio::time::timeout(timeout, service.list_all_tools())
76        .await
77        .map_err(|_| "downstream tools/list timed out".to_string())
78        .and_then(|r| r.map_err(|e| format!("downstream tools/list failed: {e}")))
79}
80
81/// Call a tool on an already-connected session (bounded by `timeout`).
82pub async fn call_tool_on(
83    service: &ClientService,
84    tool: &str,
85    arguments: Map<String, Value>,
86    timeout: Duration,
87) -> Result<CallToolResult, String> {
88    let param = CallToolRequestParams::new(tool.to_string()).with_arguments(arguments);
89    tokio::time::timeout(timeout, service.call_tool(param))
90        .await
91        .map_err(|_| "downstream tools/call timed out".to_string())
92        .and_then(|r| r.map_err(|e| format!("downstream tools/call failed: {e}")))
93}
94
95/// List a downstream server's tools (connect → `tools/list` → disconnect).
96pub async fn fetch_tools(
97    transport: &ResolvedTransport,
98    timeout: Duration,
99) -> Result<Vec<Tool>, String> {
100    let service = open(transport, timeout).await?;
101    let listed = list_tools_on(&service, timeout).await;
102    let _ = service.cancel().await;
103    listed
104}
105
106/// Proxy a single tool call to a downstream server (connect → `tools/call` →
107/// disconnect).
108pub async fn proxy_call(
109    transport: &ResolvedTransport,
110    tool: &str,
111    arguments: Map<String, Value>,
112    timeout: Duration,
113) -> Result<CallToolResult, String> {
114    let service = open(transport, timeout).await?;
115    let called = call_tool_on(&service, tool, arguments, timeout).await;
116    let _ = service.cancel().await;
117    called
118}
119
120/// Flatten a downstream [`CallToolResult`] into plain text. Text blocks are
121/// concatenated; non-text blocks (images/resources) are summarized so the proxy
122/// never returns binary blobs into the model context.
123pub fn result_to_text(result: &CallToolResult) -> String {
124    let mut parts: Vec<String> = Vec::new();
125    for c in &result.content {
126        if let Some(t) = c.as_text() {
127            parts.push(t.text.clone());
128        } else if c.as_image().is_some() {
129            parts.push("[image content omitted by gateway]".to_string());
130        } else {
131            parts.push("[non-text content omitted by gateway]".to_string());
132        }
133    }
134    parts.join("\n")
135}
lean_ctx/core/gateway/client.rs

lean_ctx/core/gateway/
client.rs
