Skip to main content

embacle_server/
mcp_client.rs

1// ABOUTME: MCP stdio client pool that connects to downstream MCP tool servers
2// ABOUTME: Discovers their tools and routes tool calls via rmcp for server-side agent execution
3//
4// SPDX-License-Identifier: Apache-2.0
5// Copyright (c) 2026 dravr.ai
6
7//! # MCP Client Pool
8//!
9//! Connects to one or more downstream MCP tool servers (configured via
10//! `[[mcp_servers]]`) as an MCP **client** over stdio, using the official
11//! [`rmcp`] SDK. On connection it discovers each server's tools and builds a
12//! routing table from tool name to owning server.
13//!
14//! The pool implements [`embacle::McpToolExecutor`], so it can be handed to the
15//! text-based tool loop ([`embacle::mcp_tool_bridge`]) / [`embacle::agent::AgentExecutor`]
16//! to power server-side tool execution on `/v1/chat/completions`.
17//!
18//! Tool name collisions across servers are resolved first-wins, with a warning
19//! logged for every dropped tool so coverage is never silently reduced.
20
21use std::collections::HashMap;
22
23use async_trait::async_trait;
24use embacle::types::RunnerError;
25use embacle::{FunctionDeclaration, McpServerConfig, McpToolExecutor};
26use rmcp::model::{CallToolRequestParams, CallToolResult};
27use rmcp::service::{RoleClient, RunningService};
28use rmcp::transport::TokioChildProcess;
29use rmcp::ServiceExt;
30use serde_json::Value;
31use tokio::process::Command;
32use tracing::{info, warn};
33
34/// A live connection to a single downstream MCP server.
35struct ConnectedServer {
36    /// Logical name from configuration (for diagnostics)
37    name: String,
38    /// Running rmcp client session bound to the spawned subprocess
39    service: RunningService<RoleClient, ()>,
40}
41
42/// A pool of MCP stdio clients exposing their union of tools as a single executor.
43pub struct McpClientPool {
44    /// Connected downstream servers, indexed by position
45    servers: Vec<ConnectedServer>,
46    /// Tool name -> index into `servers`
47    routing: HashMap<String, usize>,
48    /// Tool declarations discovered across all servers (for catalog injection)
49    declarations: Vec<FunctionDeclaration>,
50}
51
52impl McpClientPool {
53    /// Connect to every configured MCP server, discover their tools, and build
54    /// the routing table.
55    ///
56    /// Connection failures for an individual server are fatal: a misconfigured
57    /// tool server should surface loudly rather than silently degrade the agent.
58    ///
59    /// # Errors
60    ///
61    /// Returns [`RunnerError::external_service`] if spawning, initializing, or
62    /// listing tools for any configured server fails.
63    pub async fn connect(configs: &[McpServerConfig]) -> Result<Self, RunnerError> {
64        let mut servers = Vec::with_capacity(configs.len());
65        let mut routing = HashMap::new();
66        let mut declarations = Vec::new();
67
68        for cfg in configs {
69            let mut command = Command::new(&cfg.command);
70            command.args(&cfg.args);
71            for (key, value) in &cfg.env {
72                command.env(key, value);
73            }
74
75            let transport = TokioChildProcess::new(command).map_err(|e| {
76                RunnerError::external_service(
77                    "mcp",
78                    format!("failed to spawn MCP server '{}': {e}", cfg.name),
79                )
80            })?;
81
82            let service = ().serve(transport).await.map_err(|e| {
83                RunnerError::external_service(
84                    "mcp",
85                    format!("failed to initialize MCP server '{}': {e}", cfg.name),
86                )
87            })?;
88
89            let tools = service.list_all_tools().await.map_err(|e| {
90                RunnerError::external_service(
91                    "mcp",
92                    format!("failed to list tools for MCP server '{}': {e}", cfg.name),
93                )
94            })?;
95
96            let server_idx = servers.len();
97            let mut registered = 0_usize;
98            for tool in tools {
99                let name = tool.name.to_string();
100                if routing.contains_key(&name) {
101                    warn!(
102                        tool = %name,
103                        server = %cfg.name,
104                        "Duplicate MCP tool name across servers; keeping first registration, dropping this one"
105                    );
106                    continue;
107                }
108                declarations.push(FunctionDeclaration {
109                    name: name.clone(),
110                    description: tool.description.map(|d| d.to_string()).unwrap_or_default(),
111                    parameters: Some(Value::Object((*tool.input_schema).clone())),
112                });
113                routing.insert(name, server_idx);
114                registered += 1;
115            }
116
117            info!(
118                server = %cfg.name,
119                command = %cfg.command,
120                tools = registered,
121                "Connected to MCP tool server"
122            );
123
124            servers.push(ConnectedServer {
125                name: cfg.name.clone(),
126                service,
127            });
128        }
129
130        info!(
131            servers = servers.len(),
132            tools = declarations.len(),
133            "MCP client pool ready"
134        );
135
136        Ok(Self {
137            servers,
138            routing,
139            declarations,
140        })
141    }
142
143    /// Tool declarations discovered across all connected servers.
144    pub fn declarations(&self) -> &[FunctionDeclaration] {
145        &self.declarations
146    }
147
148    /// Returns true if no servers are connected (no tools available).
149    pub fn is_empty(&self) -> bool {
150        self.servers.is_empty()
151    }
152
153    /// Number of tools available across all connected servers.
154    pub fn tool_count(&self) -> usize {
155        self.declarations.len()
156    }
157}
158
159#[async_trait]
160impl McpToolExecutor for McpClientPool {
161    async fn execute(&self, tool_name: &str, arguments: &Value) -> Result<Value, RunnerError> {
162        let server_idx = *self.routing.get(tool_name).ok_or_else(|| {
163            RunnerError::internal(format!(
164                "no connected MCP server provides tool '{tool_name}'"
165            ))
166        })?;
167        let server = &self.servers[server_idx];
168
169        let mut params = CallToolRequestParams::new(tool_name.to_owned());
170        if let Some(object) = arguments.as_object() {
171            params = params.with_arguments(object.clone());
172        }
173
174        let result = server.service.call_tool(params).await.map_err(|e| {
175            RunnerError::external_service(
176                "mcp",
177                format!("tool '{tool_name}' on server '{}' failed: {e}", server.name),
178            )
179        })?;
180
181        Ok(call_result_to_json(result))
182    }
183}
184
185/// Convert an rmcp `CallToolResult` into a single JSON value for the tool loop.
186///
187/// Prefers `structured_content` when present. Otherwise concatenates text
188/// content blocks, attempting to parse the joined text as JSON before falling
189/// back to a string. Tool-reported errors are wrapped as `{"error": ...}` so the
190/// model can observe and recover rather than aborting the whole request.
191fn call_result_to_json(result: CallToolResult) -> Value {
192    if let Some(structured) = result.structured_content {
193        return structured;
194    }
195
196    let text = result
197        .content
198        .iter()
199        .filter_map(|c| c.as_text().map(|t| t.text.clone()))
200        .collect::<Vec<_>>()
201        .join("\n");
202
203    if result.is_error == Some(true) {
204        return serde_json::json!({ "error": text });
205    }
206
207    serde_json::from_str::<Value>(&text).unwrap_or(Value::String(text))
208}
209
210#[cfg(test)]
211mod tests {
212    use rmcp::model::Content;
213
214    use super::*;
215
216    #[test]
217    fn call_result_prefers_structured_content() {
218        let result = CallToolResult::structured(serde_json::json!({"temp": 72}));
219        let value = call_result_to_json(result);
220        assert_eq!(value["temp"], 72);
221    }
222
223    #[test]
224    fn call_result_parses_json_text() {
225        let result = CallToolResult::success(vec![Content::text(r#"{"ok":true}"#)]);
226        let value = call_result_to_json(result);
227        assert_eq!(value["ok"], true);
228    }
229
230    #[test]
231    fn call_result_falls_back_to_string() {
232        let result = CallToolResult::success(vec![Content::text("plain text answer")]);
233        let value = call_result_to_json(result);
234        assert_eq!(value, Value::String("plain text answer".to_owned()));
235    }
236
237    #[test]
238    fn call_result_wraps_errors() {
239        let result = CallToolResult::error(vec![Content::text("boom")]);
240        let value = call_result_to_json(result);
241        assert_eq!(value["error"], "boom");
242    }
243}