Skip to main content

roboticus_agent/mcp/
mod.rs

1pub mod bridge;
2pub mod client;
3pub mod manager;
4
5use roboticus_core::Result;
6use serde::{Deserialize, Serialize};
7use std::collections::HashMap;
8use tracing::{debug, info};
9
10/// MCP tool descriptor exposed to clients.
11#[derive(Debug, Clone, Serialize, Deserialize)]
12pub struct McpToolDescriptor {
13    pub name: String,
14    pub description: String,
15    pub input_schema: serde_json::Value,
16}
17
18/// MCP resource descriptor.
19#[derive(Debug, Clone, Serialize, Deserialize)]
20pub struct McpResource {
21    pub uri: String,
22    pub name: String,
23    pub description: String,
24    #[serde(default)]
25    pub mime_type: String,
26}
27
28/// Registry for MCP-exposed tools and resources.
29#[derive(Debug, Default)]
30pub struct McpServerRegistry {
31    tools: HashMap<String, McpToolDescriptor>,
32    resources: HashMap<String, McpResource>,
33}
34
35impl McpServerRegistry {
36    pub fn new() -> Self {
37        Self::default()
38    }
39
40    pub fn register_tool(&mut self, descriptor: McpToolDescriptor) {
41        debug!(name = %descriptor.name, "registered MCP tool");
42        self.tools.insert(descriptor.name.clone(), descriptor);
43    }
44
45    pub fn register_resource(&mut self, resource: McpResource) {
46        debug!(uri = %resource.uri, "registered MCP resource");
47        self.resources.insert(resource.uri.clone(), resource);
48    }
49
50    pub fn list_tools(&self) -> Vec<&McpToolDescriptor> {
51        self.tools.values().collect()
52    }
53
54    pub fn get_tool(&self, name: &str) -> Option<&McpToolDescriptor> {
55        self.tools.get(name)
56    }
57
58    pub fn list_resources(&self) -> Vec<&McpResource> {
59        self.resources.values().collect()
60    }
61
62    pub fn get_resource(&self, uri: &str) -> Option<&McpResource> {
63        self.resources.get(uri)
64    }
65
66    pub fn tool_count(&self) -> usize {
67        self.tools.len()
68    }
69
70    pub fn resource_count(&self) -> usize {
71        self.resources.len()
72    }
73}
74
75/// Represents a connection to an external MCP server.
76#[derive(Debug, Clone)]
77pub struct McpClientConnection {
78    pub name: String,
79    pub url: String,
80    pub available_tools: Vec<McpToolDescriptor>,
81    pub available_resources: Vec<McpResource>,
82    pub connected: bool,
83}
84
85impl McpClientConnection {
86    pub fn new(name: String, url: String) -> Self {
87        Self {
88            name,
89            url,
90            available_tools: Vec::new(),
91            available_resources: Vec::new(),
92            connected: false,
93        }
94    }
95
96    /// Discover tools from the remote MCP server.
97    /// In production, this would make HTTP/SSE requests.
98    pub fn discover(&mut self) -> Result<()> {
99        info!(name = %self.name, url = %self.url, "MCP client discovering tools");
100        self.connected = true;
101        Ok(())
102    }
103
104    pub fn is_connected(&self) -> bool {
105        self.connected
106    }
107
108    pub fn disconnect(&mut self) {
109        self.connected = false;
110        self.available_tools.clear();
111        self.available_resources.clear();
112        debug!(name = %self.name, "MCP client disconnected");
113    }
114}
115
116/// Manages multiple MCP client connections.
117#[derive(Debug, Default)]
118pub struct McpClientManager {
119    connections: HashMap<String, McpClientConnection>,
120}
121
122impl McpClientManager {
123    pub fn new() -> Self {
124        Self::default()
125    }
126
127    pub fn add_connection(&mut self, conn: McpClientConnection) {
128        self.connections.insert(conn.name.clone(), conn);
129    }
130
131    pub fn get_connection(&self, name: &str) -> Option<&McpClientConnection> {
132        self.connections.get(name)
133    }
134
135    pub fn get_connection_mut(&mut self, name: &str) -> Option<&mut McpClientConnection> {
136        self.connections.get_mut(name)
137    }
138
139    pub fn list_connections(&self) -> Vec<&McpClientConnection> {
140        self.connections.values().collect()
141    }
142
143    pub fn connected_count(&self) -> usize {
144        self.connections.values().filter(|c| c.connected).count()
145    }
146
147    pub fn total_count(&self) -> usize {
148        self.connections.len()
149    }
150
151    /// Get all tools available across all connected MCP servers.
152    pub fn all_available_tools(&self) -> Vec<(&str, &McpToolDescriptor)> {
153        self.connections
154            .values()
155            .filter(|c| c.connected)
156            .flat_map(|c| c.available_tools.iter().map(move |t| (c.name.as_str(), t)))
157            .collect()
158    }
159
160    pub fn disconnect_all(&mut self) {
161        for conn in self.connections.values_mut() {
162            conn.disconnect();
163        }
164    }
165}
166
167/// Build MCP tool descriptors from the internal tool registry.
168pub fn export_tools_as_mcp(
169    tools: &[(String, String, serde_json::Value)],
170) -> Vec<McpToolDescriptor> {
171    tools
172        .iter()
173        .map(|(name, desc, schema)| McpToolDescriptor {
174            name: name.clone(),
175            description: desc.clone(),
176            input_schema: schema.clone(),
177        })
178        .collect()
179}
180
181#[cfg(test)]
182mod tests {
183    use super::*;
184
185    fn sample_tool() -> McpToolDescriptor {
186        McpToolDescriptor {
187            name: "memory_search".to_string(),
188            description: "Search agent memory".to_string(),
189            input_schema: serde_json::json!({"type": "object", "properties": {"query": {"type": "string"}}}),
190        }
191    }
192
193    fn sample_resource() -> McpResource {
194        McpResource {
195            uri: "roboticus://sessions/current".to_string(),
196            name: "Current Session".to_string(),
197            description: "The active session context".to_string(),
198            mime_type: "application/json".to_string(),
199        }
200    }
201
202    #[test]
203    fn server_registry_tools() {
204        let mut reg = McpServerRegistry::new();
205        reg.register_tool(sample_tool());
206        assert_eq!(reg.tool_count(), 1);
207        assert!(reg.get_tool("memory_search").is_some());
208        assert!(reg.get_tool("nonexistent").is_none());
209    }
210
211    #[test]
212    fn server_registry_resources() {
213        let mut reg = McpServerRegistry::new();
214        reg.register_resource(sample_resource());
215        assert_eq!(reg.resource_count(), 1);
216        assert!(reg.get_resource("roboticus://sessions/current").is_some());
217    }
218
219    #[test]
220    fn server_registry_list() {
221        let mut reg = McpServerRegistry::new();
222        reg.register_tool(sample_tool());
223        reg.register_resource(sample_resource());
224        assert_eq!(reg.list_tools().len(), 1);
225        assert_eq!(reg.list_resources().len(), 1);
226    }
227
228    #[test]
229    fn client_connection_lifecycle() {
230        let mut conn =
231            McpClientConnection::new("test-server".into(), "http://localhost:8080".into());
232        assert!(!conn.is_connected());
233
234        conn.discover().unwrap();
235        assert!(conn.is_connected());
236
237        conn.disconnect();
238        assert!(!conn.is_connected());
239    }
240
241    #[test]
242    fn client_manager_basic() {
243        let mut mgr = McpClientManager::new();
244        let conn = McpClientConnection::new("server-a".into(), "http://a.example.com".into());
245        mgr.add_connection(conn);
246
247        assert_eq!(mgr.total_count(), 1);
248        assert_eq!(mgr.connected_count(), 0);
249        assert!(mgr.get_connection("server-a").is_some());
250    }
251
252    #[test]
253    fn client_manager_discover() {
254        let mut mgr = McpClientManager::new();
255        mgr.add_connection(McpClientConnection::new(
256            "s1".into(),
257            "http://s1.local".into(),
258        ));
259        mgr.add_connection(McpClientConnection::new(
260            "s2".into(),
261            "http://s2.local".into(),
262        ));
263
264        mgr.get_connection_mut("s1").unwrap().discover().unwrap();
265        assert_eq!(mgr.connected_count(), 1);
266
267        mgr.get_connection_mut("s2").unwrap().discover().unwrap();
268        assert_eq!(mgr.connected_count(), 2);
269    }
270
271    #[test]
272    fn client_manager_disconnect_all() {
273        let mut mgr = McpClientManager::new();
274        let mut c1 = McpClientConnection::new("a".into(), "http://a".into());
275        c1.discover().unwrap();
276        mgr.add_connection(c1);
277
278        let mut c2 = McpClientConnection::new("b".into(), "http://b".into());
279        c2.discover().unwrap();
280        mgr.add_connection(c2);
281
282        assert_eq!(mgr.connected_count(), 2);
283        mgr.disconnect_all();
284        assert_eq!(mgr.connected_count(), 0);
285    }
286
287    #[test]
288    fn export_tools_as_mcp_conversion() {
289        let tools = vec![
290            (
291                "tool_a".to_string(),
292                "Description A".to_string(),
293                serde_json::json!({}),
294            ),
295            (
296                "tool_b".to_string(),
297                "Description B".to_string(),
298                serde_json::json!({"type": "object"}),
299            ),
300        ];
301        let mcp_tools = export_tools_as_mcp(&tools);
302        assert_eq!(mcp_tools.len(), 2);
303        assert_eq!(mcp_tools[0].name, "tool_a");
304        assert_eq!(mcp_tools[1].description, "Description B");
305    }
306
307    #[test]
308    fn all_available_tools_across_connections() {
309        let mut mgr = McpClientManager::new();
310        let mut c1 = McpClientConnection::new("s1".into(), "http://s1".into());
311        c1.discover().unwrap();
312        c1.available_tools.push(sample_tool());
313        mgr.add_connection(c1);
314
315        let c2 = McpClientConnection::new("s2".into(), "http://s2".into());
316        mgr.add_connection(c2);
317
318        let all_tools = mgr.all_available_tools();
319        assert_eq!(all_tools.len(), 1);
320        assert_eq!(all_tools[0].0, "s1");
321    }
322
323    #[test]
324    fn mcp_transport_default() {
325        let transport = roboticus_core::config::McpTransport::default();
326        assert!(matches!(
327            transport,
328            roboticus_core::config::McpTransport::Stdio
329        ));
330    }
331
332    #[test]
333    fn tool_descriptor_serde() {
334        let tool = sample_tool();
335        let json = serde_json::to_string(&tool).unwrap();
336        let back: McpToolDescriptor = serde_json::from_str(&json).unwrap();
337        assert_eq!(back.name, tool.name);
338    }
339
340    #[test]
341    fn resource_serde() {
342        let res = sample_resource();
343        let json = serde_json::to_string(&res).unwrap();
344        let back: McpResource = serde_json::from_str(&json).unwrap();
345        assert_eq!(back.uri, res.uri);
346    }
347}