1use 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
12pub 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 {
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 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 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 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}