telos_agent/integrations/mcp/
manager.rs

1use std::collections::HashMap;
2use std::path::Path;
3
4use serde_json::Value;
5use tokio::sync::Mutex;
6
7use crate::error::AgentError;
8use crate::integrations::mcp::client::{McpClient, McpTool};
9use crate::integrations::mcp::config::McpServerConfig;
10
11/// Manages multiple MCP server connections.
12pub struct McpManager {
13    servers: Mutex<HashMap<String, McpServerHandle>>,
14}
15
16struct McpServerHandle {
17    config: McpServerConfig,
18    client: McpClient,
19    connected: bool,
20}
21
22impl McpManager {
23    /// Create a manager from a set of named server configs.
24    pub fn new(servers: HashMap<String, McpServerConfig>) -> Self {
25        let handles = servers
26            .into_iter()
27            .map(|(id, config)| {
28                let client = McpClient::new(config.clone());
29                (id, McpServerHandle { config, client, connected: false })
30            })
31            .collect();
32        Self { servers: Mutex::new(handles) }
33    }
34
35    /// Load from `.tiny-agent/mcp.json`.
36    ///
37    /// The expected JSON format is:
38    /// ```json
39    /// { "mcpServers": { "name": { "command": "...", "args": [...], "auto_connect": true } } }
40    /// ```
41    pub fn load_config(path: &Path) -> Result<Self, AgentError> {
42        if !path.exists() {
43            return Ok(Self { servers: Mutex::new(HashMap::new()) });
44        }
45        let content = std::fs::read_to_string(path)
46            .map_err(|e| AgentError::Config(format!("failed to read MCP config: {e}")))?;
47        let config: Value = serde_json::from_str(&content)
48            .map_err(|e| AgentError::Config(format!("failed to parse MCP config: {e}")))?;
49        let servers_map = config.get("mcpServers").and_then(|v| v.as_object());
50        let mut servers = HashMap::new();
51        if let Some(map) = servers_map {
52            for (name, server_cfg) in map {
53                let command = server_cfg.get("command").and_then(|v| v.as_str()).unwrap_or("");
54                if command.is_empty() {
55                    tracing::warn!(
56                        server = %name,
57                        "MCP server config missing 'command' field — skipping"
58                    );
59                    continue;
60                }
61                let args: Vec<String> = server_cfg
62                    .get("args")
63                    .and_then(|v| v.as_array())
64                    .map(|a| a.iter().filter_map(|v| v.as_str().map(String::from)).collect())
65                    .unwrap_or_default();
66                let auto_connect =
67                    server_cfg.get("auto_connect").and_then(|v| v.as_bool()).unwrap_or(true);
68                let cfg = McpServerConfig {
69                    command: command.to_string(),
70                    args,
71                    auto_connect,
72                    ..McpServerConfig::default()
73                };
74                servers.insert(name.clone(), cfg);
75            }
76        }
77        Ok(Self::new(servers))
78    }
79
80    /// Register a new server configuration. If a server with the same id
81    /// already exists, it is replaced.
82    pub async fn register_server(&self, id: String, config: McpServerConfig) {
83        let mut servers = self.servers.lock().await;
84        let client = McpClient::new(config.clone());
85        servers.insert(id, McpServerHandle { config, client, connected: false });
86    }
87
88    /// Register multiple servers from a map of configs.
89    pub async fn register_servers(&self, new_servers: HashMap<String, McpServerConfig>) {
90        let mut servers = self.servers.lock().await;
91        for (id, config) in new_servers {
92            let client = McpClient::new(config.clone());
93            servers.insert(id, McpServerHandle { config, client, connected: false });
94        }
95    }
96
97    /// Connect all servers with `auto_connect` enabled.
98    pub async fn connect_all(&self) {
99        let mut servers = self.servers.lock().await;
100        for (id, handle) in servers.iter_mut() {
101            if handle.config.auto_connect && !handle.connected {
102                match handle.client.connect().await {
103                    Ok(()) => {
104                        handle.connected = true;
105                        tracing::info!(server = %id, "MCP server connected");
106                    }
107                    Err(e) => {
108                        tracing::warn!(server = %id, error = %e, "MCP server failed to connect");
109                    }
110                }
111            }
112        }
113    }
114
115    /// Collect all tools from all connected servers.
116    ///
117    /// Returns `(server_id, tool)` pairs.
118    pub async fn all_tools(&self) -> Vec<(String, McpTool)> {
119        let servers = self.servers.lock().await;
120        servers
121            .iter()
122            .filter(|(_, h)| h.connected)
123            .flat_map(|(id, h)| h.client.tools().into_iter().map(move |t| (id.clone(), t)))
124            .collect()
125    }
126
127    /// Call a tool on a specific server.
128    pub async fn call_tool(
129        &self,
130        server_id: &str,
131        tool_name: &str,
132        args: Value,
133    ) -> Result<Value, AgentError> {
134        let servers = self.servers.lock().await;
135        let handle = servers.get(server_id).ok_or_else(|| {
136            AgentError::ToolNotFound(format!("MCP server '{server_id}' not found"))
137        })?;
138        handle.client.call_tool(tool_name, args).await
139    }
140
141    /// Disconnect all servers.
142    pub async fn disconnect_all(&self) {
143        let servers = self.servers.lock().await;
144        for (_, handle) in servers.iter() {
145            handle.client.disconnect();
146        }
147    }
148}
149
150#[cfg(test)]
151mod tests {
152    use super::*;
153    use std::path::PathBuf;
154
155    #[test]
156    fn load_config_file_not_found() {
157        let path = PathBuf::from("/tmp/nonexistent-mcp-config.json");
158        let manager =
159            McpManager::load_config(&path).expect("load_config should succeed for missing file");
160        assert!(manager.servers.try_lock().unwrap().is_empty());
161    }
162
163    #[test]
164    fn load_config_valid_json() {
165        let dir = tempfile::tempdir().expect("failed to create temp dir");
166        let config_path = dir.path().join("mcp.json");
167        let json = r#"{
168            "mcpServers": {
169                "fs": {
170                    "command": "npx",
171                    "args": ["-y", "@modelcontextprotocol/server-filesystem", "/tmp"],
172                    "auto_connect": true
173                },
174                "noargs": {
175                    "command": "echo",
176                    "auto_connect": false
177                }
178            }
179        }"#;
180        std::fs::write(&config_path, json).unwrap();
181
182        let manager = McpManager::load_config(&config_path)
183            .expect("load_config should succeed for valid JSON");
184        let servers = manager.servers.try_lock().unwrap();
185        assert_eq!(servers.len(), 2);
186        let fs = servers.get("fs").expect("fs server should exist");
187        assert_eq!(fs.config.command, "npx");
188        assert_eq!(fs.config.args, vec!["-y", "@modelcontextprotocol/server-filesystem", "/tmp"]);
189        assert!(fs.config.auto_connect);
190
191        let noargs = servers.get("noargs").expect("noargs server should exist");
192        assert_eq!(noargs.config.command, "echo");
193        assert!(!noargs.config.auto_connect);
194    }
195
196    #[test]
197    fn load_config_invalid_json_returns_error() {
198        let dir = tempfile::tempdir().expect("failed to create temp dir");
199        let config_path = dir.path().join("bad.json");
200        std::fs::write(&config_path, "not json").unwrap();
201
202        let result = McpManager::load_config(&config_path);
203        match result {
204            Err(AgentError::Config(_)) => {} // expected
205            Err(e) => panic!("expected AgentError::Config, got: {e}"),
206            Ok(_) => panic!("expected error, got Ok"),
207        }
208    }
209
210    #[test]
211    fn new_creates_empty_manager() {
212        let manager = McpManager::new(HashMap::new());
213        let servers = manager.servers.try_lock().unwrap();
214        assert!(servers.is_empty());
215    }
216
217    #[test]
218    fn new_creates_manager_with_servers() {
219        let mut servers = HashMap::new();
220        servers.insert("test".to_string(), McpServerConfig::new("echo", vec![]));
221        let manager = McpManager::new(servers);
222        let handles = manager.servers.try_lock().unwrap();
223        assert_eq!(handles.len(), 1);
224        let handle = handles.get("test").unwrap();
225        assert_eq!(handle.config.command, "echo");
226        assert!(!handle.connected);
227    }
228}
telos_agent/integrations/mcp/manager.rs

telos_agent/integrations/mcp/
manager.rs
