Skip to main content

synaps_cli/mcp/
lazy.rs

1//! Lazy MCP connection — the mcp_connect tool that connects to servers on-demand.
2use serde_json::{json, Value};
3use std::collections::HashMap;
4use std::sync::Arc;
5use tokio::sync::Mutex;
6use crate::Tool;
7
8use super::McpServerConfig;
9use super::connection::McpConnection;
10use super::tool::McpTool;
11
12/// A tool that lazily connects to MCP servers on demand.
13/// Instead of spawning all servers at startup (burning tokens on 65 tool schemas),
14/// this registers ONE tool that the model calls to activate a specific server.
15///
16/// MCP connect gateway — discovers tools from an MCP server and registers them dynamically.
17/// Uses the ToolContext.tool_register_tx channel instead of holding a direct Arc to the registry,
18/// breaking the circular reference that previously existed.
19pub struct McpConnectTool {
20    configs: HashMap<String, McpServerConfig>,
21    connected: Arc<Mutex<std::collections::HashSet<String>>>,
22}
23
24impl McpConnectTool {
25    pub fn new(
26        configs: HashMap<String, McpServerConfig>,
27    ) -> Self {
28        Self {
29            configs,
30            connected: Arc::new(Mutex::new(std::collections::HashSet::new())),
31        }
32    }
33}
34
35#[async_trait::async_trait]
36impl Tool for McpConnectTool {
37    fn name(&self) -> &str { "connect_mcp_server" }
38
39    fn description(&self) -> &str {
40        "Connect to an MCP server and load its tools. Call this before using tools from an external MCP server. Available servers are listed in the description below. Once connected, the server's tools become available for the rest of the session."
41    }
42
43    fn parameters(&self) -> Value {
44        let server_names: Vec<&str> = self.configs.keys().map(|s| s.as_str()).collect();
45        let server_list = server_names.join(", ");
46        json!({
47            "type": "object",
48            "properties": {
49                "server": {
50                    "type": "string",
51                    "description": format!("Name of the MCP server to connect to. Available: {}", server_list)
52                }
53            },
54            "required": ["server"]
55        })
56    }
57
58    async fn execute(&self, params: Value, ctx: crate::ToolContext) -> crate::Result<String> {
59        let server_name = params["server"].as_str()
60            .ok_or_else(|| crate::RuntimeError::Tool("Missing 'server' parameter".to_string()))?;
61
62        // Atomically check-and-mark to prevent double-connect from parallel calls
63        {
64            let mut connected = self.connected.lock().await;
65            if connected.contains(server_name) {
66                return Ok(format!("Server '{}' is already connected.", server_name));
67            }
68            // Mark now — if connection fails, we'll unmark
69            connected.insert(server_name.to_string());
70        }
71
72        let config = self.configs.get(server_name)
73            .ok_or_else(|| {
74                let available: Vec<&str> = self.configs.keys().map(|s| s.as_str()).collect();
75                crate::RuntimeError::Tool(format!(
76                    "Unknown MCP server '{}'. Available: {}", server_name, available.join(", ")
77                ))
78            })?;
79
80        tracing::info!(server = %server_name, "Lazy-connecting to MCP server");
81
82        let mut conn = match McpConnection::start(config).await {
83            Ok(c) => c,
84            Err(e) => {
85                // Unmark on failure so retry is possible
86                self.connected.lock().await.remove(server_name);
87                return Err(crate::RuntimeError::Tool(format!(
88                    "Failed to connect to MCP server '{}': {}", server_name, e
89                )));
90            }
91        };
92
93        let tools = match conn.list_tools().await {
94            Ok(t) => t,
95            Err(e) => {
96                self.connected.lock().await.remove(server_name);
97                return Err(crate::RuntimeError::Tool(format!(
98                    "Failed to list tools from '{}': {}", server_name, e
99                )));
100            }
101        };
102
103        let tool_count = tools.len();
104        let connection = Arc::new(Mutex::new(conn));
105        let mut tool_names = Vec::new();
106        let mut new_tools: Vec<Arc<dyn crate::Tool>> = Vec::new();
107
108        for tool_def in tools {
109            let prefixed_name = format!("ext__{}__{}", server_name, tool_def.name);
110            tool_names.push(format!("{} — {}", tool_def.name,
111                tool_def.description.chars().take(60).collect::<String>()));
112
113            let mcp_tool = McpTool {
114                tool_name: prefixed_name,
115                server_tool_name: tool_def.name.clone(),
116                server_name: server_name.to_string(),
117                description: format!("[MCP:{}] {}", server_name, tool_def.description),
118                input_schema: tool_def.input_schema,
119                connection: Arc::clone(&connection),
120            };
121
122            new_tools.push(Arc::new(mcp_tool));
123        }
124
125        // Send new tools to the runtime for registration (via channel, no circular Arc)
126        if let Some(ref tx) = ctx.capabilities.tool_register_tx {
127            let _ = tx.send(new_tools);
128        }
129
130        tracing::info!(server = %server_name, tools = tool_count, "MCP server connected (lazy)");
131
132        let tool_list = tool_names.join("\n  • ");
133        Ok(format!(
134            "Connected to '{}' — {} tools now available:\n  • {}",
135            server_name, tool_count, tool_list
136        ))
137    }
138}