| 1 | //! Async MCP (Model Context Protocol) Implementation |
| 2 | //! |
| 3 | //! This module provides full async support for MCP servers with: |
| 4 | //! - Connection pooling for server reuse |
| 5 | //! - Automatic tool discovery via `tools/list` |
| 6 | //! - Configurable timeouts per-server and globally |
| 7 | |
| 8 | use std::collections::HashMap; |
| 9 | use std::fs; |
| 10 | use std::path::Path; |
| 11 | use std::sync::atomic::{AtomicU64, Ordering}; |
| 12 | use std::time::Duration; |
| 13 | |
| 14 | use anyhow::{Context, Result}; |
| 15 | use serde::{Deserialize, Serialize}; |
| 16 | use tokio::io::{AsyncBufReadExt, AsyncWriteExt}; |
| 17 | use tokio::process::{Child, ChildStdin, ChildStdout}; |
| 18 | |
| 19 | use crate::network_policy::{Decision, NetworkPolicyDecider, host_from_url}; |
| 20 | use crate::utils::write_atomic; |
| 21 | |
| 22 | // === Error diagnostics helpers (#71) === |
| 23 | |
| 24 | /// Bytes of a non-2xx response body to surface in connection errors. |
| 25 | const ERROR_BODY_PREVIEW_BYTES: usize = 200; |
| 26 | |
| 27 | /// Mask a URL so any embedded credentials in the userinfo portion (e.g. |
| 28 | /// `https://user:secret@host`) are replaced with `***`. Failures fall back to |
| 29 | /// the original string so we don't lose context — we never want masking to |
| 30 | /// produce an empty error. |
| 31 | fn mask_url_secrets(url: &str) -> String { |
| 32 | if let Ok(parsed) = reqwest::Url::parse(url) { |
| 33 | let mut clone = parsed.clone(); |
| 34 | if !parsed.username().is_empty() || parsed.password().is_some() { |
| 35 | let _ = clone.set_username("***"); |
| 36 | let _ = clone.set_password(Some("***")); |
| 37 | } |
| 38 | return clone.to_string(); |
| 39 | } |
| 40 | url.to_string() |
| 41 | } |
| 42 | |
| 43 | /// Mask any obvious token-like substrings in a body excerpt before surfacing |
| 44 | /// it. Conservative: replaces `Bearer <token>` and `api_key=...` shapes. |
| 45 | fn redact_body_preview(body: &str) -> String { |
| 46 | let mut out = body.to_string(); |
| 47 | if let Some(idx) = out.to_lowercase().find("bearer ") { |
| 48 | let tail_start = idx + "bearer ".len(); |
| 49 | if tail_start < out.len() { |
| 50 | let end = out[tail_start..] |
| 51 | .find(|c: char| c.is_whitespace() || c == '"' || c == ',') |
| 52 | .map_or(out.len(), |off| tail_start + off); |
| 53 | out.replace_range(tail_start..end, "***"); |
| 54 | } |
| 55 | } |
| 56 | for needle in ["api_key=", "apikey=", "api-key=", "token="] { |
| 57 | if let Some(idx) = out.to_lowercase().find(needle) { |
| 58 | let tail_start = idx + needle.len(); |
| 59 | let end = out[tail_start..] |
| 60 | .find(|c: char| c.is_whitespace() || c == '&' || c == '"' || c == ',') |
| 61 | .map_or(out.len(), |off| tail_start + off); |
| 62 | out.replace_range(tail_start..end, "***"); |
| 63 | } |
| 64 | } |
| 65 | out |
| 66 | } |
| 67 | |
| 68 | /// Read up to `max_bytes` of a reqwest Response body and produce a single-line |
| 69 | /// excerpt suitable for an error message. Best-effort — if the body can't be |
| 70 | /// read, returns the literal string `<no body>`. |
| 71 | async fn bounded_body_excerpt(response: reqwest::Response, max_bytes: usize) -> String { |
| 72 | let body_text = response.text().await.unwrap_or_default(); |
| 73 | if body_text.is_empty() { |
| 74 | return "<no body>".to_string(); |
| 75 | } |
| 76 | let trimmed: String = body_text.chars().take(max_bytes).collect(); |
| 77 | let suffix = if body_text.len() > trimmed.len() { |
| 78 | "…" |
| 79 | } else { |
| 80 | "" |
| 81 | }; |
| 82 | let one_line = trimmed.replace(['\n', '\r'], " "); |
| 83 | format!("{}{}", redact_body_preview(&one_line), suffix) |
| 84 | } |
| 85 | |
| 86 | // === Configuration Types === |
| 87 | |
| 88 | /// Full MCP configuration from mcp.json |
| 89 | #[derive(Debug, Clone, Default, Deserialize, Serialize)] |
| 90 | pub struct McpConfig { |
| 91 | #[serde(default)] |
| 92 | pub timeouts: McpTimeouts, |
| 93 | #[serde(default, alias = "mcpServers")] |
| 94 | pub servers: HashMap<String, McpServerConfig>, |
| 95 | } |
| 96 | |
| 97 | /// Global timeout configuration |
| 98 | #[derive(Debug, Clone, Copy, Deserialize, Serialize)] |
| 99 | #[allow(clippy::struct_field_names)] |
| 100 | pub struct McpTimeouts { |
| 101 | #[serde(default = "default_connect_timeout")] |
| 102 | pub connect_timeout: u64, |
| 103 | #[serde(default = "default_execute_timeout")] |
| 104 | pub execute_timeout: u64, |
| 105 | #[serde(default = "default_read_timeout")] |
| 106 | pub read_timeout: u64, |
| 107 | } |
| 108 | |
| 109 | fn default_connect_timeout() -> u64 { |
| 110 | 10 |
| 111 | } |
| 112 | fn default_execute_timeout() -> u64 { |
| 113 | 60 |
| 114 | } |
| 115 | fn default_read_timeout() -> u64 { |
| 116 | 120 |
| 117 | } |
| 118 | |
| 119 | impl Default for McpTimeouts { |
| 120 | fn default() -> Self { |
| 121 | Self { |
| 122 | connect_timeout: default_connect_timeout(), |
| 123 | execute_timeout: default_execute_timeout(), |
| 124 | read_timeout: default_read_timeout(), |
| 125 | } |
| 126 | } |
| 127 | } |
| 128 | |
| 129 | /// Configuration for a single MCP server |
| 130 | #[derive(Debug, Clone, Deserialize, Serialize)] |
| 131 | pub struct McpServerConfig { |
| 132 | pub command: Option<String>, |
| 133 | #[serde(default)] |
| 134 | pub args: Vec<String>, |
| 135 | #[serde(default)] |
| 136 | pub env: HashMap<String, String>, |
| 137 | pub url: Option<String>, |
| 138 | #[serde(default)] |
| 139 | pub connect_timeout: Option<u64>, |
| 140 | #[serde(default)] |
| 141 | pub execute_timeout: Option<u64>, |
| 142 | #[serde(default)] |
| 143 | pub read_timeout: Option<u64>, |
| 144 | #[serde(default)] |
| 145 | pub disabled: bool, |
| 146 | #[serde(default = "default_enabled")] |
| 147 | pub enabled: bool, |
| 148 | #[serde(default)] |
| 149 | pub required: bool, |
| 150 | #[serde(default)] |
| 151 | pub enabled_tools: Vec<String>, |
| 152 | #[serde(default)] |
| 153 | pub disabled_tools: Vec<String>, |
| 154 | } |
| 155 | |
| 156 | fn default_enabled() -> bool { |
| 157 | true |
| 158 | } |
| 159 | |
| 160 | impl McpServerConfig { |
| 161 | pub fn effective_connect_timeout(&self, global: &McpTimeouts) -> u64 { |
| 162 | self.connect_timeout.unwrap_or(global.connect_timeout) |
| 163 | } |
| 164 | |
| 165 | pub fn effective_execute_timeout(&self, global: &McpTimeouts) -> u64 { |
| 166 | self.execute_timeout.unwrap_or(global.execute_timeout) |
| 167 | } |
| 168 | |
| 169 | pub fn effective_read_timeout(&self, global: &McpTimeouts) -> u64 { |
| 170 | self.read_timeout.unwrap_or(global.read_timeout) |
| 171 | } |
| 172 | |
| 173 | pub fn is_enabled(&self) -> bool { |
| 174 | self.enabled && !self.disabled |
| 175 | } |
| 176 | |
| 177 | pub fn is_tool_enabled(&self, tool_name: &str) -> bool { |
| 178 | let allowed = if self.enabled_tools.is_empty() { |
| 179 | true |
| 180 | } else { |
| 181 | self.enabled_tools.iter().any(|t| t == tool_name) |
| 182 | }; |
| 183 | if !allowed { |
| 184 | return false; |
| 185 | } |
| 186 | !self.disabled_tools.iter().any(|t| t == tool_name) |
| 187 | } |
| 188 | } |
| 189 | |
| 190 | // === MCP Tool Definition === |
| 191 | |
| 192 | /// Tool discovered from an MCP server |
| 193 | #[derive(Debug, Clone, Deserialize, Serialize)] |
| 194 | pub struct McpTool { |
| 195 | pub name: String, |
| 196 | #[serde(default)] |
| 197 | pub description: Option<String>, |
| 198 | #[serde(rename = "inputSchema", default)] |
| 199 | pub input_schema: serde_json::Value, |
| 200 | } |
| 201 | |
| 202 | /// Resource discovered from an MCP server |
| 203 | #[derive(Debug, Clone, Deserialize, Serialize)] |
| 204 | pub struct McpResource { |
| 205 | pub uri: String, |
| 206 | pub name: String, |
| 207 | #[serde(default)] |
| 208 | pub description: Option<String>, |
| 209 | #[serde(rename = "mimeType", default)] |
| 210 | pub mime_type: Option<String>, |
| 211 | } |
| 212 | |
| 213 | /// Resource template discovered from an MCP server |
| 214 | #[derive(Debug, Clone, Deserialize, Serialize)] |
| 215 | pub struct McpResourceTemplate { |
| 216 | #[serde(rename = "uriTemplate")] |
| 217 | pub uri_template: String, |
| 218 | pub name: String, |
| 219 | #[serde(default)] |
| 220 | pub description: Option<String>, |
| 221 | #[serde(rename = "mimeType", default)] |
| 222 | pub mime_type: Option<String>, |
| 223 | } |
| 224 | |
| 225 | /// Prompt discovered from an MCP server |
| 226 | #[derive(Debug, Clone, Deserialize, Serialize)] |
| 227 | pub struct McpPrompt { |
| 228 | pub name: String, |
| 229 | #[serde(default)] |
| 230 | pub description: Option<String>, |
| 231 | #[serde(default)] |
| 232 | pub arguments: Vec<McpPromptArgument>, |
| 233 | } |
| 234 | |
| 235 | /// Argument for an MCP prompt |
| 236 | #[derive(Debug, Clone, Deserialize, Serialize)] |
| 237 | pub struct McpPromptArgument { |
| 238 | pub name: String, |
| 239 | #[serde(default)] |
| 240 | pub description: Option<String>, |
| 241 | #[serde(default)] |
| 242 | pub required: bool, |
| 243 | } |
| 244 | |
| 245 | // === Connection State === |
| 246 | |
| 247 | /// State of an MCP connection |
| 248 | #[derive(Debug, Clone, Copy, PartialEq, Eq)] |
| 249 | pub enum ConnectionState { |
| 250 | Connecting, |
| 251 | Ready, |
| 252 | Disconnected, |
| 253 | } |
| 254 | |
| 255 | // === McpConnection - Async Connection Management === |
| 256 | |
| 257 | // === Transport Trait === |
| 258 | |
| 259 | #[async_trait::async_trait] |
| 260 | pub trait McpTransport: Send + Sync { |
| 261 | async fn send(&mut self, msg: serde_json::Value) -> Result<()>; |
| 262 | async fn recv(&mut self) -> Result<serde_json::Value>; |
| 263 | |
| 264 | /// Graceful shutdown — stdio transports send SIGTERM to the child and |
| 265 | /// give it a brief window to exit before tokio's `kill_on_drop` fires |
| 266 | /// SIGKILL as the backstop. Default is a no-op for non-stdio transports |
| 267 | /// that have no child process. Whalescale#420. |
| 268 | async fn shutdown(&mut self) {} |
| 269 | } |
| 270 | |
| 271 | pub struct StdioTransport { |
| 272 | child: Child, |
| 273 | stdin: ChildStdin, |
| 274 | reader: tokio::io::BufReader<ChildStdout>, |
| 275 | } |
| 276 | |
| 277 | /// How long `StdioTransport::shutdown` waits for the child to exit on SIGTERM |
| 278 | /// before `kill_on_drop` fires SIGKILL. Tuned short so a hung MCP server |
| 279 | /// can't stall TUI exit; well-behaved servers almost always exit within |
| 280 | /// a few hundred ms. |
| 281 | const STDIO_SHUTDOWN_GRACE: Duration = Duration::from_millis(2_000); |
| 282 | |
| 283 | /// Best-effort SIGTERM. On Unix uses `libc::kill`; on Windows there's no |
| 284 | /// equivalent so we let `kill_on_drop` (TerminateProcess) handle it via the |
| 285 | /// subsequent Drop. Returns whether a signal was actually sent. |
| 286 | fn send_sigterm(child: &Child) -> bool { |
| 287 | #[cfg(unix)] |
| 288 | { |
| 289 | if let Some(pid) = child.id() { |
| 290 | // SAFETY: pid was just obtained from `child.id()`. `libc::kill` |
| 291 | // with `SIGTERM` is async-signal-safe and never observes invalid |
| 292 | // memory. Worst case (pid wrap / process already gone) returns |
| 293 | // ESRCH, which we deliberately ignore. |
| 294 | unsafe { |
| 295 | let _ = libc::kill(pid as i32, libc::SIGTERM); |
| 296 | } |
| 297 | return true; |
| 298 | } |
| 299 | false |
| 300 | } |
| 301 | #[cfg(not(unix))] |
| 302 | { |
| 303 | let _ = child; |
| 304 | false |
| 305 | } |
| 306 | } |
| 307 | |
| 308 | #[async_trait::async_trait] |
| 309 | impl McpTransport for StdioTransport { |
| 310 | async fn send(&mut self, msg: serde_json::Value) -> Result<()> { |
| 311 | let line = serde_json::to_string(&msg)? + "\n"; |
| 312 | self.stdin.write_all(line.as_bytes()).await?; |
| 313 | self.stdin.flush().await?; |
| 314 | Ok(()) |
| 315 | } |
| 316 | |
| 317 | async fn recv(&mut self) -> Result<serde_json::Value> { |
| 318 | let mut line = String::new(); |
| 319 | loop { |
| 320 | line.clear(); |
| 321 | let bytes = self.reader.read_line(&mut line).await?; |
| 322 | if bytes == 0 { |
| 323 | anyhow::bail!("Stdio transport closed"); |
| 324 | } |
| 325 | |
| 326 | let trimmed = line.trim(); |
| 327 | if trimmed.is_empty() { |
| 328 | continue; |
| 329 | } |
| 330 | |
| 331 | if let Ok(value) = serde_json::from_str::<serde_json::Value>(trimmed) { |
| 332 | return Ok(value); |
| 333 | } |
| 334 | } |
| 335 | } |
| 336 | |
| 337 | /// Send SIGTERM and wait up to `STDIO_SHUTDOWN_GRACE` for graceful exit |
| 338 | /// before letting Drop / `kill_on_drop` fire SIGKILL as the backstop. |
| 339 | async fn shutdown(&mut self) { |
| 340 | send_sigterm(&self.child); |
| 341 | // Give the child a window to exit cleanly. Discard the result — |
| 342 | // either it exits (success) or the timeout fires (Drop will SIGKILL). |
| 343 | let _ = tokio::time::timeout(STDIO_SHUTDOWN_GRACE, self.child.wait()).await; |
| 344 | } |
| 345 | } |
| 346 | |
| 347 | /// Drop fallback (#420): if `shutdown` was never called explicitly, still |
| 348 | /// fire SIGTERM before tokio's `kill_on_drop` sends SIGKILL. The two |
| 349 | /// signals arrive back-to-back so well-behaved servers at least see the |
| 350 | /// SIGTERM first; misbehaving ones get SIGKILL'd anyway. |
| 351 | impl Drop for StdioTransport { |
| 352 | fn drop(&mut self) { |
| 353 | send_sigterm(&self.child); |
| 354 | } |
| 355 | } |
| 356 | |
| 357 | pub struct SseTransport { |
| 358 | client: reqwest::Client, |
| 359 | base_url: String, |
| 360 | endpoint_url: Option<String>, |
| 361 | receiver: tokio::sync::mpsc::UnboundedReceiver<serde_json::Value>, |
| 362 | } |
| 363 | |
| 364 | impl SseTransport { |
| 365 | pub async fn connect( |
| 366 | client: reqwest::Client, |
| 367 | url: String, |
| 368 | cancel_token: tokio_util::sync::CancellationToken, |
| 369 | ) -> Result<Self> { |
| 370 | let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); |
| 371 | let client_clone = client.clone(); |
| 372 | let url_clone = url.clone(); |
| 373 | |
| 374 | tokio::spawn(async move { |
| 375 | if cancel_token.is_cancelled() { |
| 376 | return; |
| 377 | } |
| 378 | use futures_util::FutureExt; |
| 379 | let result = std::panic::AssertUnwindSafe(Self::run_sse_loop( |
| 380 | client_clone, |
| 381 | url_clone, |
| 382 | tx, |
| 383 | cancel_token, |
| 384 | )) |
| 385 | .catch_unwind() |
| 386 | .await; |
| 387 | match result { |
| 388 | Ok(res) => { |
| 389 | if let Err(e) = res { |
| 390 | tracing::error!("SSE loop error: {}", e); |
| 391 | } |
| 392 | } |
| 393 | Err(panic_err) => { |
| 394 | if let Some(msg) = panic_err.downcast_ref::<&str>() { |
| 395 | tracing::error!("SSE loop panicked: {}", msg); |
| 396 | } else if let Some(msg) = panic_err.downcast_ref::<String>() { |
| 397 | tracing::error!("SSE loop panicked: {}", msg); |
| 398 | } else { |
| 399 | tracing::error!("SSE loop panicked with unknown error"); |
| 400 | } |
| 401 | } |
| 402 | } |
| 403 | }); |
| 404 | |
| 405 | Ok(Self { |
| 406 | client, |
| 407 | base_url: url, |
| 408 | endpoint_url: None, |
| 409 | receiver: rx, |
| 410 | }) |
| 411 | } |
| 412 | |
| 413 | async fn run_sse_loop( |
| 414 | client: reqwest::Client, |
| 415 | url: String, |
| 416 | tx: tokio::sync::mpsc::UnboundedSender<serde_json::Value>, |
| 417 | cancel_token: tokio_util::sync::CancellationToken, |
| 418 | ) -> Result<()> { |
| 419 | let response = client.get(&url).send().await.with_context(|| { |
| 420 | format!( |
| 421 | "MCP SSE connect failed (transport=http url={})", |
| 422 | mask_url_secrets(&url), |
| 423 | ) |
| 424 | })?; |
| 425 | let status = response.status(); |
| 426 | if !status.is_success() { |
| 427 | let body_excerpt = bounded_body_excerpt(response, ERROR_BODY_PREVIEW_BYTES).await; |
| 428 | anyhow::bail!( |
| 429 | "MCP SSE rejected (transport=http url={} status={}): {}", |
| 430 | mask_url_secrets(&url), |
| 431 | status, |
| 432 | body_excerpt, |
| 433 | ); |
| 434 | } |
| 435 | |
| 436 | let mut stream = response.bytes_stream(); |
| 437 | use futures_util::StreamExt; |
| 438 | let mut buffer = String::new(); |
| 439 | |
| 440 | loop { |
| 441 | if cancel_token.is_cancelled() { |
| 442 | tracing::debug!("SSE loop cancelled"); |
| 443 | break; |
| 444 | } |
| 445 | let item = tokio::select! { |
| 446 | _ = cancel_token.cancelled() => { |
| 447 | tracing::debug!("SSE loop shutting down"); |
| 448 | break; |
| 449 | } |
| 450 | item = stream.next() => { |
| 451 | match item { |
| 452 | Some(i) => i, |
| 453 | None => break, |
| 454 | } |
| 455 | } |
| 456 | }; |
| 457 | let chunk = item?; |
| 458 | let s = String::from_utf8_lossy(&chunk); |
| 459 | buffer.push_str(&s); |
| 460 | |
| 461 | while let Some(pos) = buffer.find("\n\n") { |
| 462 | let event_block = buffer[..pos].to_string(); |
| 463 | buffer = buffer[pos + 2..].to_string(); |
| 464 | |
| 465 | let mut event_type = "message"; |
| 466 | let mut data = String::new(); |
| 467 | |
| 468 | for line in event_block.lines() { |
| 469 | if let Some(stripped) = line.strip_prefix("event: ") { |
| 470 | event_type = stripped; |
| 471 | } else if let Some(stripped) = line.strip_prefix("data: ") { |
| 472 | data.push_str(stripped); |
| 473 | } |
| 474 | } |
| 475 | |
| 476 | match event_type { |
| 477 | "endpoint" => { |
| 478 | // Special internal message to set endpoint |
| 479 | let _ = tx.send(serde_json::json!({ |
| 480 | "__internal_sse_endpoint__": data |
| 481 | })); |
| 482 | } |
| 483 | "message" => { |
| 484 | if let Ok(val) = serde_json::from_str::<serde_json::Value>(&data) { |
| 485 | let _ = tx.send(val); |
| 486 | } |
| 487 | } |
| 488 | _ => {} |
| 489 | } |
| 490 | } |
| 491 | } |
| 492 | Ok(()) |
| 493 | } |
| 494 | } |
| 495 | |
| 496 | #[async_trait::async_trait] |
| 497 | impl McpTransport for SseTransport { |
| 498 | async fn send(&mut self, msg: serde_json::Value) -> Result<()> { |
| 499 | let endpoint = self |
| 500 | .endpoint_url |
| 501 | .as_ref() |
| 502 | .context("SSE endpoint not yet discovered")?; |
| 503 | let response = self.client.post(endpoint).json(&msg).send().await?; |
| 504 | if !response.status().is_success() { |
| 505 | anyhow::bail!("Failed to send message via SSE POST: {}", response.status()); |
| 506 | } |
| 507 | Ok(()) |
| 508 | } |
| 509 | |
| 510 | async fn recv(&mut self) -> Result<serde_json::Value> { |
| 511 | loop { |
| 512 | let msg = self.receiver.recv().await.context("SSE transport closed")?; |
| 513 | if let Some(endpoint) = msg.get("__internal_sse_endpoint__") { |
| 514 | let url_str = endpoint.as_str().context("Invalid endpoint format")?; |
| 515 | // Handle relative vs absolute URLs |
| 516 | if url_str.starts_with("http") { |
| 517 | self.endpoint_url = Some(url_str.to_string()); |
| 518 | } else { |
| 519 | let base = reqwest::Url::parse(&self.base_url)?; |
| 520 | let joined = base.join(url_str)?; |
| 521 | self.endpoint_url = Some(joined.to_string()); |
| 522 | } |
| 523 | continue; |
| 524 | } |
| 525 | return Ok(msg); |
| 526 | } |
| 527 | } |
| 528 | } |
| 529 | |
| 530 | // === McpConnection - Async Connection Management === |
| 531 | |
| 532 | /// Manages a single async connection to an MCP server |
| 533 | pub struct McpConnection { |
| 534 | name: String, |
| 535 | transport: Box<dyn McpTransport>, |
| 536 | tools: Vec<McpTool>, |
| 537 | resources: Vec<McpResource>, |
| 538 | resource_templates: Vec<McpResourceTemplate>, |
| 539 | prompts: Vec<McpPrompt>, |
| 540 | request_id: AtomicU64, |
| 541 | state: ConnectionState, |
| 542 | config: McpServerConfig, |
| 543 | cancel_token: tokio_util::sync::CancellationToken, |
| 544 | } |
| 545 | |
| 546 | impl McpConnection { |
| 547 | /// Connect to an MCP server and initialize it. |
| 548 | /// |
| 549 | /// `network_policy` (added in v0.7.0 for #135) is consulted for HTTP/SSE |
| 550 | /// transports only — STDIO transports are unaffected. Pass `None` to |
| 551 | /// match pre-v0.7.0 permissive behavior. |
| 552 | pub async fn connect_with_policy( |
| 553 | name: String, |
| 554 | config: McpServerConfig, |
| 555 | global_timeouts: &McpTimeouts, |
| 556 | network_policy: Option<&NetworkPolicyDecider>, |
| 557 | ) -> Result<Self> { |
| 558 | let connect_timeout_secs = config.effective_connect_timeout(global_timeouts); |
| 559 | let cancel_token = tokio_util::sync::CancellationToken::new(); |
| 560 | |
| 561 | let transport: Box<dyn McpTransport> = if let Some(url) = &config.url { |
| 562 | // Per-domain network policy gate (#135). Only the HTTP/SSE transport |
| 563 | // is gated; STDIO MCP servers run as local subprocesses and never |
| 564 | // touch the network from this code path. |
| 565 | if let Some(decider) = network_policy |
| 566 | && let Some(host) = host_from_url(url) |
| 567 | { |
| 568 | match decider.evaluate(&host, "mcp") { |
| 569 | Decision::Allow => {} |
| 570 | Decision::Deny => { |
| 571 | anyhow::bail!( |
| 572 | "MCP server '{name}' connection to '{host}' blocked by network policy" |
| 573 | ); |
| 574 | } |
| 575 | Decision::Prompt => { |
| 576 | anyhow::bail!( |
| 577 | "MCP server '{name}' connection to '{host}' requires approval; \ |
| 578 | re-run after `/network allow {host}` or set network.default = \"allow\" in config" |
| 579 | ); |
| 580 | } |
| 581 | } |
| 582 | } |
| 583 | let client = reqwest::Client::builder() |
| 584 | .timeout(Duration::from_secs(connect_timeout_secs)) |
| 585 | .build()?; |
| 586 | Box::new(SseTransport::connect(client, url.clone(), cancel_token.clone()).await?) |
| 587 | } else if let Some(command) = &config.command { |
| 588 | let mut cmd = tokio::process::Command::new(command); |
| 589 | cmd.args(&config.args) |
| 590 | .stdin(std::process::Stdio::piped()) |
| 591 | .stdout(std::process::Stdio::piped()) |
| 592 | .stderr(std::process::Stdio::null()) |
| 593 | .kill_on_drop(true); |
| 594 | |
| 595 | for (key, value) in &config.env { |
| 596 | cmd.env(key, value); |
| 597 | } |
| 598 | |
| 599 | let mut child = cmd.spawn().with_context(|| { |
| 600 | let env_keys: Vec<&str> = config.env.keys().map(String::as_str).collect(); |
| 601 | format!( |
| 602 | "MCP stdio spawn failed (transport=stdio server={name} cmd={command:?} args={:?} env_keys={env_keys:?})", |
| 603 | config.args, |
| 604 | ) |
| 605 | })?; |
| 606 | |
| 607 | let stdin = child.stdin.take().context("Failed to get MCP stdin")?; |
| 608 | let stdout = child.stdout.take().context("Failed to get MCP stdout")?; |
| 609 | |
| 610 | Box::new(StdioTransport { |
| 611 | child, |
| 612 | stdin, |
| 613 | reader: tokio::io::BufReader::new(stdout), |
| 614 | }) |
| 615 | } else { |
| 616 | anyhow::bail!( |
| 617 | "MCP server '{}' config must have either 'command' or 'url'", |
| 618 | name |
| 619 | ); |
| 620 | }; |
| 621 | |
| 622 | let mut conn = Self { |
| 623 | name: name.clone(), |
| 624 | transport, |
| 625 | tools: Vec::new(), |
| 626 | resources: Vec::new(), |
| 627 | resource_templates: Vec::new(), |
| 628 | prompts: Vec::new(), |
| 629 | request_id: AtomicU64::new(1), |
| 630 | state: ConnectionState::Connecting, |
| 631 | config, |
| 632 | cancel_token, |
| 633 | }; |
| 634 | |
| 635 | // Initialize with timeout |
| 636 | tokio::time::timeout(Duration::from_secs(connect_timeout_secs), conn.initialize()) |
| 637 | .await |
| 638 | .with_context(|| format!("MCP server '{name}' initialization timed out"))??; |
| 639 | |
| 640 | // Discover tools, resources, and prompts with timeout |
| 641 | tokio::time::timeout( |
| 642 | Duration::from_secs(connect_timeout_secs), |
| 643 | conn.discover_all(), |
| 644 | ) |
| 645 | .await |
| 646 | .with_context(|| format!("MCP server '{name}' discovery timed out"))??; |
| 647 | |
| 648 | conn.state = ConnectionState::Ready; |
| 649 | Ok(conn) |
| 650 | } |
| 651 | |
| 652 | /// Send initialize request and wait for response |
| 653 | async fn initialize(&mut self) -> Result<()> { |
| 654 | let init_id = self.next_id(); |
| 655 | self.send(serde_json::json!({ |
| 656 | "jsonrpc": "2.0", |
| 657 | "id": init_id, |
| 658 | "method": "initialize", |
| 659 | "params": { |
| 660 | "protocolVersion": "2024-11-05", |
| 661 | "clientInfo": { |
| 662 | "name": "deepseek-tui", |
| 663 | "version": env!("CARGO_PKG_VERSION") |
| 664 | }, |
| 665 | "capabilities": { |
| 666 | "tools": {}, |
| 667 | "resources": {}, |
| 668 | "prompts": {} |
| 669 | } |
| 670 | } |
| 671 | })) |
| 672 | .await?; |
| 673 | |
| 674 | self.recv(init_id).await?; |
| 675 | |
| 676 | // Send initialized notification (no id, no response expected) |
| 677 | self.send(serde_json::json!({ |
| 678 | "jsonrpc": "2.0", |
| 679 | "method": "notifications/initialized" |
| 680 | })) |
| 681 | .await?; |
| 682 | |
| 683 | Ok(()) |
| 684 | } |
| 685 | |
| 686 | /// Discover tools, resources, and prompts |
| 687 | async fn discover_all(&mut self) -> Result<()> { |
| 688 | // We use join! to discover everything concurrently if possible, |
| 689 | // but for now let's keep it sequential for simplicity in error handling |
| 690 | self.discover_tools().await?; |
| 691 | self.discover_resources().await?; |
| 692 | self.discover_resource_templates().await?; |
| 693 | self.discover_prompts().await?; |
| 694 | Ok(()) |
| 695 | } |
| 696 | |
| 697 | /// Discover available tools from the MCP server |
| 698 | async fn discover_tools(&mut self) -> Result<()> { |
| 699 | let list_id = self.next_id(); |
| 700 | self.send(serde_json::json!({ |
| 701 | "jsonrpc": "2.0", |
| 702 | "id": list_id, |
| 703 | "method": "tools/list", |
| 704 | "params": {} |
| 705 | })) |
| 706 | .await?; |
| 707 | |
| 708 | let response = self.recv(list_id).await?; |
| 709 | |
| 710 | if let Some(result) = response.get("result") |
| 711 | && let Some(tools) = result.get("tools") |
| 712 | { |
| 713 | self.tools = serde_json::from_value(tools.clone()).unwrap_or_default(); |
| 714 | } |
| 715 | |
| 716 | Ok(()) |
| 717 | } |
| 718 | |
| 719 | /// Discover available resources from the MCP server |
| 720 | async fn discover_resources(&mut self) -> Result<()> { |
| 721 | let list_id = self.next_id(); |
| 722 | self.send(serde_json::json!({ |
| 723 | "jsonrpc": "2.0", |
| 724 | "id": list_id, |
| 725 | "method": "resources/list", |
| 726 | "params": {} |
| 727 | })) |
| 728 | .await?; |
| 729 | |
| 730 | let response = self.recv(list_id).await?; |
| 731 | |
| 732 | if let Some(result) = response.get("result") |
| 733 | && let Some(resources) = result.get("resources") |
| 734 | { |
| 735 | self.resources = serde_json::from_value(resources.clone()).unwrap_or_default(); |
| 736 | } |
| 737 | |
| 738 | Ok(()) |
| 739 | } |
| 740 | |
| 741 | /// Discover available resource templates from the MCP server |
| 742 | async fn discover_resource_templates(&mut self) -> Result<()> { |
| 743 | let list_id = self.next_id(); |
| 744 | self.send(serde_json::json!({ |
| 745 | "jsonrpc": "2.0", |
| 746 | "id": list_id, |
| 747 | "method": "resources/templates/list", |
| 748 | "params": {} |
| 749 | })) |
| 750 | .await?; |
| 751 | |
| 752 | let response = self.recv(list_id).await?; |
| 753 | |
| 754 | if let Some(result) = response.get("result") { |
| 755 | let templates = result |
| 756 | .get("resourceTemplates") |
| 757 | .or_else(|| result.get("templates")) |
| 758 | .or_else(|| result.get("resource_templates")); |
| 759 | if let Some(templates) = templates { |
| 760 | self.resource_templates = |
| 761 | serde_json::from_value(templates.clone()).unwrap_or_default(); |
| 762 | } |
| 763 | } |
| 764 | |
| 765 | Ok(()) |
| 766 | } |
| 767 | |
| 768 | /// Discover available prompts from the MCP server |
| 769 | async fn discover_prompts(&mut self) -> Result<()> { |
| 770 | let list_id = self.next_id(); |
| 771 | self.send(serde_json::json!({ |
| 772 | "jsonrpc": "2.0", |
| 773 | "id": list_id, |
| 774 | "method": "prompts/list", |
| 775 | "params": {} |
| 776 | })) |
| 777 | .await?; |
| 778 | |
| 779 | let response = self.recv(list_id).await?; |
| 780 | |
| 781 | if let Some(result) = response.get("result") |
| 782 | && let Some(prompts) = result.get("prompts") |
| 783 | { |
| 784 | self.prompts = serde_json::from_value(prompts.clone()).unwrap_or_default(); |
| 785 | } |
| 786 | |
| 787 | Ok(()) |
| 788 | } |
| 789 | |
| 790 | /// Call a tool on this MCP server |
| 791 | pub async fn call_tool( |
| 792 | &mut self, |
| 793 | tool_name: &str, |
| 794 | arguments: serde_json::Value, |
| 795 | timeout_secs: u64, |
| 796 | ) -> Result<serde_json::Value> { |
| 797 | self.call_method( |
| 798 | "tools/call", |
| 799 | serde_json::json!({ |
| 800 | "name": tool_name, |
| 801 | "arguments": arguments |
| 802 | }), |
| 803 | timeout_secs, |
| 804 | ) |
| 805 | .await |
| 806 | } |
| 807 | |
| 808 | /// Read a resource from this MCP server |
| 809 | pub async fn read_resource( |
| 810 | &mut self, |
| 811 | uri: &str, |
| 812 | timeout_secs: u64, |
| 813 | ) -> Result<serde_json::Value> { |
| 814 | self.call_method( |
| 815 | "resources/read", |
| 816 | serde_json::json!({ |
| 817 | "uri": uri |
| 818 | }), |
| 819 | timeout_secs, |
| 820 | ) |
| 821 | .await |
| 822 | } |
| 823 | |
| 824 | /// Get a prompt from this MCP server |
| 825 | pub async fn get_prompt( |
| 826 | &mut self, |
| 827 | prompt_name: &str, |
| 828 | arguments: serde_json::Value, |
| 829 | timeout_secs: u64, |
| 830 | ) -> Result<serde_json::Value> { |
| 831 | self.call_method( |
| 832 | "prompts/get", |
| 833 | serde_json::json!({ |
| 834 | "name": prompt_name, |
| 835 | "arguments": arguments |
| 836 | }), |
| 837 | timeout_secs, |
| 838 | ) |
| 839 | .await |
| 840 | } |
| 841 | |
| 842 | /// Generic method to call an MCP method |
| 843 | async fn call_method( |
| 844 | &mut self, |
| 845 | method: &str, |
| 846 | params: serde_json::Value, |
| 847 | timeout_secs: u64, |
| 848 | ) -> Result<serde_json::Value> { |
| 849 | if self.state != ConnectionState::Ready { |
| 850 | anyhow::bail!( |
| 851 | "Failed to call MCP method '{}': connection '{}' is not ready", |
| 852 | method, |
| 853 | self.name |
| 854 | ); |
| 855 | } |
| 856 | |
| 857 | let call_id = self.next_id(); |
| 858 | self.send(serde_json::json!({ |
| 859 | "jsonrpc": "2.0", |
| 860 | "id": call_id, |
| 861 | "method": method, |
| 862 | "params": params |
| 863 | })) |
| 864 | .await?; |
| 865 | |
| 866 | let response = tokio::time::timeout(Duration::from_secs(timeout_secs), self.recv(call_id)) |
| 867 | .await |
| 868 | .with_context(|| { |
| 869 | format!( |
| 870 | "MCP method '{}' on server '{}' timed out after {}s", |
| 871 | method, self.name, timeout_secs |
| 872 | ) |
| 873 | })??; |
| 874 | |
| 875 | if let Some(error) = response.get("error") { |
| 876 | return Err(anyhow::anyhow!( |
| 877 | "MCP error in '{}': {}", |
| 878 | method, |
| 879 | serde_json::to_string_pretty(error)? |
| 880 | )); |
| 881 | } |
| 882 | |
| 883 | Ok(response |
| 884 | .get("result") |
| 885 | .cloned() |
| 886 | .unwrap_or(serde_json::json!(null))) |
| 887 | } |
| 888 | |
| 889 | /// Get discovered tools |
| 890 | pub fn tools(&self) -> &[McpTool] { |
| 891 | &self.tools |
| 892 | } |
| 893 | |
| 894 | /// Get discovered resources |
| 895 | pub fn resources(&self) -> &[McpResource] { |
| 896 | &self.resources |
| 897 | } |
| 898 | |
| 899 | /// Get discovered resource templates |
| 900 | pub fn resource_templates(&self) -> &[McpResourceTemplate] { |
| 901 | &self.resource_templates |
| 902 | } |
| 903 | |
| 904 | /// Get discovered prompts |
| 905 | pub fn prompts(&self) -> &[McpPrompt] { |
| 906 | &self.prompts |
| 907 | } |
| 908 | |
| 909 | /// Get server name |
| 910 | #[allow(dead_code)] // Public API for MCP consumers |
| 911 | pub fn name(&self) -> &str { |
| 912 | &self.name |
| 913 | } |
| 914 | |
| 915 | /// Check if connection is ready |
| 916 | pub fn is_ready(&self) -> bool { |
| 917 | self.state == ConnectionState::Ready |
| 918 | } |
| 919 | |
| 920 | /// Get server config |
| 921 | pub fn config(&self) -> &McpServerConfig { |
| 922 | &self.config |
| 923 | } |
| 924 | |
| 925 | /// Get connection state |
| 926 | #[allow(dead_code)] // Public API for MCP consumers |
| 927 | pub fn state(&self) -> ConnectionState { |
| 928 | self.state |
| 929 | } |
| 930 | |
| 931 | fn next_id(&self) -> u64 { |
| 932 | self.request_id.fetch_add(1, Ordering::SeqCst) |
| 933 | } |
| 934 | |
| 935 | async fn send(&mut self, msg: serde_json::Value) -> Result<()> { |
| 936 | self.transport.send(msg).await |
| 937 | } |
| 938 | |
| 939 | async fn recv(&mut self, expected_id: u64) -> Result<serde_json::Value> { |
| 940 | loop { |
| 941 | let value = self.transport.recv().await.inspect_err(|_e| { |
| 942 | self.state = ConnectionState::Disconnected; |
| 943 | })?; |
| 944 | |
| 945 | // Check if this is a response with the expected id |
| 946 | if value.get("id").and_then(serde_json::Value::as_u64) == Some(expected_id) { |
| 947 | return Ok(value); |
| 948 | } |
| 949 | // Skip notifications (no id) and responses with different ids |
| 950 | } |
| 951 | } |
| 952 | |
| 953 | /// Gracefully close the connection |
| 954 | #[allow(dead_code)] // Public API for MCP consumers |
| 955 | pub fn close(&mut self) { |
| 956 | self.cancel_token.cancel(); |
| 957 | self.state = ConnectionState::Disconnected; |
| 958 | } |
| 959 | } |
| 960 | |
| 961 | impl Drop for McpConnection { |
| 962 | fn drop(&mut self) { |
| 963 | self.cancel_token.cancel(); |
| 964 | } |
| 965 | } |
| 966 | |
| 967 | // === McpPool - Connection Pool Management === |
| 968 | |
| 969 | /// Pool of MCP connections for reuse |
| 970 | pub struct McpPool { |
| 971 | connections: HashMap<String, McpConnection>, |
| 972 | config: McpConfig, |
| 973 | network_policy: Option<NetworkPolicyDecider>, |
| 974 | } |
| 975 | |
| 976 | impl McpPool { |
| 977 | /// Create a new pool with the given configuration |
| 978 | pub fn new(config: McpConfig) -> Self { |
| 979 | Self { |
| 980 | connections: HashMap::new(), |
| 981 | config, |
| 982 | network_policy: None, |
| 983 | } |
| 984 | } |
| 985 | |
| 986 | /// Create a pool from a configuration file path |
| 987 | pub fn from_config_path(path: &std::path::Path) -> Result<Self> { |
| 988 | let config = if path.exists() { |
| 989 | let contents = fs::read_to_string(path) |
| 990 | .with_context(|| format!("Failed to read MCP config: {}", path.display()))?; |
| 991 | serde_json::from_str(&contents) |
| 992 | .with_context(|| format!("Failed to parse MCP config: {}", path.display()))? |
| 993 | } else { |
| 994 | McpConfig::default() |
| 995 | }; |
| 996 | Ok(Self::new(config)) |
| 997 | } |
| 998 | |
| 999 | /// Attach a per-domain network policy (#135). When set, HTTP/SSE |
| 1000 | /// transports are gated through it; STDIO transports are unaffected. |
| 1001 | pub fn with_network_policy(mut self, policy: NetworkPolicyDecider) -> Self { |
| 1002 | self.network_policy = Some(policy); |
| 1003 | self |
| 1004 | } |
| 1005 | |
| 1006 | /// Get or create a connection to a server |
| 1007 | pub async fn get_or_connect(&mut self, server_name: &str) -> Result<&mut McpConnection> { |
| 1008 | let is_ready = self |
| 1009 | .connections |
| 1010 | .get(server_name) |
| 1011 | .map(|conn| conn.is_ready()) |
| 1012 | .unwrap_or(false); |
| 1013 | if is_ready { |
| 1014 | return self |
| 1015 | .connections |
| 1016 | .get_mut(server_name) |
| 1017 | .ok_or_else(|| anyhow::anyhow!("MCP connection disappeared for {server_name}")); |
| 1018 | } |
| 1019 | |
| 1020 | self.connections.remove(server_name); |
| 1021 | |
| 1022 | let server_config = self |
| 1023 | .config |
| 1024 | .servers |
| 1025 | .get(server_name) |
| 1026 | .ok_or_else(|| anyhow::anyhow!("Failed to find MCP server: {server_name}"))? |
| 1027 | .clone(); |
| 1028 | |
| 1029 | if !server_config.is_enabled() { |
| 1030 | anyhow::bail!("Failed to connect MCP server '{server_name}': server is disabled"); |
| 1031 | } |
| 1032 | |
| 1033 | let connection = McpConnection::connect_with_policy( |
| 1034 | server_name.to_string(), |
| 1035 | server_config, |
| 1036 | &self.config.timeouts, |
| 1037 | self.network_policy.as_ref(), |
| 1038 | ) |
| 1039 | .await?; |
| 1040 | |
| 1041 | self.connections.insert(server_name.to_string(), connection); |
| 1042 | self.connections |
| 1043 | .get_mut(server_name) |
| 1044 | .ok_or_else(|| anyhow::anyhow!("Failed to store MCP connection for {server_name}")) |
| 1045 | } |
| 1046 | |
| 1047 | /// Connect to all enabled servers, returning errors for failed connections |
| 1048 | pub async fn connect_all(&mut self) -> Vec<(String, anyhow::Error)> { |
| 1049 | let mut errors = Vec::new(); |
| 1050 | let names: Vec<String> = self |
| 1051 | .config |
| 1052 | .servers |
| 1053 | .keys() |
| 1054 | .filter(|n| self.config.servers[*n].is_enabled()) |
| 1055 | .cloned() |
| 1056 | .collect(); |
| 1057 | |
| 1058 | for name in names { |
| 1059 | if let Err(e) = self.get_or_connect(&name).await { |
| 1060 | errors.push((name, e)); |
| 1061 | } |
| 1062 | } |
| 1063 | |
| 1064 | for (name, server_cfg) in &self.config.servers { |
| 1065 | if server_cfg.required |
| 1066 | && server_cfg.is_enabled() |
| 1067 | && !self |
| 1068 | .connections |
| 1069 | .get(name) |
| 1070 | .is_some_and(McpConnection::is_ready) |
| 1071 | { |
| 1072 | errors.push(( |
| 1073 | name.clone(), |
| 1074 | anyhow::anyhow!("required MCP server failed to initialize"), |
| 1075 | )); |
| 1076 | } |
| 1077 | } |
| 1078 | |
| 1079 | errors |
| 1080 | } |
| 1081 | |
| 1082 | /// Get all discovered tools with server-prefixed names |
| 1083 | pub fn all_tools(&self) -> Vec<(String, &McpTool)> { |
| 1084 | let mut tools = Vec::new(); |
| 1085 | for (server, conn) in &self.connections { |
| 1086 | for tool in conn.tools() { |
| 1087 | if !conn.config().is_tool_enabled(&tool.name) { |
| 1088 | continue; |
| 1089 | } |
| 1090 | // Format: mcp_{server}_{tool} |
| 1091 | tools.push((format!("mcp_{}_{}", server, tool.name), tool)); |
| 1092 | } |
| 1093 | } |
| 1094 | tools |
| 1095 | } |
| 1096 | |
| 1097 | /// Get all discovered resources with server-prefixed names |
| 1098 | pub fn all_resources(&self) -> Vec<(String, &McpResource)> { |
| 1099 | let mut resources = Vec::new(); |
| 1100 | for (server, conn) in &self.connections { |
| 1101 | for resource in conn.resources() { |
| 1102 | // Format: mcp_{server}_{resource_name} |
| 1103 | // Note: resource names might contain spaces, we should probably slugify them |
| 1104 | let safe_name = resource.name.replace(' ', "_").to_lowercase(); |
| 1105 | resources.push((format!("mcp_{}_{}", server, safe_name), resource)); |
| 1106 | } |
| 1107 | } |
| 1108 | resources |
| 1109 | } |
| 1110 | |
| 1111 | /// Get all discovered resource templates with server-prefixed names |
| 1112 | #[allow(dead_code)] // Public API for MCP resource discovery |
| 1113 | pub fn all_resource_templates(&self) -> Vec<(String, &McpResourceTemplate)> { |
| 1114 | let mut templates = Vec::new(); |
| 1115 | for (server, conn) in &self.connections { |
| 1116 | for template in conn.resource_templates() { |
| 1117 | let safe_name = template.name.replace(' ', "_").to_lowercase(); |
| 1118 | templates.push((format!("mcp_{}_{}", server, safe_name), template)); |
| 1119 | } |
| 1120 | } |
| 1121 | templates |
| 1122 | } |
| 1123 | |
| 1124 | async fn list_resources(&mut self, server: Option<String>) -> Result<Vec<serde_json::Value>> { |
| 1125 | if let Some(server_name) = server { |
| 1126 | let conn = self.get_or_connect(&server_name).await?; |
| 1127 | let resources = conn |
| 1128 | .resources() |
| 1129 | .iter() |
| 1130 | .map(|resource| { |
| 1131 | serde_json::json!({ |
| 1132 | "server": server_name.clone(), |
| 1133 | "uri": resource.uri, |
| 1134 | "name": resource.name, |
| 1135 | "description": resource.description, |
| 1136 | "mime_type": resource.mime_type, |
| 1137 | }) |
| 1138 | }) |
| 1139 | .collect(); |
| 1140 | return Ok(resources); |
| 1141 | } |
| 1142 | |
| 1143 | let _ = self.connect_all().await; |
| 1144 | let mut items = Vec::new(); |
| 1145 | for (server, conn) in &self.connections { |
| 1146 | for resource in conn.resources() { |
| 1147 | items.push(serde_json::json!({ |
| 1148 | "server": server, |
| 1149 | "uri": resource.uri, |
| 1150 | "name": resource.name, |
| 1151 | "description": resource.description, |
| 1152 | "mime_type": resource.mime_type, |
| 1153 | })); |
| 1154 | } |
| 1155 | } |
| 1156 | Ok(items) |
| 1157 | } |
| 1158 | |
| 1159 | async fn list_resource_templates( |
| 1160 | &mut self, |
| 1161 | server: Option<String>, |
| 1162 | ) -> Result<Vec<serde_json::Value>> { |
| 1163 | if let Some(server_name) = server { |
| 1164 | let conn = self.get_or_connect(&server_name).await?; |
| 1165 | let templates = conn |
| 1166 | .resource_templates() |
| 1167 | .iter() |
| 1168 | .map(|template| { |
| 1169 | serde_json::json!({ |
| 1170 | "server": server_name.clone(), |
| 1171 | "uri_template": template.uri_template, |
| 1172 | "name": template.name, |
| 1173 | "description": template.description, |
| 1174 | "mime_type": template.mime_type, |
| 1175 | }) |
| 1176 | }) |
| 1177 | .collect(); |
| 1178 | return Ok(templates); |
| 1179 | } |
| 1180 | |
| 1181 | let _ = self.connect_all().await; |
| 1182 | let mut items = Vec::new(); |
| 1183 | for (server, conn) in &self.connections { |
| 1184 | for template in conn.resource_templates() { |
| 1185 | items.push(serde_json::json!({ |
| 1186 | "server": server, |
| 1187 | "uri_template": template.uri_template, |
| 1188 | "name": template.name, |
| 1189 | "description": template.description, |
| 1190 | "mime_type": template.mime_type, |
| 1191 | })); |
| 1192 | } |
| 1193 | } |
| 1194 | Ok(items) |
| 1195 | } |
| 1196 | |
| 1197 | /// Get all discovered prompts with server-prefixed names |
| 1198 | pub fn all_prompts(&self) -> Vec<(String, &McpPrompt)> { |
| 1199 | let mut prompts = Vec::new(); |
| 1200 | for (server, conn) in &self.connections { |
| 1201 | for prompt in conn.prompts() { |
| 1202 | // Format: mcp_{server}_{prompt} |
| 1203 | prompts.push((format!("mcp_{}_{}", server, prompt.name), prompt)); |
| 1204 | } |
| 1205 | } |
| 1206 | prompts |
| 1207 | } |
| 1208 | |
| 1209 | /// Read a resource from a specific server |
| 1210 | pub async fn read_resource( |
| 1211 | &mut self, |
| 1212 | server_name: &str, |
| 1213 | uri: &str, |
| 1214 | ) -> Result<serde_json::Value> { |
| 1215 | let global_timeouts = self.config.timeouts; |
| 1216 | let conn = self.get_or_connect(server_name).await?; |
| 1217 | let timeout = conn.config().effective_read_timeout(&global_timeouts); |
| 1218 | conn.read_resource(uri, timeout).await |
| 1219 | } |
| 1220 | |
| 1221 | /// Get a prompt from a specific server |
| 1222 | pub async fn get_prompt( |
| 1223 | &mut self, |
| 1224 | server_name: &str, |
| 1225 | prompt_name: &str, |
| 1226 | arguments: serde_json::Value, |
| 1227 | ) -> Result<serde_json::Value> { |
| 1228 | let global_timeouts = self.config.timeouts; |
| 1229 | let conn = self.get_or_connect(server_name).await?; |
| 1230 | let timeout = conn.config().effective_execute_timeout(&global_timeouts); |
| 1231 | conn.get_prompt(prompt_name, arguments, timeout).await |
| 1232 | } |
| 1233 | |
| 1234 | /// Parse a prefixed name into (server_name, tool_name) |
| 1235 | fn parse_prefixed_name<'a>(&self, prefixed_name: &'a str) -> Result<(&'a str, &'a str)> { |
| 1236 | if !prefixed_name.starts_with("mcp_") { |
| 1237 | anyhow::bail!("Invalid MCP tool name: {}", prefixed_name); |
| 1238 | } |
| 1239 | let rest = &prefixed_name[4..]; |
| 1240 | let Some((server, tool)) = rest.split_once('_') else { |
| 1241 | anyhow::bail!("Invalid MCP tool name format: {}", prefixed_name); |
| 1242 | }; |
| 1243 | Ok((server, tool)) |
| 1244 | } |
| 1245 | |
| 1246 | /// Convert discovered tools to API Tool format |
| 1247 | pub fn to_api_tools(&self) -> Vec<crate::models::Tool> { |
| 1248 | let mut api_tools = Vec::new(); |
| 1249 | |
| 1250 | // Add regular tools |
| 1251 | for (name, tool) in self.all_tools() { |
| 1252 | api_tools.push(crate::models::Tool { |
| 1253 | tool_type: None, |
| 1254 | name, |
| 1255 | description: tool.description.clone().unwrap_or_default(), |
| 1256 | input_schema: tool.input_schema.clone(), |
| 1257 | allowed_callers: Some(vec!["direct".to_string()]), |
| 1258 | defer_loading: Some(false), |
| 1259 | input_examples: None, |
| 1260 | strict: None, |
| 1261 | cache_control: None, |
| 1262 | }); |
| 1263 | } |
| 1264 | |
| 1265 | if !self.config.servers.is_empty() { |
| 1266 | api_tools.push(crate::models::Tool { |
| 1267 | tool_type: None, |
| 1268 | name: "list_mcp_resources".to_string(), |
| 1269 | description: "List available MCP resources across servers (optionally filtered by server).".to_string(), |
| 1270 | input_schema: serde_json::json!({ |
| 1271 | "type": "object", |
| 1272 | "properties": { |
| 1273 | "server": { "type": "string", "description": "Optional MCP server name to filter by" } |
| 1274 | } |
| 1275 | }), |
| 1276 | allowed_callers: Some(vec!["direct".to_string()]), |
| 1277 | defer_loading: Some(false), |
| 1278 | input_examples: None, |
| 1279 | strict: None, |
| 1280 | cache_control: None, |
| 1281 | }); |
| 1282 | api_tools.push(crate::models::Tool { |
| 1283 | tool_type: None, |
| 1284 | name: "list_mcp_resource_templates".to_string(), |
| 1285 | description: "List available MCP resource templates across servers (optionally filtered by server).".to_string(), |
| 1286 | input_schema: serde_json::json!({ |
| 1287 | "type": "object", |
| 1288 | "properties": { |
| 1289 | "server": { "type": "string", "description": "Optional MCP server name to filter by" } |
| 1290 | } |
| 1291 | }), |
| 1292 | allowed_callers: Some(vec!["direct".to_string()]), |
| 1293 | defer_loading: Some(false), |
| 1294 | input_examples: None, |
| 1295 | strict: None, |
| 1296 | cache_control: None, |
| 1297 | }); |
| 1298 | } |
| 1299 | |
| 1300 | // Add resource reading tools if resources exist |
| 1301 | let resources = self.all_resources(); |
| 1302 | if !resources.is_empty() { |
| 1303 | api_tools.push(crate::models::Tool { |
| 1304 | tool_type: None, |
| 1305 | name: "mcp_read_resource".to_string(), |
| 1306 | description: "Read a resource from an MCP server using its URI".to_string(), |
| 1307 | input_schema: serde_json::json!({ |
| 1308 | "type": "object", |
| 1309 | "properties": { |
| 1310 | "server": { "type": "string", "description": "The name of the MCP server" }, |
| 1311 | "uri": { "type": "string", "description": "The URI of the resource to read" } |
| 1312 | }, |
| 1313 | "required": ["server", "uri"] |
| 1314 | }), |
| 1315 | allowed_callers: Some(vec!["direct".to_string()]), |
| 1316 | defer_loading: Some(false), |
| 1317 | input_examples: None, |
| 1318 | strict: None, |
| 1319 | cache_control: None, |
| 1320 | }); |
| 1321 | api_tools.push(crate::models::Tool { |
| 1322 | tool_type: None, |
| 1323 | name: "read_mcp_resource".to_string(), |
| 1324 | description: "Alias for mcp_read_resource.".to_string(), |
| 1325 | input_schema: serde_json::json!({ |
| 1326 | "type": "object", |
| 1327 | "properties": { |
| 1328 | "server": { "type": "string", "description": "The name of the MCP server" }, |
| 1329 | "uri": { "type": "string", "description": "The URI of the resource to read" } |
| 1330 | }, |
| 1331 | "required": ["server", "uri"] |
| 1332 | }), |
| 1333 | allowed_callers: Some(vec!["direct".to_string()]), |
| 1334 | defer_loading: Some(false), |
| 1335 | input_examples: None, |
| 1336 | strict: None, |
| 1337 | cache_control: None, |
| 1338 | }); |
| 1339 | } |
| 1340 | |
| 1341 | // Add prompt getting tools if prompts exist |
| 1342 | let prompts = self.all_prompts(); |
| 1343 | if !prompts.is_empty() { |
| 1344 | api_tools.push(crate::models::Tool { |
| 1345 | tool_type: None, |
| 1346 | name: "mcp_get_prompt".to_string(), |
| 1347 | description: "Get a prompt from an MCP server".to_string(), |
| 1348 | input_schema: serde_json::json!({ |
| 1349 | "type": "object", |
| 1350 | "properties": { |
| 1351 | "server": { "type": "string", "description": "The name of the MCP server" }, |
| 1352 | "name": { "type": "string", "description": "The name of the prompt" }, |
| 1353 | "arguments": { |
| 1354 | "type": "object", |
| 1355 | "description": "Optional arguments for the prompt", |
| 1356 | "additionalProperties": { "type": "string" } |
| 1357 | } |
| 1358 | }, |
| 1359 | "required": ["server", "name"] |
| 1360 | }), |
| 1361 | allowed_callers: Some(vec!["direct".to_string()]), |
| 1362 | defer_loading: Some(false), |
| 1363 | input_examples: None, |
| 1364 | strict: None, |
| 1365 | cache_control: None, |
| 1366 | }); |
| 1367 | } |
| 1368 | |
| 1369 | api_tools |
| 1370 | } |
| 1371 | |
| 1372 | /// Call a tool by its prefixed name (mcp_{server}_{tool}) |
| 1373 | pub async fn call_tool( |
| 1374 | &mut self, |
| 1375 | prefixed_name: &str, |
| 1376 | arguments: serde_json::Value, |
| 1377 | ) -> Result<serde_json::Value> { |
| 1378 | if prefixed_name == "list_mcp_resources" { |
| 1379 | let server = arguments |
| 1380 | .get("server") |
| 1381 | .and_then(|v| v.as_str()) |
| 1382 | .map(str::to_string); |
| 1383 | let resources = self.list_resources(server).await?; |
| 1384 | return Ok(serde_json::json!({ "resources": resources })); |
| 1385 | } |
| 1386 | |
| 1387 | if prefixed_name == "list_mcp_resource_templates" { |
| 1388 | let server = arguments |
| 1389 | .get("server") |
| 1390 | .and_then(|v| v.as_str()) |
| 1391 | .map(str::to_string); |
| 1392 | let templates = self.list_resource_templates(server).await?; |
| 1393 | return Ok(serde_json::json!({ "templates": templates })); |
| 1394 | } |
| 1395 | |
| 1396 | if prefixed_name == "mcp_read_resource" { |
| 1397 | let server_name = arguments |
| 1398 | .get("server") |
| 1399 | .and_then(|v| v.as_str()) |
| 1400 | .context("Missing 'server' argument")?; |
| 1401 | let uri = arguments |
| 1402 | .get("uri") |
| 1403 | .and_then(|v| v.as_str()) |
| 1404 | .context("Missing 'uri' argument")?; |
| 1405 | return self.read_resource(server_name, uri).await; |
| 1406 | } |
| 1407 | |
| 1408 | if prefixed_name == "read_mcp_resource" { |
| 1409 | let server_name = arguments |
| 1410 | .get("server") |
| 1411 | .and_then(|v| v.as_str()) |
| 1412 | .context("Missing 'server' argument")?; |
| 1413 | let uri = arguments |
| 1414 | .get("uri") |
| 1415 | .and_then(|v| v.as_str()) |
| 1416 | .context("Missing 'uri' argument")?; |
| 1417 | return self.read_resource(server_name, uri).await; |
| 1418 | } |
| 1419 | |
| 1420 | if prefixed_name == "mcp_get_prompt" { |
| 1421 | let server_name = arguments |
| 1422 | .get("server") |
| 1423 | .and_then(|v| v.as_str()) |
| 1424 | .context("Missing 'server' argument")?; |
| 1425 | let name = arguments |
| 1426 | .get("name") |
| 1427 | .and_then(|v| v.as_str()) |
| 1428 | .context("Missing 'name' argument")?; |
| 1429 | let args = arguments |
| 1430 | .get("arguments") |
| 1431 | .cloned() |
| 1432 | .unwrap_or(serde_json::json!({})); |
| 1433 | return self.get_prompt(server_name, name, args).await; |
| 1434 | } |
| 1435 | |
| 1436 | let (server_name, tool_name) = self.parse_prefixed_name(prefixed_name)?; |
| 1437 | // Copy the global timeouts to avoid borrow conflict |
| 1438 | let global_timeouts = self.config.timeouts; |
| 1439 | let conn = self.get_or_connect(server_name).await?; |
| 1440 | if !conn.config().is_tool_enabled(tool_name) { |
| 1441 | anyhow::bail!("MCP tool '{tool_name}' is disabled for server '{server_name}'"); |
| 1442 | } |
| 1443 | let timeout = conn.config().effective_execute_timeout(&global_timeouts); |
| 1444 | conn.call_tool(tool_name, arguments, timeout).await |
| 1445 | } |
| 1446 | |
| 1447 | /// Get list of configured server names |
| 1448 | #[allow(dead_code)] // Public API for MCP consumers |
| 1449 | pub fn server_names(&self) -> Vec<&str> { |
| 1450 | self.config |
| 1451 | .servers |
| 1452 | .keys() |
| 1453 | .map(std::string::String::as_str) |
| 1454 | .collect() |
| 1455 | } |
| 1456 | |
| 1457 | /// Get list of connected server names |
| 1458 | pub fn connected_servers(&self) -> Vec<&str> { |
| 1459 | self.connections |
| 1460 | .iter() |
| 1461 | .filter(|(_, c)| c.is_ready()) |
| 1462 | .map(|(n, _)| n.as_str()) |
| 1463 | .collect() |
| 1464 | } |
| 1465 | |
| 1466 | /// Disconnect all connections |
| 1467 | #[allow(dead_code)] // Public API for MCP lifecycle management |
| 1468 | pub fn disconnect_all(&mut self) { |
| 1469 | self.connections.clear(); |
| 1470 | } |
| 1471 | |
| 1472 | /// Graceful shutdown of every connection in the pool: send SIGTERM to |
| 1473 | /// each stdio child and give them a short grace period before drop |
| 1474 | /// fires SIGKILL. Whalescale#420. |
| 1475 | /// |
| 1476 | /// Call from the TUI exit path *before* dropping the pool to give |
| 1477 | /// MCP servers a chance to flush state. The fallback Drop on |
| 1478 | /// `StdioTransport` still sends SIGTERM if this never runs, so even |
| 1479 | /// abnormal exits avoid leaking PIDs without a signal. |
| 1480 | #[allow(dead_code)] // Wired in by callers that want graceful shutdown |
| 1481 | pub async fn shutdown_all(&mut self) { |
| 1482 | let names: Vec<String> = self.connections.keys().cloned().collect(); |
| 1483 | for name in names { |
| 1484 | if let Some(conn) = self.connections.get_mut(&name) { |
| 1485 | conn.transport.shutdown().await; |
| 1486 | } |
| 1487 | } |
| 1488 | self.connections.clear(); |
| 1489 | } |
| 1490 | |
| 1491 | /// Get the underlying configuration |
| 1492 | #[allow(dead_code)] // Public API for MCP consumers |
| 1493 | pub fn config(&self) -> &McpConfig { |
| 1494 | &self.config |
| 1495 | } |
| 1496 | |
| 1497 | /// Check if a tool name is an MCP tool |
| 1498 | pub fn is_mcp_tool(name: &str) -> bool { |
| 1499 | name.starts_with("mcp_") |
| 1500 | || matches!( |
| 1501 | name, |
| 1502 | "list_mcp_resources" | "list_mcp_resource_templates" | "read_mcp_resource" |
| 1503 | ) |
| 1504 | } |
| 1505 | } |
| 1506 | |
| 1507 | #[derive(Debug, Clone, Copy, PartialEq, Eq)] |
| 1508 | pub enum McpWriteStatus { |
| 1509 | Created, |
| 1510 | Overwritten, |
| 1511 | SkippedExists, |
| 1512 | } |
| 1513 | |
| 1514 | #[derive(Debug, Clone, PartialEq, Eq)] |
| 1515 | pub struct McpDiscoveredItem { |
| 1516 | pub name: String, |
| 1517 | pub model_name: String, |
| 1518 | pub description: Option<String>, |
| 1519 | } |
| 1520 | |
| 1521 | #[derive(Debug, Clone, PartialEq, Eq)] |
| 1522 | pub struct McpServerSnapshot { |
| 1523 | pub name: String, |
| 1524 | pub enabled: bool, |
| 1525 | pub required: bool, |
| 1526 | pub transport: String, |
| 1527 | pub command_or_url: String, |
| 1528 | pub connect_timeout: u64, |
| 1529 | pub execute_timeout: u64, |
| 1530 | pub read_timeout: u64, |
| 1531 | pub connected: bool, |
| 1532 | pub error: Option<String>, |
| 1533 | pub tools: Vec<McpDiscoveredItem>, |
| 1534 | pub resources: Vec<McpDiscoveredItem>, |
| 1535 | pub prompts: Vec<McpDiscoveredItem>, |
| 1536 | } |
| 1537 | |
| 1538 | #[derive(Debug, Clone, PartialEq, Eq)] |
| 1539 | pub struct McpManagerSnapshot { |
| 1540 | pub config_path: std::path::PathBuf, |
| 1541 | pub config_exists: bool, |
| 1542 | pub restart_required: bool, |
| 1543 | pub servers: Vec<McpServerSnapshot>, |
| 1544 | } |
| 1545 | |
| 1546 | pub fn load_config(path: &Path) -> Result<McpConfig> { |
| 1547 | if !path.exists() { |
| 1548 | return Ok(McpConfig::default()); |
| 1549 | } |
| 1550 | let contents = fs::read_to_string(path) |
| 1551 | .with_context(|| format!("Failed to read MCP config {}", path.display()))?; |
| 1552 | serde_json::from_str(&contents) |
| 1553 | .with_context(|| format!("Failed to parse MCP config {}", path.display())) |
| 1554 | } |
| 1555 | |
| 1556 | pub fn save_config(path: &Path, cfg: &McpConfig) -> Result<()> { |
| 1557 | if let Some(parent) = path.parent() { |
| 1558 | fs::create_dir_all(parent).with_context(|| { |
| 1559 | format!("Failed to create MCP config directory {}", parent.display()) |
| 1560 | })?; |
| 1561 | } |
| 1562 | let rendered = serde_json::to_string_pretty(cfg).context("Failed to serialize MCP config")?; |
| 1563 | write_atomic(path, rendered.as_bytes()) |
| 1564 | .with_context(|| format!("Failed to write MCP config {}", path.display()))?; |
| 1565 | Ok(()) |
| 1566 | } |
| 1567 | |
| 1568 | fn mcp_template_json() -> Result<String> { |
| 1569 | let mut cfg = McpConfig::default(); |
| 1570 | cfg.servers.insert( |
| 1571 | "example".to_string(), |
| 1572 | McpServerConfig { |
| 1573 | command: Some("node".to_string()), |
| 1574 | args: vec!["./path/to/your-mcp-server.js".to_string()], |
| 1575 | env: HashMap::new(), |
| 1576 | url: None, |
| 1577 | connect_timeout: None, |
| 1578 | execute_timeout: None, |
| 1579 | read_timeout: None, |
| 1580 | disabled: true, |
| 1581 | enabled: true, |
| 1582 | required: false, |
| 1583 | enabled_tools: Vec::new(), |
| 1584 | disabled_tools: Vec::new(), |
| 1585 | }, |
| 1586 | ); |
| 1587 | serde_json::to_string_pretty(&cfg).context("Failed to render MCP template JSON") |
| 1588 | } |
| 1589 | |
| 1590 | pub fn init_config(path: &Path, force: bool) -> Result<McpWriteStatus> { |
| 1591 | if path.exists() && !force { |
| 1592 | return Ok(McpWriteStatus::SkippedExists); |
| 1593 | } |
| 1594 | let status = if path.exists() { |
| 1595 | McpWriteStatus::Overwritten |
| 1596 | } else { |
| 1597 | McpWriteStatus::Created |
| 1598 | }; |
| 1599 | if let Some(parent) = path.parent() { |
| 1600 | fs::create_dir_all(parent).with_context(|| { |
| 1601 | format!("Failed to create MCP config directory {}", parent.display()) |
| 1602 | })?; |
| 1603 | } |
| 1604 | let template = mcp_template_json()?; |
| 1605 | write_atomic(path, template.as_bytes()) |
| 1606 | .with_context(|| format!("Failed to write MCP config {}", path.display()))?; |
| 1607 | Ok(status) |
| 1608 | } |
| 1609 | |
| 1610 | pub fn add_server_config( |
| 1611 | path: &Path, |
| 1612 | name: String, |
| 1613 | command: Option<String>, |
| 1614 | url: Option<String>, |
| 1615 | args: Vec<String>, |
| 1616 | ) -> Result<()> { |
| 1617 | if command.is_none() && url.is_none() { |
| 1618 | anyhow::bail!("Provide either a command or URL for MCP server '{name}'."); |
| 1619 | } |
| 1620 | let mut cfg = load_config(path)?; |
| 1621 | cfg.servers.insert( |
| 1622 | name, |
| 1623 | McpServerConfig { |
| 1624 | command, |
| 1625 | args, |
| 1626 | env: HashMap::new(), |
| 1627 | url, |
| 1628 | connect_timeout: None, |
| 1629 | execute_timeout: None, |
| 1630 | read_timeout: None, |
| 1631 | disabled: false, |
| 1632 | enabled: true, |
| 1633 | required: false, |
| 1634 | enabled_tools: Vec::new(), |
| 1635 | disabled_tools: Vec::new(), |
| 1636 | }, |
| 1637 | ); |
| 1638 | save_config(path, &cfg) |
| 1639 | } |
| 1640 | |
| 1641 | pub fn remove_server_config(path: &Path, name: &str) -> Result<()> { |
| 1642 | let mut cfg = load_config(path)?; |
| 1643 | if cfg.servers.remove(name).is_none() { |
| 1644 | anyhow::bail!("MCP server '{name}' not found"); |
| 1645 | } |
| 1646 | save_config(path, &cfg) |
| 1647 | } |
| 1648 | |
| 1649 | pub fn set_server_enabled(path: &Path, name: &str, enabled: bool) -> Result<()> { |
| 1650 | let mut cfg = load_config(path)?; |
| 1651 | let server = cfg |
| 1652 | .servers |
| 1653 | .get_mut(name) |
| 1654 | .ok_or_else(|| anyhow::anyhow!("MCP server '{name}' not found"))?; |
| 1655 | server.enabled = enabled; |
| 1656 | server.disabled = !enabled; |
| 1657 | save_config(path, &cfg) |
| 1658 | } |
| 1659 | |
| 1660 | pub fn manager_snapshot_from_config( |
| 1661 | path: &Path, |
| 1662 | restart_required: bool, |
| 1663 | ) -> Result<McpManagerSnapshot> { |
| 1664 | let cfg = load_config(path)?; |
| 1665 | Ok(snapshot_from_config( |
| 1666 | path, |
| 1667 | path.exists(), |
| 1668 | restart_required, |
| 1669 | &cfg, |
| 1670 | None, |
| 1671 | )) |
| 1672 | } |
| 1673 | |
| 1674 | pub async fn discover_manager_snapshot( |
| 1675 | path: &Path, |
| 1676 | network_policy: Option<NetworkPolicyDecider>, |
| 1677 | restart_required: bool, |
| 1678 | ) -> Result<McpManagerSnapshot> { |
| 1679 | let cfg = load_config(path)?; |
| 1680 | let mut pool = McpPool::new(cfg.clone()); |
| 1681 | if let Some(policy) = network_policy { |
| 1682 | pool = pool.with_network_policy(policy); |
| 1683 | } |
| 1684 | let errors = pool |
| 1685 | .connect_all() |
| 1686 | .await |
| 1687 | .into_iter() |
| 1688 | .map(|(name, err)| (name, err.to_string())) |
| 1689 | .collect::<HashMap<_, _>>(); |
| 1690 | Ok(snapshot_from_config( |
| 1691 | path, |
| 1692 | path.exists(), |
| 1693 | restart_required, |
| 1694 | &cfg, |
| 1695 | Some((&pool, &errors)), |
| 1696 | )) |
| 1697 | } |
| 1698 | |
| 1699 | fn snapshot_from_config( |
| 1700 | path: &Path, |
| 1701 | config_exists: bool, |
| 1702 | restart_required: bool, |
| 1703 | cfg: &McpConfig, |
| 1704 | discovery: Option<(&McpPool, &HashMap<String, String>)>, |
| 1705 | ) -> McpManagerSnapshot { |
| 1706 | let mut servers = cfg |
| 1707 | .servers |
| 1708 | .iter() |
| 1709 | .map(|(name, server)| { |
| 1710 | let transport = if server.url.is_some() { |
| 1711 | "http/sse" |
| 1712 | } else { |
| 1713 | "stdio" |
| 1714 | }; |
| 1715 | let command_or_url = server.url.clone().unwrap_or_else(|| { |
| 1716 | let mut command = server |
| 1717 | .command |
| 1718 | .clone() |
| 1719 | .unwrap_or_else(|| "(missing)".to_string()); |
| 1720 | if !server.args.is_empty() { |
| 1721 | command.push(' '); |
| 1722 | command.push_str(&server.args.join(" ")); |
| 1723 | } |
| 1724 | command |
| 1725 | }); |
| 1726 | let mut snapshot = McpServerSnapshot { |
| 1727 | name: name.clone(), |
| 1728 | enabled: server.is_enabled(), |
| 1729 | required: server.required, |
| 1730 | transport: transport.to_string(), |
| 1731 | command_or_url, |
| 1732 | connect_timeout: server.effective_connect_timeout(&cfg.timeouts), |
| 1733 | execute_timeout: server.effective_execute_timeout(&cfg.timeouts), |
| 1734 | read_timeout: server.effective_read_timeout(&cfg.timeouts), |
| 1735 | connected: false, |
| 1736 | error: if server.is_enabled() { |
| 1737 | None |
| 1738 | } else { |
| 1739 | Some("disabled".to_string()) |
| 1740 | }, |
| 1741 | tools: Vec::new(), |
| 1742 | resources: Vec::new(), |
| 1743 | prompts: Vec::new(), |
| 1744 | }; |
| 1745 | |
| 1746 | if let Some((pool, errors)) = discovery { |
| 1747 | if let Some(error) = errors.get(name) { |
| 1748 | snapshot.error = Some(error.clone()); |
| 1749 | } |
| 1750 | if let Some(conn) = pool.connections.get(name) { |
| 1751 | snapshot.connected = conn.is_ready(); |
| 1752 | snapshot.tools = conn |
| 1753 | .tools() |
| 1754 | .iter() |
| 1755 | .filter(|tool| conn.config().is_tool_enabled(&tool.name)) |
| 1756 | .map(|tool| McpDiscoveredItem { |
| 1757 | name: tool.name.clone(), |
| 1758 | model_name: format!("mcp_{}_{}", name, tool.name), |
| 1759 | description: tool.description.clone(), |
| 1760 | }) |
| 1761 | .collect(); |
| 1762 | snapshot.resources = |
| 1763 | conn.resources() |
| 1764 | .iter() |
| 1765 | .map(|resource| McpDiscoveredItem { |
| 1766 | name: resource.name.clone(), |
| 1767 | model_name: format!( |
| 1768 | "mcp_{}_{}", |
| 1769 | name, |
| 1770 | resource.name.replace(' ', "_").to_lowercase() |
| 1771 | ), |
| 1772 | description: resource.description.clone(), |
| 1773 | }) |
| 1774 | .chain(conn.resource_templates().iter().map(|template| { |
| 1775 | McpDiscoveredItem { |
| 1776 | name: template.name.clone(), |
| 1777 | model_name: format!( |
| 1778 | "mcp_{}_{}", |
| 1779 | name, |
| 1780 | template.name.replace(' ', "_").to_lowercase() |
| 1781 | ), |
| 1782 | description: template.description.clone(), |
| 1783 | } |
| 1784 | })) |
| 1785 | .collect(); |
| 1786 | snapshot.prompts = conn |
| 1787 | .prompts() |
| 1788 | .iter() |
| 1789 | .map(|prompt| McpDiscoveredItem { |
| 1790 | name: prompt.name.clone(), |
| 1791 | model_name: format!("mcp_{}_{}", name, prompt.name), |
| 1792 | description: prompt.description.clone(), |
| 1793 | }) |
| 1794 | .collect(); |
| 1795 | } |
| 1796 | } |
| 1797 | |
| 1798 | snapshot |
| 1799 | }) |
| 1800 | .collect::<Vec<_>>(); |
| 1801 | servers.sort_by(|a, b| a.name.cmp(&b.name)); |
| 1802 | McpManagerSnapshot { |
| 1803 | config_path: path.to_path_buf(), |
| 1804 | config_exists, |
| 1805 | restart_required, |
| 1806 | servers, |
| 1807 | } |
| 1808 | } |
| 1809 | |
| 1810 | // === Helper Functions === |
| 1811 | |
| 1812 | /// Format MCP tool result for display |
| 1813 | #[allow(dead_code)] // Will be used when MCP tool results are displayed in TUI |
| 1814 | pub fn format_tool_result(result: &serde_json::Value) -> String { |
| 1815 | let is_error = result |
| 1816 | .get("isError") |
| 1817 | .and_then(serde_json::Value::as_bool) |
| 1818 | .unwrap_or(false); |
| 1819 | |
| 1820 | let content = result |
| 1821 | .get("content") |
| 1822 | .and_then(|v| v.as_array()) |
| 1823 | .map_or_else( |
| 1824 | || serde_json::to_string_pretty(result).unwrap_or_default(), |
| 1825 | |arr| { |
| 1826 | arr.iter() |
| 1827 | .filter_map(|item| match item.get("type")?.as_str()? { |
| 1828 | "text" => item.get("text")?.as_str().map(String::from), |
| 1829 | other => Some(format!("[{other} content]")), |
| 1830 | }) |
| 1831 | .collect::<Vec<_>>() |
| 1832 | .join("\n") |
| 1833 | }, |
| 1834 | ); |
| 1835 | |
| 1836 | if is_error { |
| 1837 | format!("Error: {content}") |
| 1838 | } else { |
| 1839 | content |
| 1840 | } |
| 1841 | } |
| 1842 | |
| 1843 | // === Unit Tests === |
| 1844 | |
| 1845 | #[cfg(test)] |
| 1846 | mod tests { |
| 1847 | use super::*; |
| 1848 | |
| 1849 | #[test] |
| 1850 | fn test_mcp_config_defaults() { |
| 1851 | let config = McpConfig::default(); |
| 1852 | assert_eq!(config.timeouts.connect_timeout, 10); |
| 1853 | assert_eq!(config.timeouts.execute_timeout, 60); |
| 1854 | assert_eq!(config.timeouts.read_timeout, 120); |
| 1855 | assert!(config.servers.is_empty()); |
| 1856 | } |
| 1857 | |
| 1858 | #[test] |
| 1859 | fn test_mcp_config_parse() { |
| 1860 | let json = r#"{ |
| 1861 | "timeouts": { |
| 1862 | "connect_timeout": 15, |
| 1863 | "execute_timeout": 90 |
| 1864 | }, |
| 1865 | "servers": { |
| 1866 | "test": { |
| 1867 | "command": "node", |
| 1868 | "args": ["server.js"], |
| 1869 | "env": {"FOO": "bar"} |
| 1870 | } |
| 1871 | } |
| 1872 | }"#; |
| 1873 | |
| 1874 | let config: McpConfig = serde_json::from_str(json).unwrap(); |
| 1875 | assert_eq!(config.timeouts.connect_timeout, 15); |
| 1876 | assert_eq!(config.timeouts.execute_timeout, 90); |
| 1877 | assert_eq!(config.timeouts.read_timeout, 120); // default |
| 1878 | assert!(config.servers.contains_key("test")); |
| 1879 | |
| 1880 | let server = config.servers.get("test").unwrap(); |
| 1881 | assert_eq!(server.command, Some("node".to_string())); |
| 1882 | assert_eq!(server.args, vec!["server.js"]); |
| 1883 | assert_eq!(server.env.get("FOO"), Some(&"bar".to_string())); |
| 1884 | } |
| 1885 | |
| 1886 | #[test] |
| 1887 | fn test_mcp_config_parse_mcp_servers_alias_and_snapshot() { |
| 1888 | let dir = tempfile::tempdir().unwrap(); |
| 1889 | let path = dir.path().join("mcp.json"); |
| 1890 | fs::write( |
| 1891 | &path, |
| 1892 | r#"{ |
| 1893 | "mcpServers": { |
| 1894 | "disabled": { |
| 1895 | "command": "node", |
| 1896 | "args": ["server.js"], |
| 1897 | "disabled": true |
| 1898 | } |
| 1899 | } |
| 1900 | }"#, |
| 1901 | ) |
| 1902 | .unwrap(); |
| 1903 | |
| 1904 | let cfg = load_config(&path).unwrap(); |
| 1905 | assert!(cfg.servers.contains_key("disabled")); |
| 1906 | let snapshot = manager_snapshot_from_config(&path, true).unwrap(); |
| 1907 | assert!(snapshot.restart_required); |
| 1908 | assert_eq!(snapshot.servers[0].name, "disabled"); |
| 1909 | assert!(!snapshot.servers[0].enabled); |
| 1910 | assert_eq!(snapshot.servers[0].error.as_deref(), Some("disabled")); |
| 1911 | } |
| 1912 | |
| 1913 | #[test] |
| 1914 | fn test_mcp_config_manager_actions_round_trip() { |
| 1915 | let dir = tempfile::tempdir().unwrap(); |
| 1916 | let path = dir.path().join("mcp.json"); |
| 1917 | |
| 1918 | assert_eq!(init_config(&path, false).unwrap(), McpWriteStatus::Created); |
| 1919 | assert_eq!( |
| 1920 | init_config(&path, false).unwrap(), |
| 1921 | McpWriteStatus::SkippedExists |
| 1922 | ); |
| 1923 | |
| 1924 | add_server_config( |
| 1925 | &path, |
| 1926 | "local".to_string(), |
| 1927 | Some("node".to_string()), |
| 1928 | None, |
| 1929 | vec!["server.js".to_string()], |
| 1930 | ) |
| 1931 | .unwrap(); |
| 1932 | set_server_enabled(&path, "local", false).unwrap(); |
| 1933 | let disabled = manager_snapshot_from_config(&path, true).unwrap(); |
| 1934 | let local = disabled |
| 1935 | .servers |
| 1936 | .iter() |
| 1937 | .find(|server| server.name == "local") |
| 1938 | .unwrap(); |
| 1939 | assert!(!local.enabled); |
| 1940 | assert_eq!(local.transport, "stdio"); |
| 1941 | |
| 1942 | remove_server_config(&path, "local").unwrap(); |
| 1943 | let removed = manager_snapshot_from_config(&path, true).unwrap(); |
| 1944 | assert!(removed.servers.iter().all(|server| server.name != "local")); |
| 1945 | } |
| 1946 | |
| 1947 | #[test] |
| 1948 | fn test_server_effective_timeouts() { |
| 1949 | let global = McpTimeouts::default(); |
| 1950 | |
| 1951 | let server_with_override = McpServerConfig { |
| 1952 | command: Some("test".to_string()), |
| 1953 | args: vec![], |
| 1954 | env: HashMap::new(), |
| 1955 | url: None, |
| 1956 | connect_timeout: Some(20), |
| 1957 | execute_timeout: None, |
| 1958 | read_timeout: Some(180), |
| 1959 | disabled: false, |
| 1960 | enabled: true, |
| 1961 | required: false, |
| 1962 | enabled_tools: Vec::new(), |
| 1963 | disabled_tools: Vec::new(), |
| 1964 | }; |
| 1965 | |
| 1966 | assert_eq!(server_with_override.effective_connect_timeout(&global), 20); |
| 1967 | assert_eq!(server_with_override.effective_execute_timeout(&global), 60); // global default |
| 1968 | assert_eq!(server_with_override.effective_read_timeout(&global), 180); |
| 1969 | } |
| 1970 | |
| 1971 | #[test] |
| 1972 | fn test_mcp_pool_is_mcp_tool() { |
| 1973 | assert!(McpPool::is_mcp_tool("mcp_filesystem_read")); |
| 1974 | assert!(McpPool::is_mcp_tool("mcp_git_status")); |
| 1975 | assert!(McpPool::is_mcp_tool("list_mcp_resources")); |
| 1976 | assert!(McpPool::is_mcp_tool("list_mcp_resource_templates")); |
| 1977 | assert!(McpPool::is_mcp_tool("read_mcp_resource")); |
| 1978 | assert!(!McpPool::is_mcp_tool("read_file")); |
| 1979 | assert!(!McpPool::is_mcp_tool("exec_shell")); |
| 1980 | } |
| 1981 | |
| 1982 | #[test] |
| 1983 | fn test_format_tool_result_text() { |
| 1984 | let result = serde_json::json!({ |
| 1985 | "content": [ |
| 1986 | {"type": "text", "text": "Hello, world!"} |
| 1987 | ] |
| 1988 | }); |
| 1989 | assert_eq!(format_tool_result(&result), "Hello, world!"); |
| 1990 | } |
| 1991 | |
| 1992 | #[test] |
| 1993 | fn test_format_tool_result_error() { |
| 1994 | let result = serde_json::json!({ |
| 1995 | "isError": true, |
| 1996 | "content": [ |
| 1997 | {"type": "text", "text": "Something went wrong"} |
| 1998 | ] |
| 1999 | }); |
| 2000 | assert_eq!(format_tool_result(&result), "Error: Something went wrong"); |
| 2001 | } |
| 2002 | |
| 2003 | #[test] |
| 2004 | fn test_format_tool_result_multiple_content() { |
| 2005 | let result = serde_json::json!({ |
| 2006 | "content": [ |
| 2007 | {"type": "text", "text": "Line 1"}, |
| 2008 | {"type": "text", "text": "Line 2"}, |
| 2009 | {"type": "image", "data": "base64..."} |
| 2010 | ] |
| 2011 | }); |
| 2012 | let formatted = format_tool_result(&result); |
| 2013 | assert!(formatted.contains("Line 1")); |
| 2014 | assert!(formatted.contains("Line 2")); |
| 2015 | assert!(formatted.contains("[image content]")); |
| 2016 | } |
| 2017 | |
| 2018 | #[tokio::test] |
| 2019 | async fn test_mcp_pool_empty_config() { |
| 2020 | let pool = McpPool::new(McpConfig::default()); |
| 2021 | assert!(pool.server_names().is_empty()); |
| 2022 | assert!(pool.all_tools().is_empty()); |
| 2023 | } |
| 2024 | |
| 2025 | #[test] |
| 2026 | fn mask_url_secrets_strips_userinfo() { |
| 2027 | let masked = mask_url_secrets("https://user:s3cret@host.example/api?foo=bar"); |
| 2028 | assert!(masked.contains("***"), "expected masked userinfo: {masked}"); |
| 2029 | assert!(!masked.contains("s3cret"), "secret leaked: {masked}"); |
| 2030 | assert!(masked.contains("host.example"), "host preserved: {masked}"); |
| 2031 | } |
| 2032 | |
| 2033 | #[test] |
| 2034 | fn mask_url_secrets_passes_through_clean_url() { |
| 2035 | assert_eq!( |
| 2036 | mask_url_secrets("https://api.example.com/mcp"), |
| 2037 | "https://api.example.com/mcp" |
| 2038 | ); |
| 2039 | } |
| 2040 | |
| 2041 | #[test] |
| 2042 | fn redact_body_preview_masks_bearer_token() { |
| 2043 | let redacted = redact_body_preview("Authorization: Bearer abc.def.ghi end"); |
| 2044 | assert!(redacted.contains("Bearer ***"), "redacted: {redacted}"); |
| 2045 | assert!(!redacted.contains("abc.def.ghi"), "leaked: {redacted}"); |
| 2046 | } |
| 2047 | |
| 2048 | #[test] |
| 2049 | fn redact_body_preview_masks_api_key_param() { |
| 2050 | let redacted = redact_body_preview("error message api_key=sk-12345&other=val"); |
| 2051 | assert!(redacted.contains("api_key=***"), "redacted: {redacted}"); |
| 2052 | assert!(!redacted.contains("sk-12345"), "leaked: {redacted}"); |
| 2053 | assert!( |
| 2054 | redacted.contains("other=val"), |
| 2055 | "non-secret preserved: {redacted}" |
| 2056 | ); |
| 2057 | } |
| 2058 | |
| 2059 | /// #420: `StdioTransport::shutdown` reaps the child process by sending |
| 2060 | /// SIGTERM and giving it a brief grace period before drop fires SIGKILL. |
| 2061 | /// The test spawns `cat` (which exits immediately on stdin EOF / SIGTERM) |
| 2062 | /// and verifies the transport tears down cleanly. Unix-only because |
| 2063 | /// SIGTERM doesn't exist on Windows; on Windows the test would just |
| 2064 | /// duplicate the kill_on_drop path. |
| 2065 | #[cfg(unix)] |
| 2066 | #[tokio::test] |
| 2067 | async fn stdio_transport_shutdown_terminates_child() { |
| 2068 | use tokio::process::Command as TokioCommand; |
| 2069 | let mut cmd = TokioCommand::new("cat"); |
| 2070 | cmd.stdin(std::process::Stdio::piped()) |
| 2071 | .stdout(std::process::Stdio::piped()) |
| 2072 | .stderr(std::process::Stdio::null()) |
| 2073 | .kill_on_drop(true); |
| 2074 | let mut child = cmd.spawn().expect("spawn cat"); |
| 2075 | let pid = child.id().expect("child pid"); |
| 2076 | let stdin = child.stdin.take().expect("child stdin"); |
| 2077 | let stdout = child.stdout.take().expect("child stdout"); |
| 2078 | let mut transport = StdioTransport { |
| 2079 | child, |
| 2080 | stdin, |
| 2081 | reader: tokio::io::BufReader::new(stdout), |
| 2082 | }; |
| 2083 | |
| 2084 | // shutdown() should send SIGTERM and complete within the grace window. |
| 2085 | let start = std::time::Instant::now(); |
| 2086 | transport.shutdown().await; |
| 2087 | let elapsed = start.elapsed(); |
| 2088 | assert!( |
| 2089 | elapsed < STDIO_SHUTDOWN_GRACE + Duration::from_millis(500), |
| 2090 | "shutdown blocked beyond grace window: {elapsed:?}" |
| 2091 | ); |
| 2092 | |
| 2093 | // The child should be reaped — kill(pid, 0) returning ESRCH means |
| 2094 | // the pid is gone. If it's still alive, kill(0) returns 0, which |
| 2095 | // means our shutdown didn't terminate it. |
| 2096 | // SAFETY: pid was just collected from a tokio Child we spawned. |
| 2097 | // libc::kill with signal 0 only checks pid existence and is |
| 2098 | // async-signal-safe. |
| 2099 | let still_alive = unsafe { libc::kill(pid as i32, 0) } == 0; |
| 2100 | assert!( |
| 2101 | !still_alive, |
| 2102 | "child {pid} survived StdioTransport::shutdown — SIGTERM not delivered" |
| 2103 | ); |
| 2104 | } |
| 2105 | } |
| 2106 |