embacle_server/
mcp_client.rs1use 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
34struct ConnectedServer {
36 name: String,
38 service: RunningService<RoleClient, ()>,
40}
41
42pub struct McpClientPool {
44 servers: Vec<ConnectedServer>,
46 routing: HashMap<String, usize>,
48 declarations: Vec<FunctionDeclaration>,
50}
51
52impl McpClientPool {
53 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 pub fn declarations(&self) -> &[FunctionDeclaration] {
145 &self.declarations
146 }
147
148 pub fn is_empty(&self) -> bool {
150 self.servers.is_empty()
151 }
152
153 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
185fn 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}