agent_engine/mcp/
lazy.rs

1//! Lazy MCP connection — the mcp_connect tool that connects to servers on-demand.
2use crate::Tool;
3use serde_json::{json, Value};
4use std::collections::HashMap;
5use std::sync::Arc;
6use tokio::sync::Mutex;
7
8use super::connection::McpConnection;
9use super::tool::McpTool;
10use super::McpServerConfig;
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(configs: HashMap<String, McpServerConfig>) -> Self {
26        Self {
27            configs,
28            connected: Arc::new(Mutex::new(std::collections::HashSet::new())),
29        }
30    }
31}
32
33#[async_trait::async_trait]
34impl Tool for McpConnectTool {
35    fn origin(&self) -> crate::tools::ToolOrigin {
36        crate::tools::ToolOrigin::Builtin
37    }
38
39    fn name(&self) -> &str {
40        "connect_mcp_server"
41    }
42
43    fn description(&self) -> &str {
44        "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."
45    }
46
47    fn parameters(&self) -> Value {
48        let server_names: Vec<&str> = self.configs.keys().map(|s| s.as_str()).collect();
49        let server_list = server_names.join(", ");
50        json!({
51            "type": "object",
52            "properties": {
53                "server": {
54                    "type": "string",
55                    "description": format!("Name of the MCP server to connect to. Available: {}", server_list)
56                }
57            },
58            "required": ["server"]
59        })
60    }
61
62    async fn execute(&self, params: Value, ctx: crate::ToolContext) -> crate::Result<String> {
63        let server_name = params["server"]
64            .as_str()
65            .ok_or_else(|| crate::RuntimeError::Tool("Missing 'server' parameter".to_string()))?;
66
67        // Atomically check-and-mark to prevent double-connect from parallel calls
68        {
69            let mut connected = self.connected.lock().await;
70            if connected.contains(server_name) {
71                return Ok(format!("Server '{}' is already connected.", server_name));
72            }
73            // Mark now — if connection fails, we'll unmark
74            connected.insert(server_name.to_string());
75        }
76
77        let config = self.configs.get(server_name).ok_or_else(|| {
78            let available: Vec<&str> = self.configs.keys().map(|s| s.as_str()).collect();
79            crate::RuntimeError::Tool(format!(
80                "Unknown MCP server '{}'. Available: {}",
81                server_name,
82                available.join(", ")
83            ))
84        })?;
85
86        tracing::info!(server = %server_name, "Lazy-connecting to MCP server");
87
88        let mut conn = match McpConnection::start(config).await {
89            Ok(c) => c,
90            Err(e) => {
91                // Unmark on failure so retry is possible
92                self.connected.lock().await.remove(server_name);
93                return Err(crate::RuntimeError::Tool(format!(
94                    "Failed to connect to MCP server '{}': {}",
95                    server_name, e
96                )));
97            }
98        };
99
100        let tools = match conn.list_tools().await {
101            Ok(t) => t,
102            Err(e) => {
103                self.connected.lock().await.remove(server_name);
104                return Err(crate::RuntimeError::Tool(format!(
105                    "Failed to list tools from '{}': {}",
106                    server_name, e
107                )));
108            }
109        };
110
111        let tool_count = tools.len();
112        let connection = Arc::new(Mutex::new(conn));
113        let mut tool_names = Vec::new();
114        let mut new_tools: Vec<Arc<dyn crate::Tool>> = Vec::new();
115
116        for tool_def in tools {
117            let prefixed_name = format!("ext__{}__{}", server_name, tool_def.name);
118            tool_names.push(format!(
119                "{} — {}",
120                tool_def.name,
121                tool_def.description.chars().take(60).collect::<String>()
122            ));
123
124            let mcp_tool = McpTool {
125                tool_name: prefixed_name,
126                server_tool_name: tool_def.name.clone(),
127                server_name: server_name.to_string(),
128                description: format!("[MCP:{}] {}", server_name, tool_def.description),
129                input_schema: tool_def.input_schema,
130                connection: Arc::clone(&connection),
131            };
132
133            new_tools.push(Arc::new(mcp_tool));
134        }
135
136        // Send new tools to the runtime for registration (via channel, no circular Arc)
137        if let Some(ref tx) = ctx.capabilities.tool_register_tx {
138            let _ = tx.send(new_tools);
139        }
140
141        tracing::info!(server = %server_name, tools = tool_count, "MCP server connected (lazy)");
142
143        let tool_list = tool_names.join("\n  • ");
144        Ok(format!(
145            "Connected to '{}' — {} tools now available:\n  • {}",
146            server_name, tool_count, tool_list
147        ))
148    }
149}
agent_engine/mcp/lazy.rs

agent_engine/mcp/
lazy.rs
