| 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::{BTreeSet, HashMap, HashSet}; |
| 9 | use std::ffi::{OsStr, OsString}; |
| 10 | use std::fs; |
| 11 | use std::future::Future; |
| 12 | use std::io::{Read, Seek}; |
| 13 | use std::path::{Component, Path, PathBuf}; |
| 14 | use std::sync::Arc; |
| 15 | use std::sync::atomic::{AtomicU64, Ordering}; |
| 16 | use std::time::Duration; |
| 17 | |
| 18 | use anyhow::{Context, Result}; |
| 19 | use futures_util::FutureExt; |
| 20 | use parking_lot::RwLock; |
| 21 | use serde::{Deserialize, Serialize}; |
| 22 | use sha2::Digest as _; |
| 23 | |
| 24 | pub mod external_import; |
| 25 | mod headers; |
| 26 | mod http; |
| 27 | mod http_client; |
| 28 | pub mod oauth; |
| 29 | mod sse; |
| 30 | mod stdio; |
| 31 | mod streamable_http; |
| 32 | mod wire; |
| 33 | |
| 34 | use self::http::{HttpTransport, McpHttpAuth}; |
| 35 | use self::sse::SseTransport; |
| 36 | use self::stdio::StdioTransport; |
| 37 | #[cfg(all(test, unix))] |
| 38 | use self::stdio::{STDIO_SHUTDOWN_GRACE, StderrTail}; |
| 39 | use self::wire::{is_mcp_stale_session_body, is_retriable_mcp_call_error}; |
| 40 | use crate::network_policy::{Decision, NetworkPolicyDecider, host_from_url}; |
| 41 | use crate::utils::write_atomic; |
| 42 | |
| 43 | // === Error diagnostics helpers (#71) === |
| 44 | |
| 45 | /// Bytes of a non-2xx response body to surface in connection errors. |
| 46 | const ERROR_BODY_PREVIEW_BYTES: usize = 200; |
| 47 | |
| 48 | /// Newest dated MCP protocol revision Codewhale advertises at `initialize` and |
| 49 | /// answers as an MCP server. Matches the shared MCP crate (`crates/mcp`). |
| 50 | pub(crate) const MCP_PROTOCOL_VERSION: &str = "2025-06-18"; |
| 51 | /// Dated MCP revisions accepted during negotiation, newest first. A peer |
| 52 | /// answering or requesting any of these continues the handshake. |
| 53 | pub(crate) const MCP_SUPPORTED_PROTOCOL_VERSIONS: &[&str] = |
| 54 | &[MCP_PROTOCOL_VERSION, "2025-03-26", "2024-11-05"]; |
| 55 | |
| 56 | fn validate_mcp_config_path(path: &Path) -> Result<()> { |
| 57 | if path.as_os_str().is_empty() { |
| 58 | anyhow::bail!("MCP config path cannot be empty"); |
| 59 | } |
| 60 | if path |
| 61 | .components() |
| 62 | .any(|component| matches!(component, Component::ParentDir)) |
| 63 | { |
| 64 | anyhow::bail!("MCP config path cannot contain '..' components"); |
| 65 | } |
| 66 | Ok(()) |
| 67 | } |
| 68 | |
| 69 | /// Expand `${NAME}` placeholders in an MCP config value from the process |
| 70 | /// environment. This lets secrets (API keys, bearer tokens, …) be supplied |
| 71 | /// through environment variables instead of being written in cleartext into |
| 72 | /// the MCP config file on disk. |
| 73 | /// |
| 74 | /// On a missing or malformed placeholder the error names only the offending |
| 75 | /// variable, never the surrounding value, so a secret-bearing string is never |
| 76 | /// echoed into logs or error output. |
| 77 | fn expand_env_placeholders_with( |
| 78 | value: &str, |
| 79 | environment: Option<&crate::plugins::HostEnvironment>, |
| 80 | ) -> Result<String> { |
| 81 | let mut out = String::new(); |
| 82 | let mut rest = value; |
| 83 | while let Some(start) = rest.find("${") { |
| 84 | out.push_str(&rest[..start]); |
| 85 | let after = &rest[start + 2..]; |
| 86 | let Some(end) = after.find('}') else { |
| 87 | anyhow::bail!("unterminated environment placeholder in MCP config value"); |
| 88 | }; |
| 89 | let name = &after[..end]; |
| 90 | if name.is_empty() || !name.chars().all(|c| c.is_ascii_alphanumeric() || c == '_') { |
| 91 | anyhow::bail!("invalid environment placeholder in MCP config value"); |
| 92 | } |
| 93 | let env_value = environment |
| 94 | .map_or_else(|| std::env::var(name), |env| env.var(name)) |
| 95 | .with_context(|| { |
| 96 | format!("environment variable {name} required by MCP config is not set") |
| 97 | })?; |
| 98 | out.push_str(&env_value); |
| 99 | rest = &after[end + 1..]; |
| 100 | } |
| 101 | out.push_str(rest); |
| 102 | Ok(out) |
| 103 | } |
| 104 | |
| 105 | #[cfg(test)] |
| 106 | fn expand_env_placeholders(value: &str) -> Result<String> { |
| 107 | expand_env_placeholders_with(value, None) |
| 108 | } |
| 109 | |
| 110 | /// Expand `${NAME}` placeholders across every value of an MCP config map |
| 111 | /// (e.g. the stdio child `env`). `context` only labels expansion errors so a |
| 112 | /// failure can be attributed to the right map. |
| 113 | fn expand_env_placeholders_map_with_environment( |
| 114 | values: &HashMap<String, String>, |
| 115 | context: &str, |
| 116 | environment: Option<&crate::plugins::HostEnvironment>, |
| 117 | ) -> Result<HashMap<String, String>> { |
| 118 | let mut expanded = HashMap::with_capacity(values.len()); |
| 119 | for (key, value) in values { |
| 120 | expanded.insert( |
| 121 | key.clone(), |
| 122 | expand_env_placeholders_with(value, environment) |
| 123 | .with_context(|| format!("failed to expand MCP {context} value for {key}"))?, |
| 124 | ); |
| 125 | } |
| 126 | Ok(expanded) |
| 127 | } |
| 128 | |
| 129 | #[cfg(test)] |
| 130 | fn expand_env_placeholders_map( |
| 131 | values: &HashMap<String, String>, |
| 132 | context: &str, |
| 133 | ) -> Result<HashMap<String, String>> { |
| 134 | expand_env_placeholders_map_with_environment(values, context, None) |
| 135 | } |
| 136 | |
| 137 | fn expanded_mcp_stdio_env(config: &McpServerConfig) -> Result<HashMap<String, String>> { |
| 138 | let environment = config |
| 139 | .reviewed_plugin |
| 140 | .as_ref() |
| 141 | .map(|source| source.host_environment.as_ref()); |
| 142 | expand_env_placeholders_map_with_environment(&config.env, "env", environment) |
| 143 | } |
| 144 | |
| 145 | /// Mirror the exact expanded and sanitized environment applied by the MCP |
| 146 | /// stdio spawn path, without constructing or starting a process. |
| 147 | fn mcp_stdio_child_env(config: &McpServerConfig) -> Result<Vec<(OsString, OsString)>> { |
| 148 | let expanded_env = expanded_mcp_stdio_env(config)?; |
| 149 | let overrides = crate::child_env::string_map_env(&expanded_env); |
| 150 | Ok(if let Some(source) = config.reviewed_plugin.as_ref() { |
| 151 | // Plugin reviews name every extra environment source explicitly. Do |
| 152 | // not silently widen that consent to the compatibility-oriented MCP |
| 153 | // bootstrap namespace (for example NPM_CONFIG_*). |
| 154 | crate::child_env::sanitized_plugin_mcp_env_from( |
| 155 | source.host_environment.entries().iter().cloned(), |
| 156 | overrides, |
| 157 | ) |
| 158 | } else { |
| 159 | crate::child_env::sanitized_mcp_env(overrides) |
| 160 | }) |
| 161 | } |
| 162 | |
| 163 | #[derive(Debug, Clone, Copy, PartialEq, Eq)] |
| 164 | pub(crate) enum McpCommandAvailability { |
| 165 | Available, |
| 166 | Missing, |
| 167 | NotApplicable, |
| 168 | NotChecked, |
| 169 | } |
| 170 | |
| 171 | impl McpCommandAvailability { |
| 172 | pub(crate) fn as_str(self) -> &'static str { |
| 173 | match self { |
| 174 | Self::Available => "available", |
| 175 | Self::Missing => "missing", |
| 176 | Self::NotApplicable => "not_applicable", |
| 177 | Self::NotChecked => "not_checked", |
| 178 | } |
| 179 | } |
| 180 | } |
| 181 | |
| 182 | pub(crate) fn is_relative_stdio_path_arg(value: &str) -> bool { |
| 183 | if value.is_empty() || value.starts_with('-') || value.contains("://") || value.starts_with('~') |
| 184 | { |
| 185 | return false; |
| 186 | } |
| 187 | let looks_like_path = value.contains('/') || value.contains('\\'); |
| 188 | if !looks_like_path { |
| 189 | return false; |
| 190 | } |
| 191 | let bytes = value.as_bytes(); |
| 192 | let windows_absolute = value.starts_with("\\\\") |
| 193 | || (bytes.len() >= 3 && bytes[1] == b':' && (bytes[2] == b'\\' || bytes[2] == b'/')); |
| 194 | !Path::new(value).is_absolute() && !windows_absolute |
| 195 | } |
| 196 | |
| 197 | fn env_value<'a>(env: &'a [(OsString, OsString)], name: &str) -> Option<&'a OsStr> { |
| 198 | env.iter() |
| 199 | .rev() |
| 200 | .find(|(key, _)| { |
| 201 | #[cfg(windows)] |
| 202 | { |
| 203 | key.to_string_lossy().eq_ignore_ascii_case(name) |
| 204 | } |
| 205 | #[cfg(not(windows))] |
| 206 | { |
| 207 | key == OsStr::new(name) |
| 208 | } |
| 209 | }) |
| 210 | .map(|(_, value)| value.as_os_str()) |
| 211 | } |
| 212 | |
| 213 | #[cfg(unix)] |
| 214 | fn spawnable_command_file(path: &Path) -> bool { |
| 215 | use std::os::unix::fs::PermissionsExt; |
| 216 | |
| 217 | path.is_file() |
| 218 | && fs::metadata(path) |
| 219 | .map(|metadata| metadata.permissions().mode() & 0o111 != 0) |
| 220 | .unwrap_or(false) |
| 221 | } |
| 222 | |
| 223 | #[cfg(windows)] |
| 224 | fn spawnable_command_file(path: &Path) -> bool { |
| 225 | path.is_file() || (path.extension().is_none() && path.with_extension("exe").is_file()) |
| 226 | } |
| 227 | |
| 228 | #[cfg(not(any(unix, windows)))] |
| 229 | fn spawnable_command_file(path: &Path) -> bool { |
| 230 | path.is_file() |
| 231 | } |
| 232 | |
| 233 | fn path_candidate(dir: &Path, name: &str, cwd: Option<&Path>) -> PathBuf { |
| 234 | #[cfg(unix)] |
| 235 | { |
| 236 | // Unix performs PATH lookup after applying Command::current_dir. That |
| 237 | // includes empty PATH entries, which mean the child's current dir. |
| 238 | if dir.is_relative() |
| 239 | && let Some(cwd) = cwd |
| 240 | { |
| 241 | return cwd.join(dir).join(name); |
| 242 | } |
| 243 | } |
| 244 | #[cfg(not(unix))] |
| 245 | let _ = cwd; |
| 246 | dir.join(name) |
| 247 | } |
| 248 | |
| 249 | fn command_availability_on_path( |
| 250 | name: &str, |
| 251 | env: &[(OsString, OsString)], |
| 252 | cwd: Option<&Path>, |
| 253 | ) -> McpCommandAvailability { |
| 254 | let Some(path) = env_value(env, "PATH") else { |
| 255 | // On Unix execvp falls back to an OS-defined path. On Windows Rust's |
| 256 | // resolver still checks system and parent locations. We cannot prove a |
| 257 | // miss without reproducing platform internals, so remain conservative. |
| 258 | return McpCommandAvailability::NotChecked; |
| 259 | }; |
| 260 | for dir in std::env::split_paths(path) { |
| 261 | let candidate = path_candidate(&dir, name, cwd); |
| 262 | if spawnable_command_file(&candidate) { |
| 263 | return McpCommandAvailability::Available; |
| 264 | } |
| 265 | } |
| 266 | |
| 267 | #[cfg(windows)] |
| 268 | { |
| 269 | // Windows Command resolution also checks the running executable's |
| 270 | // directory, system directories, and the parent PATH after an explicit |
| 271 | // child PATH. A static miss in the child PATH is therefore not proof |
| 272 | // that spawn will fail. PATHEXT is intentionally not consulted: Rust |
| 273 | // only supplies an omitted `.exe`; `.cmd`/`.bat` must be explicit. |
| 274 | return McpCommandAvailability::NotChecked; |
| 275 | } |
| 276 | #[cfg(not(windows))] |
| 277 | { |
| 278 | McpCommandAvailability::Missing |
| 279 | } |
| 280 | } |
| 281 | |
| 282 | /// Inspect an MCP stdio command using the same expanded, sanitized environment |
| 283 | /// as the real spawn path, without starting the configured process. |
| 284 | pub(crate) fn static_mcp_command_availability( |
| 285 | server: &McpServerConfig, |
| 286 | ) -> Result<McpCommandAvailability> { |
| 287 | if server.url.is_some() { |
| 288 | return Ok(McpCommandAvailability::NotApplicable); |
| 289 | } |
| 290 | let Some(cmd) = server.command.as_deref() else { |
| 291 | return Ok(McpCommandAvailability::NotChecked); |
| 292 | }; |
| 293 | if cmd.is_empty() { |
| 294 | return Ok(McpCommandAvailability::Missing); |
| 295 | } |
| 296 | |
| 297 | // StdioTransport expands every configured env value before spawning, even |
| 298 | // when the command itself is absolute. Mirror that failure boundary here. |
| 299 | let child_env = mcp_stdio_child_env(server)?; |
| 300 | let path = Path::new(cmd); |
| 301 | let is_absolute = path.is_absolute() || cmd.starts_with('/'); |
| 302 | if is_absolute { |
| 303 | return Ok(if spawnable_command_file(path) { |
| 304 | McpCommandAvailability::Available |
| 305 | } else { |
| 306 | McpCommandAvailability::Missing |
| 307 | }); |
| 308 | } |
| 309 | |
| 310 | if is_relative_stdio_path_arg(cmd) { |
| 311 | let Some(cwd) = server.cwd.as_deref() else { |
| 312 | return Ok(McpCommandAvailability::NotChecked); |
| 313 | }; |
| 314 | return Ok(if spawnable_command_file(&cwd.join(path)) { |
| 315 | McpCommandAvailability::Available |
| 316 | } else { |
| 317 | McpCommandAvailability::Missing |
| 318 | }); |
| 319 | } |
| 320 | |
| 321 | Ok(command_availability_on_path( |
| 322 | cmd, |
| 323 | &child_env, |
| 324 | server.cwd.as_deref(), |
| 325 | )) |
| 326 | } |
| 327 | |
| 328 | /// Mask a URL so any embedded credentials in the userinfo portion (e.g. |
| 329 | /// `https://user:secret@host`) are replaced with `***`. Failures fall back to |
| 330 | /// the original string so we don't lose context — we never want masking to |
| 331 | /// produce an empty error. |
| 332 | fn mask_url_secrets(url: &str) -> String { |
| 333 | if let Ok(parsed) = reqwest::Url::parse(url) { |
| 334 | let mut clone = parsed.clone(); |
| 335 | if !parsed.username().is_empty() || parsed.password().is_some() { |
| 336 | let _ = clone.set_username("***"); |
| 337 | let _ = clone.set_password(Some("***")); |
| 338 | } |
| 339 | if parsed.query().is_some() { |
| 340 | clone.set_query(Some("***")); |
| 341 | } |
| 342 | clone.set_fragment(None); |
| 343 | return clone.to_string(); |
| 344 | } |
| 345 | url.to_string() |
| 346 | } |
| 347 | |
| 348 | /// Redact the userinfo segment (`username[:password]@…` portion) from |
| 349 | /// a proxy URL so it can be safely included in `tracing::warn!` output |
| 350 | /// without leaking the |
| 351 | /// password into the on-disk log. URLs without userinfo are returned |
| 352 | /// unchanged. Garbage input (no `://` scheme separator) is also returned |
| 353 | /// unchanged — the malformed-URL warning path is the only caller, so an |
| 354 | /// unparseable input is already the failure case. |
| 355 | fn redact_proxy_userinfo(proxy_url: &str) -> String { |
| 356 | let Some(scheme_end) = proxy_url.find("://") else { |
| 357 | return proxy_url.to_string(); |
| 358 | }; |
| 359 | let after_scheme = scheme_end + 3; |
| 360 | // The userinfo segment ends at the next `@`, but only if that `@` |
| 361 | // comes before the next `/`, `?`, or `#` (otherwise the `@` is in a |
| 362 | // path / query and the URL has no userinfo at all). |
| 363 | let rest = &proxy_url[after_scheme..]; |
| 364 | let at_idx = rest.find('@'); |
| 365 | let path_idx = rest.find(['/', '?', '#']); |
| 366 | let userinfo_end = match (at_idx, path_idx) { |
| 367 | (Some(a), Some(p)) if a < p => Some(a), |
| 368 | (Some(a), None) => Some(a), |
| 369 | _ => None, |
| 370 | }; |
| 371 | if let Some(end) = userinfo_end { |
| 372 | let mut out = String::with_capacity(proxy_url.len()); |
| 373 | out.push_str(&proxy_url[..after_scheme]); |
| 374 | out.push_str("***@"); |
| 375 | out.push_str(&rest[end + 1..]); |
| 376 | out |
| 377 | } else { |
| 378 | proxy_url.to_string() |
| 379 | } |
| 380 | } |
| 381 | |
| 382 | fn redact_values_after_ascii_needle( |
| 383 | output: &mut String, |
| 384 | needle: &str, |
| 385 | terminates: impl Fn(char) -> bool, |
| 386 | ) { |
| 387 | let needle = needle.as_bytes(); |
| 388 | let mut search_from = 0_usize; |
| 389 | while search_from.saturating_add(needle.len()) <= output.len() { |
| 390 | let Some(relative) = output.as_bytes()[search_from..] |
| 391 | .windows(needle.len()) |
| 392 | .position(|candidate| candidate.eq_ignore_ascii_case(needle)) |
| 393 | else { |
| 394 | break; |
| 395 | }; |
| 396 | let value_start = search_from + relative + needle.len(); |
| 397 | let value_end = output[value_start..] |
| 398 | .char_indices() |
| 399 | .find(|(_, ch)| terminates(*ch)) |
| 400 | .map_or(output.len(), |(offset, _)| value_start + offset); |
| 401 | if value_end == value_start { |
| 402 | if value_start == output.len() { |
| 403 | break; |
| 404 | } |
| 405 | // The empty value is already safe. Advance over its ASCII |
| 406 | // separator so a second occurrence later in the body is found. |
| 407 | search_from = value_start + 1; |
| 408 | continue; |
| 409 | } |
| 410 | output.replace_range(value_start..value_end, "***"); |
| 411 | search_from = value_start + 3; |
| 412 | } |
| 413 | } |
| 414 | |
| 415 | /// Mask obvious token-like substrings in a body excerpt before surfacing it. |
| 416 | /// Every occurrence is replaced, not only the first one. |
| 417 | fn redact_body_preview(body: &str) -> String { |
| 418 | let mut out = body.to_string(); |
| 419 | redact_values_after_ascii_needle(&mut out, "bearer ", |ch| { |
| 420 | ch.is_whitespace() || ch == '"' || ch == ',' |
| 421 | }); |
| 422 | for needle in ["api_key=", "apikey=", "api-key=", "token="] { |
| 423 | redact_values_after_ascii_needle(&mut out, needle, |ch| { |
| 424 | ch.is_whitespace() || ch == '&' || ch == '"' || ch == ',' |
| 425 | }); |
| 426 | } |
| 427 | out |
| 428 | } |
| 429 | |
| 430 | /// Read at most `max_bytes` of a reqwest response body and produce a |
| 431 | /// single-line excerpt suitable for an error message. The stream is dropped as |
| 432 | /// soon as the cap is reached, so an unbounded or never-ending error response |
| 433 | /// cannot make diagnostics retain the entire body. Best-effort — if the body |
| 434 | /// can't be read, returns the literal string `<no body>`. |
| 435 | async fn bounded_body_excerpt(response: reqwest::Response, max_bytes: usize) -> String { |
| 436 | use futures_util::StreamExt; |
| 437 | |
| 438 | let declared_truncated = response |
| 439 | .content_length() |
| 440 | .is_some_and(|length| length > max_bytes as u64); |
| 441 | let mut stream = response.bytes_stream(); |
| 442 | let mut body = Vec::with_capacity(max_bytes.min(8 * 1024)); |
| 443 | let mut truncated = declared_truncated; |
| 444 | |
| 445 | while body.len() < max_bytes { |
| 446 | let Some(chunk) = stream.next().await else { |
| 447 | break; |
| 448 | }; |
| 449 | let Ok(chunk) = chunk else { |
| 450 | break; |
| 451 | }; |
| 452 | let remaining = max_bytes - body.len(); |
| 453 | if chunk.len() > remaining { |
| 454 | body.extend_from_slice(&chunk[..remaining]); |
| 455 | truncated = true; |
| 456 | break; |
| 457 | } |
| 458 | body.extend_from_slice(&chunk); |
| 459 | if body.len() == max_bytes { |
| 460 | // For a chunked response there is no length that proves EOF. Stop |
| 461 | // now rather than polling an attacker-controlled stream again. |
| 462 | truncated = true; |
| 463 | } |
| 464 | } |
| 465 | |
| 466 | if body.is_empty() { |
| 467 | return "<no body>".to_string(); |
| 468 | } |
| 469 | |
| 470 | let one_line = String::from_utf8_lossy(&body).replace(['\n', '\r'], " "); |
| 471 | let suffix = if truncated { "…" } else { "" }; |
| 472 | format!("{}{}", redact_body_preview(&one_line), suffix) |
| 473 | } |
| 474 | |
| 475 | fn invalid_json_preview(bytes: &[u8]) -> String { |
| 476 | let body_text = String::from_utf8_lossy(bytes); |
| 477 | if body_text.is_empty() { |
| 478 | return "<empty>".to_string(); |
| 479 | } |
| 480 | |
| 481 | let trimmed: String = body_text.chars().take(ERROR_BODY_PREVIEW_BYTES).collect(); |
| 482 | let suffix = if body_text.chars().count() > ERROR_BODY_PREVIEW_BYTES { |
| 483 | "…" |
| 484 | } else { |
| 485 | "" |
| 486 | }; |
| 487 | let one_line = trimmed.replace(['\n', '\r'], " "); |
| 488 | format!("{}{}", redact_body_preview(&one_line), suffix) |
| 489 | } |
| 490 | |
| 491 | // === Configuration Types === |
| 492 | |
| 493 | /// Full MCP configuration from mcp.json |
| 494 | #[derive(Debug, Clone, Default, Deserialize, Serialize)] |
| 495 | pub struct McpConfig { |
| 496 | #[serde(default)] |
| 497 | pub timeouts: McpTimeouts, |
| 498 | #[serde(default, alias = "mcpServers")] |
| 499 | pub servers: HashMap<String, McpServerConfig>, |
| 500 | } |
| 501 | |
| 502 | /// Global timeout configuration |
| 503 | #[derive(Debug, Clone, Copy, Deserialize, Serialize)] |
| 504 | #[allow(clippy::struct_field_names)] |
| 505 | pub struct McpTimeouts { |
| 506 | #[serde(default = "default_connect_timeout")] |
| 507 | pub connect_timeout: u64, |
| 508 | #[serde(default = "default_execute_timeout")] |
| 509 | pub execute_timeout: u64, |
| 510 | #[serde(default = "default_read_timeout")] |
| 511 | pub read_timeout: u64, |
| 512 | } |
| 513 | |
| 514 | fn default_connect_timeout() -> u64 { |
| 515 | 10 |
| 516 | } |
| 517 | fn default_execute_timeout() -> u64 { |
| 518 | 60 |
| 519 | } |
| 520 | fn default_read_timeout() -> u64 { |
| 521 | 120 |
| 522 | } |
| 523 | |
| 524 | impl Default for McpTimeouts { |
| 525 | fn default() -> Self { |
| 526 | Self { |
| 527 | connect_timeout: default_connect_timeout(), |
| 528 | execute_timeout: default_execute_timeout(), |
| 529 | read_timeout: default_read_timeout(), |
| 530 | } |
| 531 | } |
| 532 | } |
| 533 | |
| 534 | /// Configuration for a single MCP server |
| 535 | #[derive(Debug, Clone, Deserialize, Serialize)] |
| 536 | pub struct McpServerConfig { |
| 537 | pub command: Option<String>, |
| 538 | #[serde(default)] |
| 539 | pub args: Vec<String>, |
| 540 | #[serde(default)] |
| 541 | pub env: HashMap<String, String>, |
| 542 | #[serde(default)] |
| 543 | #[serde(skip_serializing_if = "Option::is_none")] |
| 544 | pub cwd: Option<PathBuf>, |
| 545 | pub url: Option<String>, |
| 546 | /// Explicit operator authority for private DNS names at this exact origin. |
| 547 | /// Ignored for model-added runtime servers. |
| 548 | #[serde(default, skip_serializing_if = "std::ops::Not::not")] |
| 549 | pub allow_private_network: bool, |
| 550 | /// Optional explicit HTTP transport override. |
| 551 | /// |
| 552 | /// By default URL-based MCP servers use Streamable HTTP first and fall |
| 553 | /// back to legacy SSE only when the server rejects Streamable HTTP with |
| 554 | /// a known incompatible status. Set this to `"sse"` for legacy SSE |
| 555 | /// endpoints that must start with a long-lived GET endpoint discovery |
| 556 | /// stream and cannot accept an initial POST to the configured URL. |
| 557 | #[serde(default)] |
| 558 | #[serde(skip_serializing_if = "Option::is_none")] |
| 559 | pub transport: Option<String>, |
| 560 | #[serde(default)] |
| 561 | pub connect_timeout: Option<u64>, |
| 562 | #[serde(default)] |
| 563 | pub execute_timeout: Option<u64>, |
| 564 | #[serde(default)] |
| 565 | pub read_timeout: Option<u64>, |
| 566 | #[serde(default)] |
| 567 | pub disabled: bool, |
| 568 | #[serde(default = "default_enabled")] |
| 569 | pub enabled: bool, |
| 570 | #[serde(default)] |
| 571 | pub required: bool, |
| 572 | #[serde(default)] |
| 573 | pub enabled_tools: Vec<String>, |
| 574 | #[serde(default)] |
| 575 | pub disabled_tools: Vec<String>, |
| 576 | /// Extra HTTP headers sent with every request to this MCP server. |
| 577 | /// Only the HTTP transports (streamable HTTP today; SSE in a |
| 578 | /// follow-up) honor this — `command`-based stdio servers ignore it. |
| 579 | /// |
| 580 | /// Mirrors the `headers` field that Claude Code, Codex, and |
| 581 | /// OpenCode already accept in their MCP config formats. Use it to |
| 582 | /// authenticate against gateways that require a Bearer token or |
| 583 | /// API key, e.g.: |
| 584 | /// |
| 585 | /// ```jsonc |
| 586 | /// "huggingface": { |
| 587 | /// "url": "https://huggingface.co/api/mcp", |
| 588 | /// "headers": { "Authorization": "Bearer ${HF_TOKEN}" } |
| 589 | /// } |
| 590 | /// ``` |
| 591 | /// |
| 592 | /// Header keys and values are passed through as-is — we do not |
| 593 | /// substitute environment variables in v0.8.31. If you store a |
| 594 | /// real token here, the value lives in plain text in |
| 595 | /// `~/.deepseek/mcp.json`; treat that file with the same care |
| 596 | /// as any other secret-bearing config. |
| 597 | #[serde(default)] |
| 598 | #[serde(skip_serializing_if = "HashMap::is_empty")] |
| 599 | pub headers: HashMap<String, String>, |
| 600 | /// HTTP headers whose values are read from environment variables at request |
| 601 | /// time. This keeps common bearer/API-token integrations out of mcp.json. |
| 602 | #[serde(default, alias = "env_http_headers")] |
| 603 | #[serde(skip_serializing_if = "HashMap::is_empty")] |
| 604 | pub env_headers: HashMap<String, String>, |
| 605 | /// Environment variable containing a bearer token. When present and set, |
| 606 | /// CodeWhale sends `Authorization: Bearer <value>` for URL-based servers. |
| 607 | #[serde(default)] |
| 608 | #[serde(skip_serializing_if = "Option::is_none")] |
| 609 | pub bearer_token_env_var: Option<String>, |
| 610 | /// OAuth scopes requested during `codewhale mcp login`. |
| 611 | #[serde(default)] |
| 612 | #[serde(skip_serializing_if = "Vec::is_empty")] |
| 613 | pub scopes: Vec<String>, |
| 614 | /// OAuth client override for MCP servers that require a pre-registered |
| 615 | /// public client instead of dynamic registration. |
| 616 | #[serde(default)] |
| 617 | #[serde(skip_serializing_if = "Option::is_none")] |
| 618 | pub oauth: Option<McpServerOAuthConfig>, |
| 619 | /// Optional RFC 8707 resource parameter appended to the authorization URL. |
| 620 | #[serde(default)] |
| 621 | #[serde(skip_serializing_if = "Option::is_none")] |
| 622 | pub oauth_resource: Option<String>, |
| 623 | /// In-memory provenance for MCP servers contributed by a reviewed plugin |
| 624 | /// bundle. This is never deserialized from or serialized into user config: |
| 625 | /// only the trusted plugin merge adapter may attach it. |
| 626 | #[serde(skip)] |
| 627 | pub(crate) reviewed_plugin: Option<ReviewedPluginMcpSource>, |
| 628 | /// Only the runtime registration boundary can attach this provenance. |
| 629 | #[serde(skip)] |
| 630 | pub(crate) runtime_added: bool, |
| 631 | } |
| 632 | |
| 633 | #[derive(Debug, Clone)] |
| 634 | pub(crate) struct ReviewedPluginMcpSource { |
| 635 | authority: crate::plugins::types::PluginAuthority, |
| 636 | approved_remote_endpoint: Option<String>, |
| 637 | approved_remote_origin: Option<String>, |
| 638 | host_environment: Arc<crate::plugins::HostEnvironment>, |
| 639 | } |
| 640 | |
| 641 | impl ReviewedPluginMcpSource { |
| 642 | fn from_authority( |
| 643 | authority: crate::plugins::types::PluginAuthority, |
| 644 | remote_endpoint: Option<&str>, |
| 645 | host_environment: Arc<crate::plugins::HostEnvironment>, |
| 646 | ) -> Result<Self> { |
| 647 | let (approved_remote_endpoint, approved_remote_origin) = match remote_endpoint { |
| 648 | Some(endpoint) => reviewed_remote_endpoint_identity(endpoint) |
| 649 | .map(|(endpoint, origin)| (Some(endpoint), Some(origin)))?, |
| 650 | None => (None, None), |
| 651 | }; |
| 652 | Ok(Self { |
| 653 | authority, |
| 654 | approved_remote_endpoint, |
| 655 | approved_remote_origin, |
| 656 | host_environment, |
| 657 | }) |
| 658 | } |
| 659 | |
| 660 | pub(crate) fn validate_before_stdio_spawn(&self, server_name: &str) -> Result<()> { |
| 661 | self.validate_before_use(server_name, "spawn") |
| 662 | } |
| 663 | |
| 664 | pub(crate) fn prepare_stdio_launch( |
| 665 | &self, |
| 666 | server_name: &str, |
| 667 | command: &str, |
| 668 | args: &[String], |
| 669 | cwd: Option<&Path>, |
| 670 | ) -> Result<ReviewedStdioLaunch> { |
| 671 | self.validate_before_stdio_spawn(server_name)?; |
| 672 | let staged_root = self |
| 673 | .authority |
| 674 | .staged_manifest |
| 675 | .parent() |
| 676 | .context("reviewed plugin stage manifest has no parent")?; |
| 677 | let validated = crate::plugins::manifest::PluginManifest::validate_from_path( |
| 678 | &self.authority.staged_manifest, |
| 679 | ) |
| 680 | .map_err(|_| anyhow::anyhow!("reviewed plugin stage could not be opened for launch"))?; |
| 681 | if validated.content_hash != self.authority.content_hash |
| 682 | || validated.capability_hash != self.authority.capability_hash |
| 683 | { |
| 684 | anyhow::bail!("reviewed plugin stage changed before stdio launch"); |
| 685 | } |
| 686 | |
| 687 | let mut launch = ReviewedStdioLaunch { |
| 688 | command: std::ffi::OsString::from(command), |
| 689 | args: args.iter().map(std::ffi::OsString::from).collect(), |
| 690 | cwd: cwd.map(Path::to_path_buf), |
| 691 | opened_files: Vec::new(), |
| 692 | #[cfg(unix)] |
| 693 | cwd_fd: None, |
| 694 | }; |
| 695 | if Path::new(command).is_absolute() { |
| 696 | launch.bind_command(staged_root, Path::new(command), &validated.file_hashes)?; |
| 697 | } |
| 698 | // Darwin descriptor paths lose Node's module filename, package.json |
| 699 | // context and relative-import directory (#5916). Keep the staged path |
| 700 | // for .js/.cjs and multi-file .mjs entries, after bind_file verifies |
| 701 | // their bytes. Node reopens these paths; this is not atomic descriptor |
| 702 | // execution. Sibling modules already use the reviewed staged paths. |
| 703 | #[cfg(target_os = "macos")] |
| 704 | let node_entry_index = is_node_command(command) |
| 705 | .then(|| node_script_entry_index(args)) |
| 706 | .flatten() |
| 707 | .filter(|&index| { |
| 708 | let path = Path::new(&args[index]); |
| 709 | path.is_absolute() |
| 710 | && path.starts_with(staged_root) |
| 711 | && path.extension().is_some_and(|extension| { |
| 712 | matches!(extension.to_str(), Some("mjs" | "js" | "cjs")) |
| 713 | }) |
| 714 | }); |
| 715 | #[cfg(target_os = "macos")] |
| 716 | let node_entry_keeps_path = node_entry_index.is_some_and(|index| { |
| 717 | node_entry_needs_staged_path( |
| 718 | staged_root, |
| 719 | Path::new(&args[index]), |
| 720 | &validated.file_hashes, |
| 721 | ) |
| 722 | }); |
| 723 | for (index, argument) in args.iter().enumerate() { |
| 724 | let path = Path::new(argument); |
| 725 | if path.is_absolute() && path.starts_with(staged_root) && path.is_file() { |
| 726 | let bound = launch.bind_file(staged_root, path, &validated.file_hashes)?; |
| 727 | #[cfg(target_os = "macos")] |
| 728 | if node_entry_keeps_path && node_entry_index == Some(index) { |
| 729 | continue; |
| 730 | } |
| 731 | launch.args[index] = bound; |
| 732 | } |
| 733 | } |
| 734 | #[cfg(target_os = "macos")] |
| 735 | if let Some(entry_index) = node_entry_index |
| 736 | && !node_entry_keeps_path |
| 737 | { |
| 738 | launch.args = node_esm_descriptor_args(&launch.args, entry_index); |
| 739 | } |
| 740 | if let Some(cwd) = cwd { |
| 741 | if !cwd.starts_with(staged_root) { |
| 742 | anyhow::bail!("reviewed plugin stdio cwd escaped its staged root"); |
| 743 | } |
| 744 | launch.bind_cwd(cwd)?; |
| 745 | } |
| 746 | // A final authority pass detects source/stage and capability drift |
| 747 | // while handles were opened. Retained Node entry paths and imports |
| 748 | // are reopened after this check; their owner-only, read-only stage is |
| 749 | // not an atomic handle binding or an OS sandbox. |
| 750 | self.validate_before_stdio_spawn(server_name)?; |
| 751 | Ok(launch) |
| 752 | } |
| 753 | |
| 754 | fn required_capability(&self) -> crate::plugins::activation::PluginActivationCapability { |
| 755 | if self.approved_remote_endpoint.is_some() { |
| 756 | crate::plugins::activation::PluginActivationCapability::McpRemote |
| 757 | } else { |
| 758 | crate::plugins::activation::PluginActivationCapability::McpStdio |
| 759 | } |
| 760 | } |
| 761 | |
| 762 | fn validate_before_use(&self, server_name: &str, operation: &str) -> Result<()> { |
| 763 | let remediation = format!( |
| 764 | "Run `/plugin reload`, inspect `/plugin show {0}`, then repeat the displayed trust command and `/plugin enable {0}` before retrying", |
| 765 | self.authority.plugin_name |
| 766 | ); |
| 767 | crate::plugins::registry::verify_plugin_component_authority( |
| 768 | &self.authority, |
| 769 | self.required_capability(), |
| 770 | ) |
| 771 | .map_err(|reason| { |
| 772 | anyhow::anyhow!( |
| 773 | "Refusing to {operation} MCP server '{server_name}' from plugin bundle `{}`: {reason}. {remediation}", |
| 774 | self.authority.plugin_name |
| 775 | ) |
| 776 | }) |
| 777 | } |
| 778 | |
| 779 | fn validate_remote_endpoint(&self, server_name: &str, endpoint: &str) -> Result<()> { |
| 780 | let (endpoint, origin) = reviewed_remote_endpoint_identity(endpoint)?; |
| 781 | if self.approved_remote_endpoint.as_deref() != Some(endpoint.as_str()) |
| 782 | || self.approved_remote_origin.as_deref() != Some(origin.as_str()) |
| 783 | { |
| 784 | anyhow::bail!( |
| 785 | "Refusing MCP server '{server_name}': its remote endpoint no longer matches the reviewed plugin origin" |
| 786 | ); |
| 787 | } |
| 788 | Ok(()) |
| 789 | } |
| 790 | |
| 791 | fn catalog_is_current(&self) -> bool { |
| 792 | // Catalog exposure is an authority boundary too: stale tool, prompt, |
| 793 | // or resource descriptions can steer the model even when the later |
| 794 | // operation would be denied. Revalidate both the mutable reviewed |
| 795 | // source and the Codewhale-owned stage before publishing any entry. |
| 796 | crate::plugins::registry::verify_plugin_component_authority( |
| 797 | &self.authority, |
| 798 | self.required_capability(), |
| 799 | ) |
| 800 | .is_ok() |
| 801 | } |
| 802 | } |
| 803 | |
| 804 | /// Preserve .js package type lookup, .cjs module semantics and multi-file |
| 805 | /// .mjs relative imports. A lone .mjs retains the existing descriptor launch. |
| 806 | #[cfg(target_os = "macos")] |
| 807 | fn node_entry_needs_staged_path( |
| 808 | staged_root: &Path, |
| 809 | entry: &Path, |
| 810 | file_hashes: &std::collections::BTreeMap<PathBuf, String>, |
| 811 | ) -> bool { |
| 812 | let Ok(entry) = entry.strip_prefix(staged_root) else { |
| 813 | return false; |
| 814 | }; |
| 815 | if entry |
| 816 | .extension() |
| 817 | .is_some_and(|extension| matches!(extension.to_str(), Some("js" | "cjs"))) |
| 818 | { |
| 819 | return true; |
| 820 | } |
| 821 | file_hashes.keys().any(|path| { |
| 822 | path != entry |
| 823 | && path.extension().is_some_and(|extension| { |
| 824 | matches!( |
| 825 | extension.to_string_lossy().as_ref(), |
| 826 | "mjs" | "js" | "cjs" | "node" | "wasm" |
| 827 | ) |
| 828 | }) |
| 829 | }) |
| 830 | } |
| 831 | |
| 832 | /// Find the script operand, never a preload's value or an argument belonging |
| 833 | /// to an earlier script. Unknown option layouts get no Node-specific rewrite; |
| 834 | /// the generic reviewed-file binder still applies. |
| 835 | #[cfg(target_os = "macos")] |
| 836 | fn node_script_entry_index(args: &[impl AsRef<std::ffi::OsStr>]) -> Option<usize> { |
| 837 | let mut index = 0; |
| 838 | while let Some(argument) = args.get(index) { |
| 839 | let argument = argument.as_ref().to_str()?; |
| 840 | match argument { |
| 841 | "--" => return (index + 1 < args.len()).then_some(index + 1), |
| 842 | "-" | "-e" | "--eval" | "-p" | "--print" | "--run" | "--test" | "-c" | "--check" |
| 843 | | "-i" | "--interactive" | "--input-type" => return None, |
| 844 | "-r" |
| 845 | | "--require" |
| 846 | | "--import" |
| 847 | | "--loader" |
| 848 | | "--experimental-loader" |
| 849 | | "-C" |
| 850 | | "--conditions" |
| 851 | | "--max-old-space-size" |
| 852 | | "--stack-size" => index += 2, |
| 853 | "--no-warnings" |
| 854 | | "--trace-warnings" |
| 855 | | "--trace-deprecation" |
| 856 | | "--no-deprecation" |
| 857 | | "--enable-source-maps" |
| 858 | | "--preserve-symlinks" |
| 859 | | "--preserve-symlinks-main" |
| 860 | | "--abort-on-uncaught-exception" |
| 861 | | "--expose-gc" |
| 862 | | "--jitless" => index += 1, |
| 863 | _ if argument.starts_with("--eval=") |
| 864 | || argument.starts_with("--print=") |
| 865 | || argument.starts_with("--run=") |
| 866 | || argument.starts_with("--input-type=") => |
| 867 | { |
| 868 | return None; |
| 869 | } |
| 870 | _ if argument.starts_with("--") && argument.contains('=') => index += 1, |
| 871 | _ if argument.starts_with('-') => return None, |
| 872 | _ => return Some(index), |
| 873 | } |
| 874 | } |
| 875 | None |
| 876 | } |
| 877 | |
| 878 | fn is_node_command(command: &str) -> bool { |
| 879 | Path::new(command) |
| 880 | .file_name() |
| 881 | .and_then(|name| name.to_str()) |
| 882 | .is_some_and(|name| matches!(name, "node" | "nodejs" | "node.exe" | "nodejs.exe")) |
| 883 | } |
| 884 | |
| 885 | /// Rewrite a Node launch so a reviewed `.mjs` entrypoint keeps ESM semantics |
| 886 | /// after Darwin's reviewed-launch binding replaced its staged path with an |
| 887 | /// inherited `/dev/fd/N` descriptor. |
| 888 | /// |
| 889 | /// Node determines the entrypoint module type from its filename and a |
| 890 | /// descriptor path has no extension, so `node /dev/fd/N` exits without |
| 891 | /// evaluating the module. `--experimental-default-type=module` used to fix |
| 892 | /// that but was removed from current Node releases (Node 25 rejects it as a |
| 893 | /// bad option, which killed the child before the MCP handshake). `--import` |
| 894 | /// has loaded its specifier as an ES module on every supported release, so |
| 895 | /// the reviewed bytes are imported by descriptor once, `-e ""` supplies an |
| 896 | /// empty main, and the descriptor path is echoed after `--` so |
| 897 | /// `process.argv[1]` and the script's own arguments keep the file-mode shape. |
| 898 | /// Node options that preceded the entrypoint stay in front; anything after it |
| 899 | /// is passed through untouched. When the `.mjs` file is not the first |
| 900 | /// positional argument it is not the entrypoint and the launch is left alone. |
| 901 | #[cfg(target_os = "macos")] |
| 902 | fn node_esm_descriptor_args( |
| 903 | args: &[std::ffi::OsString], |
| 904 | entry_index: usize, |
| 905 | ) -> Vec<std::ffi::OsString> { |
| 906 | if node_script_entry_index(args) != Some(entry_index) { |
| 907 | return args.to_vec(); |
| 908 | } |
| 909 | let bound_entry = args[entry_index].clone(); |
| 910 | let prefix_end = entry_index - usize::from(entry_index > 0 && args[entry_index - 1] == "--"); |
| 911 | let mut rewritten: Vec<std::ffi::OsString> = args[..prefix_end].to_vec(); |
| 912 | rewritten.push(std::ffi::OsString::from("--import")); |
| 913 | rewritten.push(bound_entry.clone()); |
| 914 | rewritten.push(std::ffi::OsString::from("-e")); |
| 915 | rewritten.push(std::ffi::OsString::from("")); |
| 916 | rewritten.push(std::ffi::OsString::from("--")); |
| 917 | rewritten.push(bound_entry); |
| 918 | rewritten.extend(args[entry_index + 1..].iter().cloned()); |
| 919 | rewritten |
| 920 | } |
| 921 | |
| 922 | pub(crate) struct ReviewedStdioLaunch { |
| 923 | pub(crate) command: std::ffi::OsString, |
| 924 | pub(crate) args: Vec<std::ffi::OsString>, |
| 925 | pub(crate) cwd: Option<PathBuf>, |
| 926 | /// Kept for the child lifetime. Windows opens deny write/delete sharing; |
| 927 | /// Unix normally uses inherited descriptors; macOS Node entries needing |
| 928 | /// module path context are hash-checked here, then reopened by path. |
| 929 | pub(crate) opened_files: Vec<fs::File>, |
| 930 | #[cfg(unix)] |
| 931 | pub(crate) cwd_fd: Option<fs::File>, |
| 932 | } |
| 933 | |
| 934 | impl ReviewedStdioLaunch { |
| 935 | fn bind_command( |
| 936 | &mut self, |
| 937 | staged_root: &Path, |
| 938 | path: &Path, |
| 939 | expected_hashes: &std::collections::BTreeMap<PathBuf, String>, |
| 940 | ) -> Result<()> { |
| 941 | let bound_path = self.bind_file(staged_root, path, expected_hashes)?; |
| 942 | #[cfg(not(target_os = "macos"))] |
| 943 | { |
| 944 | self.command = bound_path; |
| 945 | Ok(()) |
| 946 | } |
| 947 | #[cfg(target_os = "macos")] |
| 948 | { |
| 949 | use std::os::unix::fs::FileExt as _; |
| 950 | |
| 951 | // Darwin devfs deliberately rejects execve("/dev/fd/N"). Bind |
| 952 | // reviewed scripts by running the interpreter declared in their |
| 953 | // exact hashed shebang and passing the inherited descriptor as |
| 954 | // input. Native Mach-O bundle commands have no fexecve/execveat |
| 955 | // equivalent on Darwin, so fail closed and require the manifest |
| 956 | // to name a bare interpreter with the bundle file as an argument. |
| 957 | let file = self |
| 958 | .opened_files |
| 959 | .last() |
| 960 | .context("reviewed command handle disappeared")?; |
| 961 | let mut prefix = [0_u8; 4_096]; |
| 962 | let read = file |
| 963 | .read_at(&mut prefix, 0) |
| 964 | .context("read reviewed command shebang")?; |
| 965 | let prefix = &prefix[..read]; |
| 966 | let line_end = prefix |
| 967 | .iter() |
| 968 | .position(|byte| *byte == b'\n') |
| 969 | .unwrap_or(prefix.len()); |
| 970 | let line = std::str::from_utf8(&prefix[..line_end]) |
| 971 | .context("reviewed script shebang is not UTF-8")?; |
| 972 | let shebang = line.strip_prefix("#!").map(str::trim).filter(|s| !s.is_empty()) |
| 973 | .context( |
| 974 | "Darwin cannot execute a reviewed native bundle command by descriptor; use a shebang script or declare a bare interpreter command plus the script argument", |
| 975 | )?; |
| 976 | let mut words = shlex::split(shebang) |
| 977 | .context("reviewed script shebang could not be parsed safely")?; |
| 978 | let interpreter = words |
| 979 | .first() |
| 980 | .filter(|word| Path::new(word).is_absolute()) |
| 981 | .context("reviewed script shebang interpreter must be absolute")? |
| 982 | .clone(); |
| 983 | words.remove(0); |
| 984 | let mut args = words |
| 985 | .into_iter() |
| 986 | .map(std::ffi::OsString::from) |
| 987 | .collect::<Vec<_>>(); |
| 988 | args.push(bound_path); |
| 989 | args.append(&mut self.args); |
| 990 | self.command = std::ffi::OsString::from(interpreter); |
| 991 | self.args = args; |
| 992 | Ok(()) |
| 993 | } |
| 994 | } |
| 995 | |
| 996 | fn bind_file( |
| 997 | &mut self, |
| 998 | staged_root: &Path, |
| 999 | path: &Path, |
| 1000 | expected_hashes: &std::collections::BTreeMap<PathBuf, String>, |
| 1001 | ) -> Result<std::ffi::OsString> { |
| 1002 | let relative = path |
| 1003 | .strip_prefix(staged_root) |
| 1004 | .context("reviewed plugin executable escaped its staged root")?; |
| 1005 | let expected = expected_hashes |
| 1006 | .get(relative) |
| 1007 | .context("reviewed plugin executable is absent from its byte inventory")?; |
| 1008 | let mut file = open_reviewed_launch_file(path)?; |
| 1009 | let mut hasher = sha2::Sha256::new(); |
| 1010 | hasher.update(b"codewhale-plugin-file-bytes-v1\0"); |
| 1011 | let mut buffer = [0_u8; 64 * 1024]; |
| 1012 | loop { |
| 1013 | let read = file |
| 1014 | .read(&mut buffer) |
| 1015 | .context("read reviewed launch file")?; |
| 1016 | if read == 0 { |
| 1017 | break; |
| 1018 | } |
| 1019 | hasher.update(&buffer[..read]); |
| 1020 | } |
| 1021 | let actual = hasher |
| 1022 | .finalize() |
| 1023 | .iter() |
| 1024 | .map(|byte| format!("{byte:02x}")) |
| 1025 | .collect::<String>(); |
| 1026 | if &actual != expected { |
| 1027 | anyhow::bail!("reviewed plugin executable bytes changed before spawn"); |
| 1028 | } |
| 1029 | file.seek(std::io::SeekFrom::Start(0)) |
| 1030 | .context("rewind reviewed launch file after verification")?; |
| 1031 | |
| 1032 | #[cfg(unix)] |
| 1033 | let launch_path = { |
| 1034 | use std::os::fd::AsRawFd as _; |
| 1035 | let fd = file.as_raw_fd(); |
| 1036 | // SAFETY: `fd` is owned by `file`; clearing only FD_CLOEXEC keeps |
| 1037 | // that same descriptor available across the imminent exec. |
| 1038 | let flags = unsafe { libc::fcntl(fd, libc::F_GETFD) }; |
| 1039 | if flags < 0 || unsafe { libc::fcntl(fd, libc::F_SETFD, flags & !libc::FD_CLOEXEC) } < 0 |
| 1040 | { |
| 1041 | anyhow::bail!("failed to inherit reviewed plugin executable descriptor"); |
| 1042 | } |
| 1043 | #[cfg(target_os = "linux")] |
| 1044 | let prefix = "/proc/self/fd"; |
| 1045 | #[cfg(not(target_os = "linux"))] |
| 1046 | let prefix = "/dev/fd"; |
| 1047 | std::ffi::OsString::from(format!("{prefix}/{fd}")) |
| 1048 | }; |
| 1049 | |
| 1050 | #[cfg(not(unix))] |
| 1051 | let launch_path = path.as_os_str().to_os_string(); |
| 1052 | |
| 1053 | self.opened_files.push(file); |
| 1054 | Ok(launch_path) |
| 1055 | } |
| 1056 | |
| 1057 | fn bind_cwd(&mut self, cwd: &Path) -> Result<()> { |
| 1058 | #[cfg(unix)] |
| 1059 | { |
| 1060 | use std::os::unix::fs::OpenOptionsExt as _; |
| 1061 | let file = fs::OpenOptions::new() |
| 1062 | .read(true) |
| 1063 | .custom_flags(libc::O_DIRECTORY | libc::O_NOFOLLOW | libc::O_CLOEXEC) |
| 1064 | .open(cwd) |
| 1065 | .context("open reviewed plugin cwd without following links")?; |
| 1066 | self.cwd_fd = Some(file); |
| 1067 | self.cwd = None; |
| 1068 | } |
| 1069 | #[cfg(windows)] |
| 1070 | { |
| 1071 | use std::os::windows::fs::{MetadataExt as _, OpenOptionsExt as _}; |
| 1072 | let file = fs::OpenOptions::new() |
| 1073 | .read(true) |
| 1074 | .share_mode(0x0000_0001) // FILE_SHARE_READ only |
| 1075 | .custom_flags(0x0220_0000) // BACKUP_SEMANTICS | OPEN_REPARSE_POINT |
| 1076 | .open(cwd) |
| 1077 | .context("open reviewed plugin cwd without write/delete sharing")?; |
| 1078 | let metadata = file |
| 1079 | .metadata() |
| 1080 | .context("inspect reviewed plugin cwd handle")?; |
| 1081 | if !metadata.is_dir() || metadata.file_attributes() & 0x0000_0400 != 0 { |
| 1082 | anyhow::bail!("reviewed plugin cwd is a reparse point or non-directory"); |
| 1083 | } |
| 1084 | self.opened_files.push(file); |
| 1085 | } |
| 1086 | Ok(()) |
| 1087 | } |
| 1088 | } |
| 1089 | |
| 1090 | #[cfg(unix)] |
| 1091 | fn open_reviewed_launch_file(path: &Path) -> Result<fs::File> { |
| 1092 | crate::plugins::manifest::open_bundle_file(path) |
| 1093 | .context("open reviewed launch file without following links") |
| 1094 | } |
| 1095 | |
| 1096 | #[cfg(windows)] |
| 1097 | fn open_reviewed_launch_file(path: &Path) -> Result<fs::File> { |
| 1098 | crate::plugins::manifest::open_bundle_file(path) |
| 1099 | .context("open reviewed launch file without links, hard links, or write/delete sharing") |
| 1100 | } |
| 1101 | |
| 1102 | #[cfg(all(not(unix), not(windows)))] |
| 1103 | fn open_reviewed_launch_file(path: &Path) -> Result<fs::File> { |
| 1104 | fs::File::open(path).context("open reviewed launch file") |
| 1105 | } |
| 1106 | |
| 1107 | fn reviewed_remote_endpoint_identity(endpoint: &str) -> Result<(String, String)> { |
| 1108 | let endpoint = |
| 1109 | reqwest::Url::parse(endpoint).context("reviewed plugin MCP endpoint is invalid")?; |
| 1110 | if !endpoint.username().is_empty() || endpoint.password().is_some() { |
| 1111 | anyhow::bail!("reviewed plugin MCP endpoint must not contain user information"); |
| 1112 | } |
| 1113 | if endpoint.query().is_some() || endpoint.fragment().is_some() { |
| 1114 | anyhow::bail!("reviewed plugin MCP endpoint must not contain a query or fragment"); |
| 1115 | } |
| 1116 | let origin = reviewed_remote_origin(&endpoint) |
| 1117 | .ok_or_else(|| anyhow::anyhow!("reviewed plugin MCP endpoint has an unsafe origin"))?; |
| 1118 | Ok((endpoint.to_string(), origin)) |
| 1119 | } |
| 1120 | |
| 1121 | fn reviewed_remote_origin(endpoint: &reqwest::Url) -> Option<String> { |
| 1122 | if !endpoint.username().is_empty() || endpoint.password().is_some() { |
| 1123 | return None; |
| 1124 | } |
| 1125 | let host = endpoint.host_str()?; |
| 1126 | let allowed_scheme = endpoint.scheme() == "https" |
| 1127 | || (endpoint.scheme() == "http" |
| 1128 | && (host.eq_ignore_ascii_case("localhost") |
| 1129 | || host |
| 1130 | .trim_matches(['[', ']']) |
| 1131 | .parse::<std::net::IpAddr>() |
| 1132 | .is_ok_and(|address| address.is_loopback()))); |
| 1133 | allowed_scheme.then(|| endpoint.origin().ascii_serialization()) |
| 1134 | } |
| 1135 | |
| 1136 | fn reviewed_redirect_matches_origin(endpoint: &reqwest::Url, approved_origin: &str) -> bool { |
| 1137 | reviewed_remote_origin(endpoint).as_deref() == Some(approved_origin) |
| 1138 | } |
| 1139 | |
| 1140 | #[derive(Debug, Clone, Default, Deserialize, Serialize, PartialEq, Eq)] |
| 1141 | pub struct McpServerOAuthConfig { |
| 1142 | #[serde(default)] |
| 1143 | #[serde(skip_serializing_if = "Option::is_none")] |
| 1144 | pub client_id: Option<String>, |
| 1145 | } |
| 1146 | |
| 1147 | fn default_enabled() -> bool { |
| 1148 | true |
| 1149 | } |
| 1150 | |
| 1151 | impl McpServerConfig { |
| 1152 | pub fn effective_connect_timeout(&self, global: &McpTimeouts) -> u64 { |
| 1153 | self.connect_timeout.unwrap_or(global.connect_timeout) |
| 1154 | } |
| 1155 | |
| 1156 | pub fn effective_execute_timeout(&self, global: &McpTimeouts) -> u64 { |
| 1157 | self.execute_timeout.unwrap_or(global.execute_timeout) |
| 1158 | } |
| 1159 | |
| 1160 | pub fn effective_read_timeout(&self, global: &McpTimeouts) -> u64 { |
| 1161 | self.read_timeout.unwrap_or(global.read_timeout) |
| 1162 | } |
| 1163 | |
| 1164 | pub fn is_enabled(&self) -> bool { |
| 1165 | self.enabled && !self.disabled |
| 1166 | } |
| 1167 | |
| 1168 | pub fn is_tool_enabled(&self, tool_name: &str) -> bool { |
| 1169 | let allowed = if self.enabled_tools.is_empty() { |
| 1170 | true |
| 1171 | } else { |
| 1172 | self.enabled_tools.iter().any(|t| t == tool_name) |
| 1173 | }; |
| 1174 | if !allowed { |
| 1175 | return false; |
| 1176 | } |
| 1177 | !self.disabled_tools.iter().any(|t| t == tool_name) |
| 1178 | } |
| 1179 | } |
| 1180 | |
| 1181 | // === MCP Tool Definition === |
| 1182 | |
| 1183 | /// Tool discovered from an MCP server |
| 1184 | #[derive(Debug, Clone, Deserialize, Serialize)] |
| 1185 | pub struct McpTool { |
| 1186 | pub name: String, |
| 1187 | #[serde(default)] |
| 1188 | pub description: Option<String>, |
| 1189 | #[serde(rename = "inputSchema", default)] |
| 1190 | pub input_schema: serde_json::Value, |
| 1191 | } |
| 1192 | |
| 1193 | const MCP_TOOL_DESCRIPTION_MAX_CHARS: usize = 80; |
| 1194 | |
| 1195 | /// Format an optional MCP tool description for terminal list surfaces. |
| 1196 | /// |
| 1197 | /// CLI and TUI callers share this helper so both stay single-line and truncate |
| 1198 | /// on Unicode scalar boundaries rather than slicing UTF-8 bytes. |
| 1199 | pub(crate) fn format_mcp_tool_description(description: Option<&str>) -> String { |
| 1200 | let Some(first_line) = description |
| 1201 | .and_then(|description| description.split(['\r', '\n']).next()) |
| 1202 | .map(str::trim) |
| 1203 | .filter(|description| !description.is_empty()) |
| 1204 | else { |
| 1205 | return String::new(); |
| 1206 | }; |
| 1207 | |
| 1208 | let mut chars = first_line.chars(); |
| 1209 | let summary: String = chars |
| 1210 | .by_ref() |
| 1211 | .take(MCP_TOOL_DESCRIPTION_MAX_CHARS) |
| 1212 | .collect(); |
| 1213 | if chars.next().is_some() { |
| 1214 | format!(": {summary}...") |
| 1215 | } else { |
| 1216 | format!(": {summary}") |
| 1217 | } |
| 1218 | } |
| 1219 | |
| 1220 | /// Resource discovered from an MCP server |
| 1221 | #[derive(Debug, Clone, Deserialize, Serialize)] |
| 1222 | pub struct McpResource { |
| 1223 | pub uri: String, |
| 1224 | pub name: String, |
| 1225 | #[serde(default)] |
| 1226 | pub description: Option<String>, |
| 1227 | #[serde(rename = "mimeType", default)] |
| 1228 | pub mime_type: Option<String>, |
| 1229 | } |
| 1230 | |
| 1231 | /// Resource template discovered from an MCP server |
| 1232 | #[derive(Debug, Clone, Deserialize, Serialize)] |
| 1233 | pub struct McpResourceTemplate { |
| 1234 | #[serde(rename = "uriTemplate")] |
| 1235 | pub uri_template: String, |
| 1236 | pub name: String, |
| 1237 | #[serde(default)] |
| 1238 | pub description: Option<String>, |
| 1239 | #[serde(rename = "mimeType", default)] |
| 1240 | pub mime_type: Option<String>, |
| 1241 | } |
| 1242 | |
| 1243 | /// Fail-closed RFC 6570 subset used only as an authorization check. Literal, |
| 1244 | /// simple (`{id}`), and reserved (`{+path}`) expansions cover the common MCP |
| 1245 | /// resource templates. More elaborate operators remain listable but are not |
| 1246 | /// callable until their expansion semantics are implemented exactly. |
| 1247 | /// |
| 1248 | /// `None` is the fail-closed answer: a template this subset cannot express |
| 1249 | /// matches nothing. |
| 1250 | fn resource_template_pattern(template: &str) -> Option<String> { |
| 1251 | let mut pattern = String::from("^"); |
| 1252 | let mut rest = template; |
| 1253 | while let Some(start) = rest.find('{') { |
| 1254 | pattern.push_str(®ex::escape(&rest[..start])); |
| 1255 | let end = rest[start + 1..].find('}')?; |
| 1256 | let expression = &rest[start + 1..start + 1 + end]; |
| 1257 | let (reserved, variables) = match expression.strip_prefix('+') { |
| 1258 | Some(variables) => (true, variables), |
| 1259 | None => (false, expression), |
| 1260 | }; |
| 1261 | if variables.is_empty() |
| 1262 | || variables.split(',').any(|variable| { |
| 1263 | variable.is_empty() |
| 1264 | || !variable |
| 1265 | .chars() |
| 1266 | .all(|ch| ch.is_ascii_alphanumeric() || matches!(ch, '_' | '-' | '.')) |
| 1267 | }) |
| 1268 | { |
| 1269 | return None; |
| 1270 | } |
| 1271 | let atom = if reserved { ".+" } else { "[^/?#]+" }; |
| 1272 | for (index, _) in variables.split(',').enumerate() { |
| 1273 | if index > 0 { |
| 1274 | pattern.push(','); |
| 1275 | } |
| 1276 | pattern.push_str(atom); |
| 1277 | } |
| 1278 | rest = &rest[start + end + 2..]; |
| 1279 | } |
| 1280 | if rest.contains('}') { |
| 1281 | return None; |
| 1282 | } |
| 1283 | pattern.push_str(®ex::escape(rest)); |
| 1284 | pattern.push('$'); |
| 1285 | Some(pattern) |
| 1286 | } |
| 1287 | |
| 1288 | /// `template`'s anchored pattern, compiled once and reused. |
| 1289 | /// |
| 1290 | /// This runs per URI per advertised template, while the template itself is |
| 1291 | /// fixed by the server's listing, so compiling it on every call was pure |
| 1292 | /// repetition. `None` still means "matches nothing" (#6213 T7). |
| 1293 | fn compiled_resource_template(template: &str) -> Option<Arc<regex::Regex>> { |
| 1294 | static CACHE: std::sync::OnceLock< |
| 1295 | std::sync::Mutex<HashMap<String, Option<Arc<regex::Regex>>>>, |
| 1296 | > = std::sync::OnceLock::new(); |
| 1297 | let cache = CACHE.get_or_init(|| std::sync::Mutex::new(HashMap::new())); |
| 1298 | let mut cache = cache |
| 1299 | .lock() |
| 1300 | .unwrap_or_else(|poisoned| poisoned.into_inner()); |
| 1301 | cache |
| 1302 | .entry(template.to_string()) |
| 1303 | .or_insert_with(|| { |
| 1304 | resource_template_pattern(template) |
| 1305 | .and_then(|pattern| regex::Regex::new(&pattern).ok()) |
| 1306 | .map(Arc::new) |
| 1307 | }) |
| 1308 | .clone() |
| 1309 | } |
| 1310 | |
| 1311 | fn resource_uri_matches_template(uri: &str, template: &str) -> bool { |
| 1312 | compiled_resource_template(template).is_some_and(|regex| regex.is_match(uri)) |
| 1313 | } |
| 1314 | |
| 1315 | /// Prompt discovered from an MCP server |
| 1316 | #[derive(Debug, Clone, Deserialize, Serialize)] |
| 1317 | pub struct McpPrompt { |
| 1318 | pub name: String, |
| 1319 | #[serde(default)] |
| 1320 | pub description: Option<String>, |
| 1321 | #[serde(default)] |
| 1322 | pub arguments: Vec<McpPromptArgument>, |
| 1323 | } |
| 1324 | |
| 1325 | /// Argument for an MCP prompt |
| 1326 | #[derive(Debug, Clone, Deserialize, Serialize)] |
| 1327 | pub struct McpPromptArgument { |
| 1328 | pub name: String, |
| 1329 | #[serde(default)] |
| 1330 | pub description: Option<String>, |
| 1331 | #[serde(default)] |
| 1332 | pub required: bool, |
| 1333 | } |
| 1334 | |
| 1335 | // === Connection State === |
| 1336 | |
| 1337 | /// State of an MCP connection |
| 1338 | #[derive(Debug, Clone, Copy, PartialEq, Eq)] |
| 1339 | pub enum ConnectionState { |
| 1340 | Connecting, |
| 1341 | Ready, |
| 1342 | Disconnected, |
| 1343 | } |
| 1344 | |
| 1345 | /// MCP server capabilities advertised in the initialize response. |
| 1346 | /// |
| 1347 | /// Each flag records presence of the corresponding MCP capability object. The |
| 1348 | /// surrounding [`McpServerCapabilityMetadata`] preserves the important |
| 1349 | /// distinction between an advertised empty set and a legacy server that did |
| 1350 | /// not send capability metadata at all. |
| 1351 | #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] |
| 1352 | pub struct McpServerCapabilities { |
| 1353 | pub tools: bool, |
| 1354 | pub resources: bool, |
| 1355 | pub prompts: bool, |
| 1356 | } |
| 1357 | |
| 1358 | impl McpServerCapabilities { |
| 1359 | fn from_initialize_response(response: &serde_json::Value) -> Option<Self> { |
| 1360 | let capabilities = response.get("result")?.get("capabilities")?.as_object()?; |
| 1361 | Some(Self { |
| 1362 | tools: capabilities.contains_key("tools"), |
| 1363 | resources: capabilities.contains_key("resources"), |
| 1364 | prompts: capabilities.contains_key("prompts"), |
| 1365 | }) |
| 1366 | } |
| 1367 | } |
| 1368 | |
| 1369 | /// Provenance-aware capability metadata for a manager snapshot. |
| 1370 | #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] |
| 1371 | pub enum McpServerCapabilityMetadata { |
| 1372 | /// The server supplied a spec-shaped `capabilities` object at initialize. |
| 1373 | Advertised(McpServerCapabilities), |
| 1374 | /// A connected legacy server omitted metadata, so bounded discovery probes |
| 1375 | /// remain enabled for backward compatibility. |
| 1376 | LegacyFallback, |
| 1377 | /// The server has not completed initialization, so no truthful capability |
| 1378 | /// claim can be made yet. |
| 1379 | #[default] |
| 1380 | NotObserved, |
| 1381 | } |
| 1382 | |
| 1383 | fn response_result<'a>( |
| 1384 | response: &'a serde_json::Value, |
| 1385 | method: &str, |
| 1386 | suppress_server_details: bool, |
| 1387 | ) -> Result<Option<&'a serde_json::Value>> { |
| 1388 | if let Some(error) = response.get("error") { |
| 1389 | if suppress_server_details { |
| 1390 | anyhow::bail!( |
| 1391 | "Reviewed plugin MCP server returned an error in '{method}' (server details suppressed to protect environment-backed credentials)" |
| 1392 | ); |
| 1393 | } |
| 1394 | anyhow::bail!("MCP error in '{method}': {error}"); |
| 1395 | } |
| 1396 | Ok(response.get("result")) |
| 1397 | } |
| 1398 | |
| 1399 | async fn run_optional_discovery<F>(server: &str, method: &str, timeout: Duration, discovery: F) |
| 1400 | where |
| 1401 | F: Future<Output = Result<()>>, |
| 1402 | { |
| 1403 | match tokio::time::timeout(timeout, discovery).await { |
| 1404 | Ok(Ok(())) => {} |
| 1405 | Ok(Err(error)) => { |
| 1406 | tracing::warn!( |
| 1407 | target: "mcp", |
| 1408 | server, |
| 1409 | method, |
| 1410 | error = %error, |
| 1411 | "optional MCP discovery failed; continuing with available capabilities" |
| 1412 | ); |
| 1413 | } |
| 1414 | Err(error) => { |
| 1415 | tracing::warn!( |
| 1416 | target: "mcp", |
| 1417 | server, |
| 1418 | method, |
| 1419 | ?timeout, |
| 1420 | error = %error, |
| 1421 | "optional MCP discovery timed out; continuing with available capabilities" |
| 1422 | ); |
| 1423 | } |
| 1424 | } |
| 1425 | } |
| 1426 | |
| 1427 | // === McpConnection - Async Connection Management === |
| 1428 | |
| 1429 | // === Transport Trait === |
| 1430 | |
| 1431 | #[async_trait::async_trait] |
| 1432 | pub trait McpTransport: Send + Sync { |
| 1433 | async fn send(&mut self, msg: Vec<u8>) -> Result<()>; |
| 1434 | async fn recv(&mut self) -> Result<Vec<u8>>; |
| 1435 | |
| 1436 | /// Record the protocol revision negotiated at `initialize`. Only the |
| 1437 | /// Streamable HTTP transport uses it (the `MCP-Protocol-Version` header on |
| 1438 | /// subsequent requests); stdio and legacy SSE have no header channel, so |
| 1439 | /// the default is a no-op. |
| 1440 | fn set_protocol_version(&mut self, _version: &str) {} |
| 1441 | |
| 1442 | /// Synchronous, best-effort liveness probe consulted by |
| 1443 | /// [`McpConnection::is_ready`] so a crashed stdio child stops reading |
| 1444 | /// as "ready" before the next call fails (#6187). Must never block and |
| 1445 | /// never spawn — a contended lock reads as alive; the next call observes |
| 1446 | /// the death. HTTP/SSE transports have no child to observe, so the |
| 1447 | /// default is "alive". |
| 1448 | fn probe_dead(&self) -> bool { |
| 1449 | false |
| 1450 | } |
| 1451 | |
| 1452 | /// Graceful shutdown — stdio transports send SIGTERM to the child and |
| 1453 | /// give it a brief window to exit before tokio's `kill_on_drop` fires |
| 1454 | /// SIGKILL as the backstop. Default is a no-op for non-stdio transports |
| 1455 | /// that have no child process. Whalescale#420. |
| 1456 | async fn shutdown(&mut self) {} |
| 1457 | } |
| 1458 | |
| 1459 | const MAX_MCP_CATALOG_PAGES: usize = 64; |
| 1460 | const MAX_MCP_CATALOG_ITEMS: usize = 4_096; |
| 1461 | const MAX_MCP_CATALOG_BYTES: usize = 32 * 1024 * 1024; |
| 1462 | |
| 1463 | struct McpCatalogBudget { |
| 1464 | method: &'static str, |
| 1465 | pages: usize, |
| 1466 | items: usize, |
| 1467 | bytes: usize, |
| 1468 | seen_cursors: HashSet<String>, |
| 1469 | } |
| 1470 | |
| 1471 | impl McpCatalogBudget { |
| 1472 | fn new(method: &'static str) -> Self { |
| 1473 | Self { |
| 1474 | method, |
| 1475 | pages: 0, |
| 1476 | items: 0, |
| 1477 | bytes: 0, |
| 1478 | seen_cursors: HashSet::new(), |
| 1479 | } |
| 1480 | } |
| 1481 | |
| 1482 | fn observe_page( |
| 1483 | &mut self, |
| 1484 | result: &serde_json::Value, |
| 1485 | item_count: usize, |
| 1486 | ) -> Result<Option<String>> { |
| 1487 | self.pages = self.pages.saturating_add(1); |
| 1488 | self.items = self.items.saturating_add(item_count); |
| 1489 | self.bytes = self.bytes.saturating_add(serde_json::to_vec(result)?.len()); |
| 1490 | if self.pages > MAX_MCP_CATALOG_PAGES { |
| 1491 | anyhow::bail!( |
| 1492 | "{} exceeded the {}-page catalogue limit", |
| 1493 | self.method, |
| 1494 | MAX_MCP_CATALOG_PAGES |
| 1495 | ); |
| 1496 | } |
| 1497 | if self.items > MAX_MCP_CATALOG_ITEMS { |
| 1498 | anyhow::bail!( |
| 1499 | "{} exceeded the {}-item catalogue limit", |
| 1500 | self.method, |
| 1501 | MAX_MCP_CATALOG_ITEMS |
| 1502 | ); |
| 1503 | } |
| 1504 | if self.bytes > MAX_MCP_CATALOG_BYTES { |
| 1505 | anyhow::bail!( |
| 1506 | "{} exceeded the {}-byte aggregate catalogue limit", |
| 1507 | self.method, |
| 1508 | MAX_MCP_CATALOG_BYTES |
| 1509 | ); |
| 1510 | } |
| 1511 | let cursor = result |
| 1512 | .get("nextCursor") |
| 1513 | .and_then(|value| value.as_str()) |
| 1514 | .map(str::to_owned); |
| 1515 | if let Some(cursor) = cursor.as_ref() |
| 1516 | && !self.seen_cursors.insert(cursor.clone()) |
| 1517 | { |
| 1518 | anyhow::bail!("{} repeated pagination cursor; aborting", self.method); |
| 1519 | } |
| 1520 | Ok(cursor) |
| 1521 | } |
| 1522 | } |
| 1523 | |
| 1524 | fn is_legacy_sse_transport(config: &McpServerConfig) -> bool { |
| 1525 | config |
| 1526 | .transport |
| 1527 | .as_deref() |
| 1528 | .map(|transport| transport.trim().eq_ignore_ascii_case("sse")) |
| 1529 | .unwrap_or(false) |
| 1530 | } |
| 1531 | |
| 1532 | pub fn validate_mcp_transport(transport: Option<&str>) -> Result<()> { |
| 1533 | let Some(transport) = transport else { |
| 1534 | return Ok(()); |
| 1535 | }; |
| 1536 | if transport.trim().eq_ignore_ascii_case("sse") { |
| 1537 | return Ok(()); |
| 1538 | } |
| 1539 | anyhow::bail!("Unsupported MCP transport '{transport}'. Supported values: sse"); |
| 1540 | } |
| 1541 | |
| 1542 | fn response_id_matches(id: Option<&serde_json::Value>, expected_id: &str) -> bool { |
| 1543 | let Some(id) = id else { |
| 1544 | return false; |
| 1545 | }; |
| 1546 | if id.as_str() == Some(expected_id) { |
| 1547 | return true; |
| 1548 | } |
| 1549 | id.as_u64() |
| 1550 | .map(|id| id.to_string() == expected_id) |
| 1551 | .unwrap_or(false) |
| 1552 | } |
| 1553 | |
| 1554 | // === McpConnection - Async Connection Management === |
| 1555 | |
| 1556 | /// Manages a single async connection to an MCP server |
| 1557 | pub struct McpConnection { |
| 1558 | name: String, |
| 1559 | transport: Box<dyn McpTransport>, |
| 1560 | tools: Vec<McpTool>, |
| 1561 | resources: Vec<McpResource>, |
| 1562 | resource_templates: Vec<McpResourceTemplate>, |
| 1563 | prompts: Vec<McpPrompt>, |
| 1564 | request_id: AtomicU64, |
| 1565 | state: ConnectionState, |
| 1566 | config: McpServerConfig, |
| 1567 | server_capabilities: Option<McpServerCapabilities>, |
| 1568 | discovery_timeout: Duration, |
| 1569 | read_timeout_secs: u64, |
| 1570 | cancel_token: tokio_util::sync::CancellationToken, |
| 1571 | authority_revocation_reason: Arc<std::sync::Mutex<Option<String>>>, |
| 1572 | authority_watch: Option<tokio::task::JoinHandle<()>>, |
| 1573 | /// Pool catalog generation that created/last authorized this connection. |
| 1574 | /// Directly constructed test connections use zero until inserted. |
| 1575 | catalog_generation: u64, |
| 1576 | } |
| 1577 | |
| 1578 | struct PendingAuthorityWatch { |
| 1579 | handle: Option<tokio::task::JoinHandle<()>>, |
| 1580 | cancel: tokio_util::sync::CancellationToken, |
| 1581 | armed: bool, |
| 1582 | } |
| 1583 | |
| 1584 | impl PendingAuthorityWatch { |
| 1585 | fn start( |
| 1586 | source: ReviewedPluginMcpSource, |
| 1587 | cancel: tokio_util::sync::CancellationToken, |
| 1588 | reason_slot: Arc<std::sync::Mutex<Option<String>>>, |
| 1589 | ) -> Self { |
| 1590 | let task_cancel = cancel.clone(); |
| 1591 | // The watch stays per-connection by design (#6211 R7a): it is born |
| 1592 | // with the connect attempt (covering the pre-insertion window) and |
| 1593 | // dies with the connection, so a watched server can neither be |
| 1594 | // missed nor leak. A pool-level task would need the pool lock — |
| 1595 | // held across in-flight calls — and regress the mid-call trip this |
| 1596 | // exists for. What moves is the check itself: synchronous |
| 1597 | // state fs has no place on the executor at 20Hz, so it runs on the |
| 1598 | // blocking pool while the 50ms revocation cadence is unchanged. |
| 1599 | let source = Arc::new(source); |
| 1600 | let handle = tokio::spawn(async move { |
| 1601 | loop { |
| 1602 | let source = Arc::clone(&source); |
| 1603 | let check = tokio::task::spawn_blocking(move || { |
| 1604 | crate::plugins::registry::verify_plugin_state_authority(&source.authority) |
| 1605 | }) |
| 1606 | .await; |
| 1607 | let reason = match check { |
| 1608 | Ok(Err(reason)) => Some(reason), |
| 1609 | Ok(Ok(())) => None, |
| 1610 | Err(_) => { |
| 1611 | Some("plugin authority check failed to run; failing closed".to_string()) |
| 1612 | } |
| 1613 | }; |
| 1614 | if let Some(reason) = reason { |
| 1615 | if let Ok(mut slot) = reason_slot.lock() { |
| 1616 | *slot = Some(reason); |
| 1617 | } |
| 1618 | task_cancel.cancel(); |
| 1619 | break; |
| 1620 | } |
| 1621 | tokio::select! { |
| 1622 | _ = task_cancel.cancelled() => break, |
| 1623 | _ = tokio::time::sleep(Duration::from_millis(50)) => {} |
| 1624 | } |
| 1625 | } |
| 1626 | }); |
| 1627 | Self { |
| 1628 | handle: Some(handle), |
| 1629 | cancel, |
| 1630 | armed: true, |
| 1631 | } |
| 1632 | } |
| 1633 | |
| 1634 | fn disarm(mut self) -> tokio::task::JoinHandle<()> { |
| 1635 | self.armed = false; |
| 1636 | self.handle.take().expect("authority watch must exist") |
| 1637 | } |
| 1638 | } |
| 1639 | |
| 1640 | impl Drop for PendingAuthorityWatch { |
| 1641 | fn drop(&mut self) { |
| 1642 | if self.armed { |
| 1643 | self.cancel.cancel(); |
| 1644 | if let Some(handle) = self.handle.take() { |
| 1645 | handle.abort(); |
| 1646 | } |
| 1647 | } |
| 1648 | } |
| 1649 | } |
| 1650 | |
| 1651 | impl McpConnection { |
| 1652 | /// Connect to an MCP server and initialize it. |
| 1653 | /// |
| 1654 | /// `network_policy` (added in v0.7.0 for #135) is consulted for HTTP/SSE |
| 1655 | /// transports only — STDIO transports are unaffected. Pass `None` to |
| 1656 | /// match pre-v0.7.0 permissive behavior. |
| 1657 | pub async fn connect_with_policy( |
| 1658 | name: String, |
| 1659 | config: McpServerConfig, |
| 1660 | global_timeouts: &McpTimeouts, |
| 1661 | network_policy: Option<&NetworkPolicyDecider>, |
| 1662 | ) -> Result<Self> { |
| 1663 | let connect_timeout_secs = config.effective_connect_timeout(global_timeouts); |
| 1664 | let read_timeout_secs = config.effective_read_timeout(global_timeouts); |
| 1665 | let cancel_token = tokio_util::sync::CancellationToken::new(); |
| 1666 | let authority_revocation_reason = Arc::new(std::sync::Mutex::new(None)); |
| 1667 | if let Some(source) = config.reviewed_plugin.as_ref() { |
| 1668 | source.validate_before_use(&name, "connect")?; |
| 1669 | if let Some(url) = config.url.as_deref() { |
| 1670 | source.validate_remote_endpoint(&name, url)?; |
| 1671 | } |
| 1672 | } |
| 1673 | // Start the cross-process generation watch before any network request |
| 1674 | // or child spawn. The guard cancels and aborts itself on every early |
| 1675 | // return; a successful connection transfers the task into `Self`. |
| 1676 | let authority_watch = config.reviewed_plugin.clone().map(|source| { |
| 1677 | PendingAuthorityWatch::start( |
| 1678 | source, |
| 1679 | cancel_token.clone(), |
| 1680 | Arc::clone(&authority_revocation_reason), |
| 1681 | ) |
| 1682 | }); |
| 1683 | let transport: Box<dyn McpTransport> = if let Some(url) = &config.url { |
| 1684 | // Per-domain network policy gate (#135). Only the HTTP/SSE transport |
| 1685 | // is gated; STDIO MCP servers run as local subprocesses and never |
| 1686 | // touch the network from this code path. |
| 1687 | if let Some(decider) = network_policy |
| 1688 | && let Some(host) = host_from_url(url) |
| 1689 | { |
| 1690 | match decider.evaluate(&host, "mcp") { |
| 1691 | Decision::Allow => {} |
| 1692 | Decision::Deny => { |
| 1693 | anyhow::bail!( |
| 1694 | "MCP server '{name}' connection to '{host}' blocked by network policy" |
| 1695 | ); |
| 1696 | } |
| 1697 | Decision::Prompt => { |
| 1698 | anyhow::bail!( |
| 1699 | "MCP server '{name}' connection to '{host}' requires approval; \ |
| 1700 | re-run after `/network allow {host}` or set network.default = \"allow\" in config" |
| 1701 | ); |
| 1702 | } |
| 1703 | } |
| 1704 | } |
| 1705 | let client = http_client::McpHttpClient::new( |
| 1706 | url, |
| 1707 | config.runtime_added, |
| 1708 | config.reviewed_plugin.is_some(), |
| 1709 | config.allow_private_network, |
| 1710 | network_policy, |
| 1711 | Duration::from_secs(connect_timeout_secs), |
| 1712 | Duration::from_secs(read_timeout_secs), |
| 1713 | )?; |
| 1714 | let oauth_runtime = if config.reviewed_plugin.is_some() { |
| 1715 | None |
| 1716 | } else { |
| 1717 | match oauth::build_default_headers(&config.headers, &config.env_headers) { |
| 1718 | Ok(default_headers) => { |
| 1719 | let prepared = tokio::select! { |
| 1720 | biased; |
| 1721 | _ = cancel_token.cancelled() => { |
| 1722 | anyhow::bail!( |
| 1723 | "MCP OAuth setup cancelled after plugin authority changed" |
| 1724 | ) |
| 1725 | } |
| 1726 | prepared = oauth::McpOAuthRuntime::from_server_config_with_client( |
| 1727 | &name, |
| 1728 | &config, |
| 1729 | default_headers, |
| 1730 | client.clone(), |
| 1731 | ) => prepared, |
| 1732 | }; |
| 1733 | match prepared { |
| 1734 | Ok(runtime) => runtime, |
| 1735 | Err(err) => { |
| 1736 | if config.reviewed_plugin.is_some() { |
| 1737 | tracing::warn!( |
| 1738 | target: "mcp", |
| 1739 | server = %name, |
| 1740 | "failed to prepare reviewed plugin MCP OAuth runtime; provider details suppressed; continuing without stored OAuth token" |
| 1741 | ); |
| 1742 | } else { |
| 1743 | tracing::warn!( |
| 1744 | target: "mcp", |
| 1745 | server = %name, |
| 1746 | error = %err, |
| 1747 | "failed to prepare MCP OAuth runtime; continuing without stored OAuth token" |
| 1748 | ); |
| 1749 | } |
| 1750 | None |
| 1751 | } |
| 1752 | } |
| 1753 | } |
| 1754 | Err(err) => { |
| 1755 | if config.reviewed_plugin.is_some() { |
| 1756 | tracing::warn!( |
| 1757 | target: "mcp", |
| 1758 | server = %name, |
| 1759 | "failed to prepare reviewed plugin MCP OAuth headers; details suppressed; continuing without stored OAuth token" |
| 1760 | ); |
| 1761 | } else { |
| 1762 | tracing::warn!( |
| 1763 | target: "mcp", |
| 1764 | server = %name, |
| 1765 | error = %err, |
| 1766 | "failed to prepare MCP OAuth default headers; continuing without stored OAuth token" |
| 1767 | ); |
| 1768 | } |
| 1769 | None |
| 1770 | } |
| 1771 | } |
| 1772 | }; |
| 1773 | let http_auth = McpHttpAuth::from_config(&name, &config, oauth_runtime); |
| 1774 | if is_legacy_sse_transport(&config) { |
| 1775 | Box::new( |
| 1776 | SseTransport::connect( |
| 1777 | client, |
| 1778 | url.clone(), |
| 1779 | http_auth, |
| 1780 | cancel_token.clone(), |
| 1781 | Duration::from_secs(connect_timeout_secs), |
| 1782 | ) |
| 1783 | .await?, |
| 1784 | ) |
| 1785 | } else { |
| 1786 | let mut http = HttpTransport::new( |
| 1787 | client, |
| 1788 | url.clone(), |
| 1789 | http_auth, |
| 1790 | cancel_token.clone(), |
| 1791 | Duration::from_secs(connect_timeout_secs), |
| 1792 | ); |
| 1793 | // Best-effort session preflight for servers that require |
| 1794 | // a session ID on every POST including `initialize` |
| 1795 | // (e.g. Hindsight, #1629). Failures are non-fatal — the |
| 1796 | // `initialize` POST will proceed and may capture a session |
| 1797 | // ID from the response instead. |
| 1798 | if let Err(e) = http.try_establish_session().await { |
| 1799 | tracing::debug!( |
| 1800 | target: "mcp", |
| 1801 | server = %name, |
| 1802 | error = %e, |
| 1803 | "session-establishment GET skipped; proceeding with POST initialize" |
| 1804 | ); |
| 1805 | } |
| 1806 | Box::new(http) |
| 1807 | } |
| 1808 | } else if let Some(command) = &config.command { |
| 1809 | Box::new(StdioTransport::spawn( |
| 1810 | &name, |
| 1811 | command, |
| 1812 | &config, |
| 1813 | cancel_token.clone(), |
| 1814 | )?) |
| 1815 | } else { |
| 1816 | anyhow::bail!("MCP server '{name}' config must have either 'command' or 'url'"); |
| 1817 | }; |
| 1818 | // Revalidate after transport construction as well: remote setup may |
| 1819 | // await DNS/TLS/SSE preflight, and a concurrent process can revoke the |
| 1820 | // receipt during that interval. Initialization and catalog discovery |
| 1821 | // never start under a stale generation. |
| 1822 | if let Some(source) = config.reviewed_plugin.as_ref() { |
| 1823 | source.validate_before_use(&name, "initialize")?; |
| 1824 | } |
| 1825 | let authority_watch = authority_watch.map(PendingAuthorityWatch::disarm); |
| 1826 | |
| 1827 | let mut conn = Self { |
| 1828 | name: name.clone(), |
| 1829 | transport, |
| 1830 | tools: Vec::new(), |
| 1831 | resources: Vec::new(), |
| 1832 | resource_templates: Vec::new(), |
| 1833 | prompts: Vec::new(), |
| 1834 | request_id: AtomicU64::new(1), |
| 1835 | state: ConnectionState::Connecting, |
| 1836 | config, |
| 1837 | server_capabilities: None, |
| 1838 | discovery_timeout: Duration::from_secs(connect_timeout_secs), |
| 1839 | read_timeout_secs, |
| 1840 | cancel_token, |
| 1841 | authority_revocation_reason, |
| 1842 | authority_watch, |
| 1843 | catalog_generation: 0, |
| 1844 | }; |
| 1845 | |
| 1846 | // Initialize with timeout |
| 1847 | tokio::time::timeout(Duration::from_secs(connect_timeout_secs), conn.initialize()) |
| 1848 | .await |
| 1849 | .with_context(|| format!("MCP server '{name}' initialization timed out"))??; |
| 1850 | |
| 1851 | conn.discover_all() |
| 1852 | .await |
| 1853 | .with_context(|| format!("MCP server '{name}' discovery failed"))?; |
| 1854 | |
| 1855 | conn.state = ConnectionState::Ready; |
| 1856 | Ok(conn) |
| 1857 | } |
| 1858 | |
| 1859 | /// Send initialize request and wait for response |
| 1860 | async fn initialize(&mut self) -> Result<()> { |
| 1861 | let init_id = self.next_id(); |
| 1862 | self.send(serde_json::json!({ |
| 1863 | "jsonrpc": "2.0", |
| 1864 | "id": &init_id, |
| 1865 | "method": "initialize", |
| 1866 | "params": { |
| 1867 | "protocolVersion": MCP_PROTOCOL_VERSION, |
| 1868 | "clientInfo": { |
| 1869 | "name": "codewhale-tui", |
| 1870 | "version": env!("CARGO_PKG_VERSION") |
| 1871 | }, |
| 1872 | "capabilities": { |
| 1873 | "tools": {}, |
| 1874 | "resources": {}, |
| 1875 | "prompts": {} |
| 1876 | } |
| 1877 | } |
| 1878 | })) |
| 1879 | .await?; |
| 1880 | |
| 1881 | let response = self.recv(init_id).await?; |
| 1882 | let result = response_result( |
| 1883 | &response, |
| 1884 | "initialize", |
| 1885 | self.config.reviewed_plugin.is_some(), |
| 1886 | )?; |
| 1887 | // Per spec, a server that cannot speak the advertised revision answers |
| 1888 | // with one it does support. Accept any dated revision we still |
| 1889 | // implement; anything else ends the handshake. |
| 1890 | let negotiated = result |
| 1891 | .and_then(|result| result.get("protocolVersion")) |
| 1892 | .and_then(|version| version.as_str()) |
| 1893 | .ok_or_else(|| { |
| 1894 | anyhow::anyhow!( |
| 1895 | "MCP server '{}' initialize result omitted protocolVersion", |
| 1896 | self.name |
| 1897 | ) |
| 1898 | })?; |
| 1899 | anyhow::ensure!( |
| 1900 | MCP_SUPPORTED_PROTOCOL_VERSIONS.contains(&negotiated), |
| 1901 | "MCP server '{}' negotiated unsupported protocol version '{negotiated}' (supported: {})", |
| 1902 | self.name, |
| 1903 | MCP_SUPPORTED_PROTOCOL_VERSIONS.join(", ") |
| 1904 | ); |
| 1905 | self.transport.set_protocol_version(negotiated); |
| 1906 | self.server_capabilities = McpServerCapabilities::from_initialize_response(&response); |
| 1907 | |
| 1908 | // Send initialized notification (no id, no response expected) |
| 1909 | self.send(serde_json::json!({ |
| 1910 | "jsonrpc": "2.0", |
| 1911 | "method": "notifications/initialized" |
| 1912 | })) |
| 1913 | .await?; |
| 1914 | |
| 1915 | Ok(()) |
| 1916 | } |
| 1917 | |
| 1918 | /// Discover tools, resources, and prompts |
| 1919 | async fn discover_all(&mut self) -> Result<()> { |
| 1920 | let capabilities = self.server_capabilities; |
| 1921 | let server = self.name.clone(); |
| 1922 | let discovery_timeout = self.discovery_timeout; |
| 1923 | |
| 1924 | // Missing initialize metadata is treated as a legacy/unknown server: |
| 1925 | // retain tool discovery and bounded best-effort probes for compatibility. |
| 1926 | // When capabilities are advertised, do not call methods the server says |
| 1927 | // it does not implement (notably JetBrains tools-only MCP servers). |
| 1928 | if capabilities.is_none_or(|capabilities| capabilities.tools) { |
| 1929 | tokio::time::timeout(discovery_timeout, self.discover_tools()) |
| 1930 | .await |
| 1931 | .with_context(|| { |
| 1932 | format!( |
| 1933 | "MCP server '{}' tool discovery timed out after {:?}", |
| 1934 | server, discovery_timeout |
| 1935 | ) |
| 1936 | })??; |
| 1937 | } |
| 1938 | |
| 1939 | // Keep all three optional calls within one discovery-timeout budget in |
| 1940 | // the worst case while also respecting a tighter transport read timeout. |
| 1941 | let optional_timeout = |
| 1942 | (discovery_timeout / 3).min(Duration::from_secs(self.read_timeout_secs)); |
| 1943 | if capabilities.is_none_or(|capabilities| capabilities.resources) { |
| 1944 | run_optional_discovery( |
| 1945 | &server, |
| 1946 | "resources/list", |
| 1947 | optional_timeout, |
| 1948 | self.discover_resources(), |
| 1949 | ) |
| 1950 | .await; |
| 1951 | run_optional_discovery( |
| 1952 | &server, |
| 1953 | "resources/templates/list", |
| 1954 | optional_timeout, |
| 1955 | self.discover_resource_templates(), |
| 1956 | ) |
| 1957 | .await; |
| 1958 | } |
| 1959 | if capabilities.is_none_or(|capabilities| capabilities.prompts) { |
| 1960 | run_optional_discovery( |
| 1961 | &server, |
| 1962 | "prompts/list", |
| 1963 | optional_timeout, |
| 1964 | self.discover_prompts(), |
| 1965 | ) |
| 1966 | .await; |
| 1967 | } |
| 1968 | Ok(()) |
| 1969 | } |
| 1970 | |
| 1971 | /// Discover available tools from the MCP server |
| 1972 | async fn discover_tools(&mut self) -> Result<()> { |
| 1973 | let mut cursor: Option<String> = None; |
| 1974 | let mut budget = McpCatalogBudget::new("tools/list"); |
| 1975 | let mut discovered = Vec::new(); |
| 1976 | loop { |
| 1977 | let list_id = self.next_id(); |
| 1978 | let params = match &cursor { |
| 1979 | Some(c) => serde_json::json!({ "cursor": c }), |
| 1980 | None => serde_json::json!({}), |
| 1981 | }; |
| 1982 | self.send(serde_json::json!({ |
| 1983 | "jsonrpc": "2.0", |
| 1984 | "id": &list_id, |
| 1985 | "method": "tools/list", |
| 1986 | "params": params |
| 1987 | })) |
| 1988 | .await?; |
| 1989 | |
| 1990 | let response = self.recv(list_id).await?; |
| 1991 | let Some(result) = response_result( |
| 1992 | &response, |
| 1993 | "tools/list", |
| 1994 | self.config.reviewed_plugin.is_some(), |
| 1995 | )? |
| 1996 | else { |
| 1997 | break; |
| 1998 | }; |
| 1999 | |
| 2000 | let items = result |
| 2001 | .get("tools") |
| 2002 | .and_then(|tools| tools.as_array()) |
| 2003 | .map_or(0, Vec::len); |
| 2004 | if let Some(arr) = result.get("tools").and_then(|t| t.as_array()) { |
| 2005 | for item in arr { |
| 2006 | match serde_json::from_value::<McpTool>(item.clone()) { |
| 2007 | Ok(tool) => discovered.push(tool), |
| 2008 | Err(err) => { |
| 2009 | // Skip individual malformed entries instead of |
| 2010 | // dropping the whole page (#1410). The old |
| 2011 | // `unwrap_or_default()` would silently throw |
| 2012 | // away every tool when one was misshapen. |
| 2013 | tracing::debug!(target: "mcp", ?err, "skipping malformed tool item"); |
| 2014 | } |
| 2015 | } |
| 2016 | } |
| 2017 | } |
| 2018 | |
| 2019 | cursor = budget.observe_page(result, items)?; |
| 2020 | if cursor.is_none() { |
| 2021 | break; |
| 2022 | } |
| 2023 | } |
| 2024 | // Sort by tool name so the order the model sees doesn't depend on |
| 2025 | // server-side pagination ordering — keeps the prompt prefix stable |
| 2026 | // for cache-hit purposes (#1319). |
| 2027 | discovered.sort_by(|a, b| a.name.cmp(&b.name)); |
| 2028 | self.tools = discovered; |
| 2029 | Ok(()) |
| 2030 | } |
| 2031 | |
| 2032 | /// Discover available resources from the MCP server |
| 2033 | async fn discover_resources(&mut self) -> Result<()> { |
| 2034 | let mut cursor: Option<String> = None; |
| 2035 | let mut budget = McpCatalogBudget::new("resources/list"); |
| 2036 | let mut discovered = Vec::new(); |
| 2037 | loop { |
| 2038 | let list_id = self.next_id(); |
| 2039 | let params = match &cursor { |
| 2040 | Some(c) => serde_json::json!({ "cursor": c }), |
| 2041 | None => serde_json::json!({}), |
| 2042 | }; |
| 2043 | self.send(serde_json::json!({ |
| 2044 | "jsonrpc": "2.0", |
| 2045 | "id": &list_id, |
| 2046 | "method": "resources/list", |
| 2047 | "params": params |
| 2048 | })) |
| 2049 | .await?; |
| 2050 | |
| 2051 | let response = self.recv(list_id).await?; |
| 2052 | let Some(result) = response_result( |
| 2053 | &response, |
| 2054 | "resources/list", |
| 2055 | self.config.reviewed_plugin.is_some(), |
| 2056 | )? |
| 2057 | else { |
| 2058 | break; |
| 2059 | }; |
| 2060 | |
| 2061 | let items = result |
| 2062 | .get("resources") |
| 2063 | .and_then(|resources| resources.as_array()) |
| 2064 | .map_or(0, Vec::len); |
| 2065 | if let Some(arr) = result.get("resources").and_then(|r| r.as_array()) { |
| 2066 | for item in arr { |
| 2067 | match serde_json::from_value::<McpResource>(item.clone()) { |
| 2068 | Ok(resource) => discovered.push(resource), |
| 2069 | Err(err) => { |
| 2070 | tracing::debug!(target: "mcp", ?err, "skipping malformed resource item"); |
| 2071 | } |
| 2072 | } |
| 2073 | } |
| 2074 | } |
| 2075 | |
| 2076 | cursor = budget.observe_page(result, items)?; |
| 2077 | if cursor.is_none() { |
| 2078 | break; |
| 2079 | } |
| 2080 | } |
| 2081 | self.resources = discovered; |
| 2082 | Ok(()) |
| 2083 | } |
| 2084 | |
| 2085 | /// Discover available resource templates from the MCP server |
| 2086 | async fn discover_resource_templates(&mut self) -> Result<()> { |
| 2087 | let mut cursor: Option<String> = None; |
| 2088 | let mut budget = McpCatalogBudget::new("resources/templates/list"); |
| 2089 | let mut discovered = Vec::new(); |
| 2090 | loop { |
| 2091 | let list_id = self.next_id(); |
| 2092 | let params = match &cursor { |
| 2093 | Some(c) => serde_json::json!({ "cursor": c }), |
| 2094 | None => serde_json::json!({}), |
| 2095 | }; |
| 2096 | self.send(serde_json::json!({ |
| 2097 | "jsonrpc": "2.0", |
| 2098 | "id": &list_id, |
| 2099 | "method": "resources/templates/list", |
| 2100 | "params": params |
| 2101 | })) |
| 2102 | .await?; |
| 2103 | |
| 2104 | let response = self.recv(list_id).await?; |
| 2105 | let Some(result) = response_result( |
| 2106 | &response, |
| 2107 | "resources/templates/list", |
| 2108 | self.config.reviewed_plugin.is_some(), |
| 2109 | )? |
| 2110 | else { |
| 2111 | break; |
| 2112 | }; |
| 2113 | |
| 2114 | let templates = result |
| 2115 | .get("resourceTemplates") |
| 2116 | .or_else(|| result.get("templates")) |
| 2117 | .or_else(|| result.get("resource_templates")); |
| 2118 | let items = templates |
| 2119 | .and_then(|templates| templates.as_array()) |
| 2120 | .map_or(0, Vec::len); |
| 2121 | if let Some(arr) = templates.and_then(|t| t.as_array()) { |
| 2122 | for item in arr { |
| 2123 | match serde_json::from_value::<McpResourceTemplate>(item.clone()) { |
| 2124 | Ok(tmpl) => discovered.push(tmpl), |
| 2125 | Err(err) => { |
| 2126 | tracing::debug!(target: "mcp", ?err, "skipping malformed resource_template item"); |
| 2127 | } |
| 2128 | } |
| 2129 | } |
| 2130 | } |
| 2131 | |
| 2132 | cursor = budget.observe_page(result, items)?; |
| 2133 | if cursor.is_none() { |
| 2134 | break; |
| 2135 | } |
| 2136 | } |
| 2137 | self.resource_templates = discovered; |
| 2138 | Ok(()) |
| 2139 | } |
| 2140 | |
| 2141 | /// Discover available prompts from the MCP server |
| 2142 | async fn discover_prompts(&mut self) -> Result<()> { |
| 2143 | let mut cursor: Option<String> = None; |
| 2144 | let mut budget = McpCatalogBudget::new("prompts/list"); |
| 2145 | let mut discovered = Vec::new(); |
| 2146 | loop { |
| 2147 | let list_id = self.next_id(); |
| 2148 | let params = match &cursor { |
| 2149 | Some(c) => serde_json::json!({ "cursor": c }), |
| 2150 | None => serde_json::json!({}), |
| 2151 | }; |
| 2152 | self.send(serde_json::json!({ |
| 2153 | "jsonrpc": "2.0", |
| 2154 | "id": &list_id, |
| 2155 | "method": "prompts/list", |
| 2156 | "params": params |
| 2157 | })) |
| 2158 | .await?; |
| 2159 | |
| 2160 | let response = self.recv(list_id).await?; |
| 2161 | let Some(result) = response_result( |
| 2162 | &response, |
| 2163 | "prompts/list", |
| 2164 | self.config.reviewed_plugin.is_some(), |
| 2165 | )? |
| 2166 | else { |
| 2167 | break; |
| 2168 | }; |
| 2169 | |
| 2170 | let items = result |
| 2171 | .get("prompts") |
| 2172 | .and_then(|prompts| prompts.as_array()) |
| 2173 | .map_or(0, Vec::len); |
| 2174 | if let Some(arr) = result.get("prompts").and_then(|p| p.as_array()) { |
| 2175 | for item in arr { |
| 2176 | match serde_json::from_value::<McpPrompt>(item.clone()) { |
| 2177 | Ok(prompt) => discovered.push(prompt), |
| 2178 | Err(err) => { |
| 2179 | tracing::debug!(target: "mcp", ?err, "skipping malformed prompt item"); |
| 2180 | } |
| 2181 | } |
| 2182 | } |
| 2183 | } |
| 2184 | |
| 2185 | cursor = budget.observe_page(result, items)?; |
| 2186 | if cursor.is_none() { |
| 2187 | break; |
| 2188 | } |
| 2189 | } |
| 2190 | self.prompts = discovered; |
| 2191 | Ok(()) |
| 2192 | } |
| 2193 | |
| 2194 | /// Call a tool on this MCP server |
| 2195 | pub async fn call_tool( |
| 2196 | &mut self, |
| 2197 | tool_name: &str, |
| 2198 | arguments: serde_json::Value, |
| 2199 | timeout_secs: u64, |
| 2200 | ) -> Result<serde_json::Value> { |
| 2201 | self.call_method( |
| 2202 | "tools/call", |
| 2203 | serde_json::json!({ |
| 2204 | "name": tool_name, |
| 2205 | "arguments": arguments |
| 2206 | }), |
| 2207 | timeout_secs, |
| 2208 | ) |
| 2209 | .await |
| 2210 | } |
| 2211 | |
| 2212 | /// Read a resource from this MCP server |
| 2213 | pub async fn read_resource( |
| 2214 | &mut self, |
| 2215 | uri: &str, |
| 2216 | timeout_secs: u64, |
| 2217 | ) -> Result<serde_json::Value> { |
| 2218 | self.call_method( |
| 2219 | "resources/read", |
| 2220 | serde_json::json!({ |
| 2221 | "uri": uri |
| 2222 | }), |
| 2223 | timeout_secs, |
| 2224 | ) |
| 2225 | .await |
| 2226 | } |
| 2227 | |
| 2228 | /// Get a prompt from this MCP server |
| 2229 | pub async fn get_prompt( |
| 2230 | &mut self, |
| 2231 | prompt_name: &str, |
| 2232 | arguments: serde_json::Value, |
| 2233 | timeout_secs: u64, |
| 2234 | ) -> Result<serde_json::Value> { |
| 2235 | self.call_method( |
| 2236 | "prompts/get", |
| 2237 | serde_json::json!({ |
| 2238 | "name": prompt_name, |
| 2239 | "arguments": arguments |
| 2240 | }), |
| 2241 | timeout_secs, |
| 2242 | ) |
| 2243 | .await |
| 2244 | } |
| 2245 | |
| 2246 | /// Generic method to call an MCP method |
| 2247 | async fn call_method( |
| 2248 | &mut self, |
| 2249 | method: &str, |
| 2250 | params: serde_json::Value, |
| 2251 | timeout_secs: u64, |
| 2252 | ) -> Result<serde_json::Value> { |
| 2253 | if self.state != ConnectionState::Ready { |
| 2254 | anyhow::bail!( |
| 2255 | "Failed to call MCP method '{}': connection '{}' is not ready", |
| 2256 | method, |
| 2257 | self.name |
| 2258 | ); |
| 2259 | } |
| 2260 | if let Some(source) = self.config.reviewed_plugin.as_ref() { |
| 2261 | source.validate_before_use(&self.name, method)?; |
| 2262 | } |
| 2263 | |
| 2264 | let call_id = self.next_id(); |
| 2265 | if let Err(error) = self |
| 2266 | .send(serde_json::json!({ |
| 2267 | "jsonrpc": "2.0", |
| 2268 | "id": &call_id, |
| 2269 | "method": method, |
| 2270 | "params": params |
| 2271 | })) |
| 2272 | .await |
| 2273 | { |
| 2274 | return self.finish_guarded_error(error).await; |
| 2275 | } |
| 2276 | |
| 2277 | let response = |
| 2278 | match tokio::time::timeout(Duration::from_secs(timeout_secs), self.recv(call_id)) |
| 2279 | .await |
| 2280 | .with_context(|| { |
| 2281 | format!( |
| 2282 | "MCP method '{}' on server '{}' timed out after {}s", |
| 2283 | method, self.name, timeout_secs |
| 2284 | ) |
| 2285 | }) { |
| 2286 | Ok(Ok(response)) => response, |
| 2287 | Ok(Err(error)) => return self.finish_guarded_error(error).await, |
| 2288 | Err(error) => return self.finish_guarded_error(error).await, |
| 2289 | }; |
| 2290 | |
| 2291 | if let Some(error) = response.get("error") { |
| 2292 | if self.config.reviewed_plugin.is_some() { |
| 2293 | anyhow::bail!( |
| 2294 | "Reviewed plugin MCP server returned an error in '{method}' (server details suppressed to protect environment-backed credentials)" |
| 2295 | ); |
| 2296 | } |
| 2297 | return Err(anyhow::anyhow!( |
| 2298 | "MCP error in '{}': {}", |
| 2299 | method, |
| 2300 | serde_json::to_string_pretty(error)? |
| 2301 | )); |
| 2302 | } |
| 2303 | |
| 2304 | // JSON-RPC requires exactly one of `result` / `error`. Treating a |
| 2305 | // response carrying neither as an empty success handed the model a |
| 2306 | // `null` tool result that is indistinguishable from a tool that |
| 2307 | // genuinely returned nothing. An explicit `"result": null` is still a |
| 2308 | // valid empty success and passes through unchanged. |
| 2309 | response.get("result").cloned().with_context(|| { |
| 2310 | format!( |
| 2311 | "MCP response from server '{}' for '{method}' contained neither a result nor an error", |
| 2312 | self.name |
| 2313 | ) |
| 2314 | }) |
| 2315 | } |
| 2316 | |
| 2317 | /// Get discovered tools |
| 2318 | pub fn tools(&self) -> &[McpTool] { |
| 2319 | &self.tools |
| 2320 | } |
| 2321 | |
| 2322 | /// Get discovered resources |
| 2323 | pub fn resources(&self) -> &[McpResource] { |
| 2324 | &self.resources |
| 2325 | } |
| 2326 | |
| 2327 | /// Get discovered resource templates |
| 2328 | pub fn resource_templates(&self) -> &[McpResourceTemplate] { |
| 2329 | &self.resource_templates |
| 2330 | } |
| 2331 | |
| 2332 | /// Get discovered prompts |
| 2333 | pub fn prompts(&self) -> &[McpPrompt] { |
| 2334 | &self.prompts |
| 2335 | } |
| 2336 | |
| 2337 | /// Get server name |
| 2338 | #[allow(dead_code)] // Public API for MCP consumers |
| 2339 | pub fn name(&self) -> &str { |
| 2340 | &self.name |
| 2341 | } |
| 2342 | |
| 2343 | /// Ready to dispatch: the transport is live **and** the plugin bundle |
| 2344 | /// backing it still carries the authority it was reviewed with. |
| 2345 | pub fn is_ready(&self) -> bool { |
| 2346 | self.is_transport_ready() && self.catalog_authorized() |
| 2347 | } |
| 2348 | |
| 2349 | /// Liveness only — no authority check. |
| 2350 | /// |
| 2351 | /// The Ready flag alone can't see a stdio child that exited between |
| 2352 | /// calls; the probe closes that gap so the pool rebuilds the connection |
| 2353 | /// instead of handing a dead transport back (#6187). |
| 2354 | /// |
| 2355 | /// Only for callers that have just run `validate_before_use` on this same |
| 2356 | /// source, where `is_ready`'s authority half would re-walk and re-hash the |
| 2357 | /// plugin bundle it already verified one statement earlier (#6209). Every |
| 2358 | /// other caller must use `is_ready`: dropping the authority half without |
| 2359 | /// a preceding check silently dispatches to a revoked or altered bundle. |
| 2360 | pub(crate) fn is_transport_ready(&self) -> bool { |
| 2361 | self.state == ConnectionState::Ready && !self.transport.probe_dead() |
| 2362 | } |
| 2363 | |
| 2364 | /// Get server config |
| 2365 | pub fn config(&self) -> &McpServerConfig { |
| 2366 | &self.config |
| 2367 | } |
| 2368 | |
| 2369 | /// Get connection state |
| 2370 | #[allow(dead_code)] // Public API for MCP consumers |
| 2371 | pub fn state(&self) -> ConnectionState { |
| 2372 | self.state |
| 2373 | } |
| 2374 | |
| 2375 | fn next_id(&self) -> String { |
| 2376 | self.request_id.fetch_add(1, Ordering::SeqCst).to_string() |
| 2377 | } |
| 2378 | |
| 2379 | async fn send(&mut self, msg: serde_json::Value) -> Result<()> { |
| 2380 | let bytes = serde_json::to_vec(&msg).context("Failed to serialize MCP JSON-RPC message")?; |
| 2381 | let cancel_token = self.cancel_token.clone(); |
| 2382 | let name = self.name.clone(); |
| 2383 | let result = tokio::select! { |
| 2384 | biased; |
| 2385 | _ = cancel_token.cancelled() => { |
| 2386 | Err(anyhow::anyhow!("MCP connection '{name}' was cancelled")) |
| 2387 | } |
| 2388 | result = self.transport.send(bytes) => result, |
| 2389 | }; |
| 2390 | if result.is_err() { |
| 2391 | // A dead write side is as fatal as a dead read side: the pool |
| 2392 | // reuses any connection whose `is_ready()` is true, so leaving |
| 2393 | // this one in `Ready` would hand the same broken transport back |
| 2394 | // on every later call instead of rebuilding it. |
| 2395 | self.state = ConnectionState::Disconnected; |
| 2396 | } |
| 2397 | result |
| 2398 | } |
| 2399 | |
| 2400 | async fn recv(&mut self, expected_id: String) -> Result<serde_json::Value> { |
| 2401 | loop { |
| 2402 | let bytes = match tokio::time::timeout( |
| 2403 | Duration::from_secs(self.read_timeout_secs), |
| 2404 | async { |
| 2405 | tokio::select! { |
| 2406 | biased; |
| 2407 | _ = self.cancel_token.cancelled() => { |
| 2408 | anyhow::bail!("MCP connection '{}' was cancelled", self.name) |
| 2409 | } |
| 2410 | result = self.transport.recv() => result, |
| 2411 | } |
| 2412 | }, |
| 2413 | ) |
| 2414 | .await |
| 2415 | { |
| 2416 | Ok(result) => result.inspect_err(|_e| { |
| 2417 | self.state = ConnectionState::Disconnected; |
| 2418 | })?, |
| 2419 | Err(_) => { |
| 2420 | self.state = ConnectionState::Disconnected; |
| 2421 | anyhow::bail!( |
| 2422 | "Timed out waiting for MCP JSON-RPC response from server '{}' after {}s", |
| 2423 | self.name, |
| 2424 | self.read_timeout_secs |
| 2425 | ); |
| 2426 | } |
| 2427 | }; |
| 2428 | let value: serde_json::Value = match serde_json::from_slice(&bytes) { |
| 2429 | Ok(value) => value, |
| 2430 | Err(err) => { |
| 2431 | self.state = ConnectionState::Disconnected; |
| 2432 | let preview = if self.config.reviewed_plugin.is_some() { |
| 2433 | "<server details suppressed for reviewed plugin>".to_string() |
| 2434 | } else { |
| 2435 | invalid_json_preview(&bytes) |
| 2436 | }; |
| 2437 | return Err(err).with_context(|| { |
| 2438 | format!( |
| 2439 | "Invalid MCP JSON-RPC message from server '{}': {}", |
| 2440 | self.name, preview |
| 2441 | ) |
| 2442 | }); |
| 2443 | } |
| 2444 | }; |
| 2445 | |
| 2446 | // Check if this is a response with the expected id. We emit |
| 2447 | // string IDs because some MCP gateways reject numeric JSON-RPC |
| 2448 | // IDs, but accept numeric echoes for compatibility with older |
| 2449 | // servers and tests. |
| 2450 | if response_id_matches(value.get("id"), &expected_id) { |
| 2451 | if let Some(error) = value.get("error") |
| 2452 | && is_mcp_stale_session_body(&error.to_string()) |
| 2453 | { |
| 2454 | anyhow::bail!("MCP session expired: {error}"); |
| 2455 | } |
| 2456 | return Ok(value); |
| 2457 | } |
| 2458 | // Skip notifications (no id) and responses with different ids |
| 2459 | } |
| 2460 | } |
| 2461 | |
| 2462 | /// Gracefully close the connection |
| 2463 | #[allow(dead_code)] // Public API for MCP consumers |
| 2464 | pub fn close(&mut self) { |
| 2465 | self.cancel_token.cancel(); |
| 2466 | self.state = ConnectionState::Disconnected; |
| 2467 | } |
| 2468 | |
| 2469 | fn catalog_authorized(&self) -> bool { |
| 2470 | self.config |
| 2471 | .reviewed_plugin |
| 2472 | .as_ref() |
| 2473 | .is_none_or(ReviewedPluginMcpSource::catalog_is_current) |
| 2474 | } |
| 2475 | |
| 2476 | async fn finish_guarded_error<T>(&mut self, error: anyhow::Error) -> Result<T> { |
| 2477 | let reason = self |
| 2478 | .authority_revocation_reason |
| 2479 | .lock() |
| 2480 | .ok() |
| 2481 | .and_then(|reason| reason.clone()); |
| 2482 | if let Some(reason) = reason { |
| 2483 | self.transport.shutdown().await; |
| 2484 | self.state = ConnectionState::Disconnected; |
| 2485 | anyhow::bail!( |
| 2486 | "MCP operation on plugin server '{}' was cancelled after authority changed: {reason}", |
| 2487 | self.name |
| 2488 | ); |
| 2489 | } |
| 2490 | Err(error) |
| 2491 | } |
| 2492 | } |
| 2493 | |
| 2494 | /// Resolve the operator's proxy route for this exact request using the same |
| 2495 | /// matcher as reqwest. A NO_PROXY match returns None: direct requests must keep |
| 2496 | /// their public DNS validation and pins. Model and reviewed-plugin requests |
| 2497 | /// return before even reading proxy credentials. |
| 2498 | fn configured_mcp_proxy<F>( |
| 2499 | url: &reqwest::Url, |
| 2500 | disallow_ambient_proxy: bool, |
| 2501 | mut read_environment: F, |
| 2502 | ) -> Result<Option<reqwest::Proxy>> |
| 2503 | where |
| 2504 | F: FnMut(&str) -> std::result::Result<String, std::env::VarError>, |
| 2505 | { |
| 2506 | if disallow_ambient_proxy { |
| 2507 | return Ok(None); |
| 2508 | } |
| 2509 | let proxy_url = read_environment("HTTPS_PROXY") |
| 2510 | .or_else(|_| read_environment("https_proxy")) |
| 2511 | .or_else(|_| read_environment("HTTP_PROXY")) |
| 2512 | .or_else(|_| read_environment("http_proxy")) |
| 2513 | .ok() |
| 2514 | .filter(|value| !value.trim().is_empty()); |
| 2515 | let Some(proxy_url) = proxy_url else { |
| 2516 | return Ok(None); |
| 2517 | }; |
| 2518 | // Normalize userinfo and Unicode with the URL parser before passing the |
| 2519 | // URL to reqwest's own underlying matcher. Keep its missing-scheme support. |
| 2520 | let normalized = reqwest::Url::parse(&proxy_url) |
| 2521 | .ok() |
| 2522 | .filter(|url| url.has_host()) |
| 2523 | .or_else(|| reqwest::Url::parse(&format!("http://{proxy_url}")).ok()); |
| 2524 | let Some(normalized) = normalized else { |
| 2525 | tracing::warn!(target: "mcp", proxy = %redact_proxy_userinfo(&proxy_url), "ignoring malformed HTTP(S)_PROXY URL"); |
| 2526 | return Ok(None); |
| 2527 | }; |
| 2528 | let no_proxy = read_environment("NO_PROXY") |
| 2529 | .or_else(|_| read_environment("no_proxy")) |
| 2530 | .unwrap_or_default(); |
| 2531 | let matcher = hyper_util::client::proxy::matcher::Matcher::builder() |
| 2532 | .all(normalized.as_str()) |
| 2533 | .no(no_proxy) |
| 2534 | .build(); |
| 2535 | let destination: oauth2::http::Uri = url.as_str().parse()?; |
| 2536 | let Some(route) = matcher.intercept(&destination) else { |
| 2537 | return Ok(None); |
| 2538 | }; |
| 2539 | // Build the actual proxy from the matched route itself so the decision |
| 2540 | // that grants delegated DNS authority cannot diverge from the transport. |
| 2541 | let mut proxy = reqwest::Proxy::all(route.uri().to_string())?; |
| 2542 | if let Some(auth) = route.basic_auth() { |
| 2543 | proxy = proxy.custom_http_auth(auth.clone()); |
| 2544 | } |
| 2545 | if let Some((user, password)) = route.raw_auth() { |
| 2546 | proxy = proxy.basic_auth(user, password); |
| 2547 | } |
| 2548 | Ok(Some(proxy)) |
| 2549 | } |
| 2550 | |
| 2551 | impl Drop for McpConnection { |
| 2552 | fn drop(&mut self) { |
| 2553 | self.cancel_token.cancel(); |
| 2554 | if let Some(watch) = self.authority_watch.take() { |
| 2555 | watch.abort(); |
| 2556 | } |
| 2557 | } |
| 2558 | } |
| 2559 | |
| 2560 | // === McpPool - Connection Pool Management === |
| 2561 | |
| 2562 | #[derive(Debug, Clone)] |
| 2563 | struct McpToolRoute { |
| 2564 | server_name: String, |
| 2565 | tool_name: String, |
| 2566 | catalog_generation: u64, |
| 2567 | plugin_authority: Option<crate::plugins::types::PluginAuthority>, |
| 2568 | } |
| 2569 | |
| 2570 | /// Model-facing name suffix of the synthetic self-serve OAuth login tool |
| 2571 | /// (`mcp_<server>_authenticate`). Registered by [`McpPool::to_api_tools`] for |
| 2572 | /// servers whose last connect failed auth-required; executed by |
| 2573 | /// [`McpPool::call_tool`] through the same flow `/mcp login` uses. |
| 2574 | pub(crate) const AUTHENTICATE_TOOL_NAME: &str = "authenticate"; |
| 2575 | |
| 2576 | /// Result of [`McpPool::begin_authenticate_tool`]: either the shared token |
| 2577 | /// store already holds a usable credential (a login completed elsewhere since |
| 2578 | /// the catalog was built) or a browser login has been started and must be |
| 2579 | /// finished outside the pool lock. |
| 2580 | pub(crate) enum AuthenticateToolStart { |
| 2581 | AlreadyAuthorized, |
| 2582 | Login(Box<oauth::McpOAuthToolLogin>), |
| 2583 | } |
| 2584 | |
| 2585 | /// How the login phase of the synthetic authenticate tool concluded, fed to |
| 2586 | /// [`McpPool::finish_authenticate_tool`]. |
| 2587 | pub(crate) enum AuthenticateToolOutcome { |
| 2588 | AlreadyAuthorized, |
| 2589 | Authenticated { authorization_url: String }, |
| 2590 | } |
| 2591 | |
| 2592 | /// Execute the synthetic `mcp_<server>_authenticate` tool against a shared |
| 2593 | /// pool without holding the pool lock during the browser wait: lock to start |
| 2594 | /// the flow, release, wait for the loopback callback, then lock again to |
| 2595 | /// reconnect. `on_authorization_url` fires as soon as the URL exists — before |
| 2596 | /// the wait — so the runtime can show it to the user while the call blocks; |
| 2597 | /// the model only sees the URL in the result, after the flow has already |
| 2598 | /// finished, so this hook is the user's real path to the sign-in page when |
| 2599 | /// the browser did not open. |
| 2600 | pub(crate) async fn authenticate_tool_via_pool( |
| 2601 | pool: &Arc<tokio::sync::Mutex<McpPool>>, |
| 2602 | server_name: &str, |
| 2603 | on_authorization_url: impl FnOnce(&str), |
| 2604 | ) -> Result<serde_json::Value> { |
| 2605 | let start = pool |
| 2606 | .lock() |
| 2607 | .await |
| 2608 | .begin_authenticate_tool(server_name) |
| 2609 | .await?; |
| 2610 | let outcome = match start { |
| 2611 | AuthenticateToolStart::AlreadyAuthorized => AuthenticateToolOutcome::AlreadyAuthorized, |
| 2612 | AuthenticateToolStart::Login(login) => { |
| 2613 | let authorization_url = login.authorization_url().to_string(); |
| 2614 | on_authorization_url(&authorization_url); |
| 2615 | login.finish().await?; |
| 2616 | AuthenticateToolOutcome::Authenticated { authorization_url } |
| 2617 | } |
| 2618 | }; |
| 2619 | pool.lock() |
| 2620 | .await |
| 2621 | .finish_authenticate_tool(server_name, outcome) |
| 2622 | .await |
| 2623 | } |
| 2624 | |
| 2625 | /// Pool of MCP connections for reuse |
| 2626 | pub struct McpPool { |
| 2627 | /// Immutable operator ceiling; source reloads and shared child pools cannot relax it. |
| 2628 | disallowed_tools: Vec<String>, |
| 2629 | connections: HashMap<String, McpConnection>, |
| 2630 | config: McpConfig, |
| 2631 | network_policy: Option<NetworkPolicyDecider>, |
| 2632 | /// Source paths the config was loaded from. Empty for pools constructed |
| 2633 | /// directly via `new` (tests, ad-hoc snapshots). Workspace-aware pools |
| 2634 | /// track both global and project-level MCP config paths so lazy reload sees |
| 2635 | /// either file appear or change. |
| 2636 | config_sources: Vec<PathBuf>, |
| 2637 | workspace: Option<PathBuf>, |
| 2638 | plugin_registry: Option<Arc<crate::plugins::PluginRegistry>>, |
| 2639 | /// 64-bit content hash of the active config (`hash_mcp_config`). Compared |
| 2640 | /// against the freshly-loaded config after an mtime change to skip |
| 2641 | /// reloading when the file was merely touched. |
| 2642 | config_hash: u64, |
| 2643 | /// Monotonic identity for the exact config/plugin catalog generation that |
| 2644 | /// advertised a callable MCP item. Resolution captures this value and the |
| 2645 | /// call boundary rejects any intervening lazy reload or dynamic mutation. |
| 2646 | catalog_generation: AtomicU64, |
| 2647 | /// Most recently observed mtime for `config_sources`. |
| 2648 | last_mtimes: Vec<Option<std::time::SystemTime>>, |
| 2649 | /// Dynamically added MCP servers (from tool calls at runtime). |
| 2650 | /// These are not persisted to disk and live for the process lifetime. |
| 2651 | pub(crate) dynamic_servers: Arc<RwLock<HashMap<String, McpServerConfig>>>, |
| 2652 | /// Servers whose most recent connect attempt failed auth-required (401 / |
| 2653 | /// OAuth not logged in). Each gets a synthetic `mcp_<server>_authenticate` |
| 2654 | /// tool in the model catalog so the model can self-serve the OAuth login |
| 2655 | /// instead of dead-ending on the error item. BTreeSet keeps catalog |
| 2656 | /// construction deterministic. |
| 2657 | needs_auth_servers: BTreeSet<String>, |
| 2658 | /// Configured OAuth callback overrides, so the synthetic self-serve |
| 2659 | /// login tool honors the same pre-registered redirect URI `/mcp login` |
| 2660 | /// uses (`mcp_oauth_callback_port` / `mcp_oauth_callback_url`). |
| 2661 | oauth_callback_port: Option<u16>, |
| 2662 | oauth_callback_url: Option<String>, |
| 2663 | /// Bumped on every `needs_auth_servers` mutation. The engine reads it |
| 2664 | /// around a tool call so a live 401 that flips the auth surface can |
| 2665 | /// flag `mcp_catalog_changed` on the failed result and the turn loop |
| 2666 | /// replaces the pool's catalog slice before the next model request. |
| 2667 | needs_auth_generation: u64, |
| 2668 | /// Per-server cooldown after a failed connect, keyed by server name. |
| 2669 | /// |
| 2670 | /// The turn loop rebuilds the tool catalog on every user message, and |
| 2671 | /// that used to re-attempt every server that was not ready — so twenty |
| 2672 | /// configured servers with four dead ones paid four connect timeouts |
| 2673 | /// before the first token, every single turn, forever. A failure now |
| 2674 | /// buys a growing cooldown; the recorded diagnosis is replayed while it |
| 2675 | /// holds, so a skipped server still reads as failing and never as |
| 2676 | /// healthy. Explicit intent (`retry_connection`, `get_or_connect`, a |
| 2677 | /// config reload) ignores the cooldown. |
| 2678 | connect_backoff: HashMap<String, ConnectBackoff>, |
| 2679 | /// Servers the supervisor last saw dead. Death is reported once, on the |
| 2680 | /// transition, so status surfaces flip exactly when liveness does instead |
| 2681 | /// of re-emitting every sweep (#6187). |
| 2682 | supervised_dead: HashSet<String>, |
| 2683 | /// Servers the supervisor stopped auto-reconnecting after |
| 2684 | /// [`SUPERVISOR_PARK_AFTER_CONSECUTIVE_FAILURES`] consecutive failures. |
| 2685 | /// A stored-ready connection or an explicit `/mcp retry` clears the park. |
| 2686 | supervised_parked: HashSet<String>, |
| 2687 | /// Servers with a spawned connect in flight right now. `connect_all`, |
| 2688 | /// the session boot pass, and explicit tool-selection connects all mark |
| 2689 | /// names here and clear them on resolution, so status surfaces never |
| 2690 | /// have to infer "connecting" from "enabled but not connected yet" |
| 2691 | /// (#6033): under lazy boot an unconnected server is one nobody has |
| 2692 | /// asked for, not one mid-handshake. |
| 2693 | connecting: HashSet<String>, |
| 2694 | } |
| 2695 | |
| 2696 | /// One server's cooldown: when to try again, and what to say until then. |
| 2697 | struct ConnectBackoff { |
| 2698 | consecutive_failures: u32, |
| 2699 | retry_after: std::time::Instant, |
| 2700 | last_error: String, |
| 2701 | } |
| 2702 | |
| 2703 | /// One supervised reconnect candidate: the name, the config to redial, and |
| 2704 | /// whether this sweep newly observed the death. |
| 2705 | pub(crate) struct SupervisionDue { |
| 2706 | pub name: String, |
| 2707 | pub config: McpServerConfig, |
| 2708 | pub fresh_death: bool, |
| 2709 | } |
| 2710 | |
| 2711 | /// One supervisor sweep's plan: candidates to redial plus the transitions |
| 2712 | /// the plan phase already knows (recoveries and newly parked servers). |
| 2713 | pub(crate) struct SupervisionPlan { |
| 2714 | pub due: Vec<SupervisionDue>, |
| 2715 | pub recovered: Vec<String>, |
| 2716 | pub parked: Vec<String>, |
| 2717 | pub timeouts: McpTimeouts, |
| 2718 | pub network_policy: Option<NetworkPolicyDecider>, |
| 2719 | pub catalog_generation: u64, |
| 2720 | } |
| 2721 | |
| 2722 | /// One supervisor sweep's transitions. Death, recovery, failed attempts, and |
| 2723 | /// parking are reported on transition only, so the engine emits a snapshot |
| 2724 | /// update exactly when something changed (#6187). |
| 2725 | #[derive(Debug, Default)] |
| 2726 | pub(crate) struct McpSupervisorUpdate { |
| 2727 | /// Newly observed dead, with the reconnect failure that confirmed it. |
| 2728 | pub died: Vec<(String, String)>, |
| 2729 | /// Reconnect attempt failed for an already-dead server, with last error. |
| 2730 | pub failed: Vec<(String, String)>, |
| 2731 | /// Dead last sweep, alive now. |
| 2732 | pub recovered: Vec<String>, |
| 2733 | /// Newly parked after repeated failures; explicit `/mcp retry` resumes. |
| 2734 | pub parked: Vec<String>, |
| 2735 | } |
| 2736 | |
| 2737 | impl McpSupervisorUpdate { |
| 2738 | pub(crate) fn is_empty(&self) -> bool { |
| 2739 | self.died.is_empty() |
| 2740 | && self.failed.is_empty() |
| 2741 | && self.recovered.is_empty() |
| 2742 | && self.parked.is_empty() |
| 2743 | } |
| 2744 | |
| 2745 | fn merge(&mut self, other: McpSupervisorUpdate) { |
| 2746 | self.died.extend(other.died); |
| 2747 | self.failed.extend(other.failed); |
| 2748 | self.recovered.extend(other.recovered); |
| 2749 | self.parked.extend(other.parked); |
| 2750 | } |
| 2751 | } |
| 2752 | |
| 2753 | /// Cooldown after `failures` consecutive failed connects. |
| 2754 | /// |
| 2755 | /// Doubling from 30s to a 10-minute ceiling: long enough that a wall of dead |
| 2756 | /// servers costs nothing per turn, short enough that a server coming back |
| 2757 | /// (a laptop rejoining a network, a local server restarted) is picked up |
| 2758 | /// within one coffee break without the user touching anything. |
| 2759 | fn connect_backoff_delay(failures: u32) -> std::time::Duration { |
| 2760 | const BASE: std::time::Duration = std::time::Duration::from_secs(30); |
| 2761 | const CAP: std::time::Duration = std::time::Duration::from_secs(600); |
| 2762 | BASE.saturating_mul(1u32 << failures.saturating_sub(1).min(5)) |
| 2763 | .min(CAP) |
| 2764 | } |
| 2765 | |
| 2766 | type McpPendingConnect = (String, McpServerConfig); |
| 2767 | type McpConnectError = (String, anyhow::Error); |
| 2768 | |
| 2769 | /// Whether an explicit tool selection (`tools_always_load`, a turn's |
| 2770 | /// `allowed_tools`) covers `server`: either an exact `mcp_<server>_<tool>` |
| 2771 | /// name or an `mcp_<prefix>*` glob whose prefix reaches the server name. |
| 2772 | /// One definition shared by the lazy boot pass and the per-turn |
| 2773 | /// explicit-connect wait so both agree on what a selection starts (#6033). |
| 2774 | pub(crate) fn tool_selection_covers_server(requested: &[String], server: &str) -> bool { |
| 2775 | let prefix = format!("mcp_{}_", server.to_ascii_lowercase()); |
| 2776 | requested.iter().any(|name| { |
| 2777 | name.starts_with(&prefix) |
| 2778 | || name |
| 2779 | .strip_suffix('*') |
| 2780 | .is_some_and(|rule| prefix.starts_with(rule)) |
| 2781 | }) |
| 2782 | } |
| 2783 | |
| 2784 | impl McpPool { |
| 2785 | /// Create a new pool with the given configuration |
| 2786 | pub fn new(config: McpConfig) -> Self { |
| 2787 | let config_hash = hash_mcp_config(&config); |
| 2788 | Self { |
| 2789 | connections: HashMap::new(), |
| 2790 | disallowed_tools: Vec::new(), |
| 2791 | config, |
| 2792 | network_policy: None, |
| 2793 | oauth_callback_port: None, |
| 2794 | oauth_callback_url: None, |
| 2795 | config_sources: Vec::new(), |
| 2796 | workspace: None, |
| 2797 | plugin_registry: None, |
| 2798 | config_hash, |
| 2799 | catalog_generation: AtomicU64::new(1), |
| 2800 | connect_backoff: HashMap::new(), |
| 2801 | supervised_dead: HashSet::new(), |
| 2802 | supervised_parked: HashSet::new(), |
| 2803 | connecting: HashSet::new(), |
| 2804 | last_mtimes: Vec::new(), |
| 2805 | dynamic_servers: Arc::new(RwLock::new(HashMap::new())), |
| 2806 | needs_auth_servers: BTreeSet::new(), |
| 2807 | needs_auth_generation: 0, |
| 2808 | } |
| 2809 | } |
| 2810 | |
| 2811 | /// Create a pool from a configuration file path. |
| 2812 | #[cfg(test)] |
| 2813 | pub fn from_config_path(path: &std::path::Path) -> Result<Self> { |
| 2814 | let config = load_config(path)?; |
| 2815 | let mut pool = Self::new(config); |
| 2816 | pool.config_sources = vec![path.to_path_buf()]; |
| 2817 | pool.last_mtimes = vec![mcp_config_mtime(path)]; |
| 2818 | Ok(pool) |
| 2819 | } |
| 2820 | |
| 2821 | /// Create a pool from global MCP config plus workspace-local |
| 2822 | /// `.codewhale/mcp.json`. Project servers override same-name global |
| 2823 | /// servers and default stdio `cwd` to the workspace root. |
| 2824 | #[cfg(test)] |
| 2825 | pub fn from_config_path_with_workspace( |
| 2826 | path: &std::path::Path, |
| 2827 | workspace: &Path, |
| 2828 | ) -> Result<Self> { |
| 2829 | let plugins = Arc::new(crate::plugins::PluginRegistry::empty(workspace)); |
| 2830 | Self::from_config_path_with_workspace_and_plugins(path, workspace, plugins) |
| 2831 | } |
| 2832 | |
| 2833 | pub fn from_config_path_with_workspace_and_plugins( |
| 2834 | path: &std::path::Path, |
| 2835 | workspace: &Path, |
| 2836 | plugins: Arc<crate::plugins::PluginRegistry>, |
| 2837 | ) -> Result<Self> { |
| 2838 | if plugins.workspace() != workspace { |
| 2839 | anyhow::bail!("plugin registry workspace does not match MCP pool workspace"); |
| 2840 | } |
| 2841 | let config = load_config_with_workspace_and_plugins(path, workspace, plugins.as_ref())?; |
| 2842 | let workspace = checked_workspace_path(workspace)?; |
| 2843 | let mut pool = Self::new(config); |
| 2844 | pool.config_sources = vec![ |
| 2845 | path.to_path_buf(), |
| 2846 | checked_workspace_mcp_config_path(&workspace)?, |
| 2847 | ]; |
| 2848 | pool.config_sources |
| 2849 | .extend(crate::config::workspace_trust_config_candidate_paths()); |
| 2850 | pool.last_mtimes = pool |
| 2851 | .config_sources |
| 2852 | .iter() |
| 2853 | .map(|source| mcp_config_mtime(source)) |
| 2854 | .collect(); |
| 2855 | pool.workspace = Some(workspace); |
| 2856 | pool.plugin_registry = Some(plugins); |
| 2857 | Ok(pool) |
| 2858 | } |
| 2859 | |
| 2860 | /// Construct a source-aware empty pool after the initial config load |
| 2861 | /// failed. Keeping the source paths means a later edit or explicit |
| 2862 | /// `/mcp reload` can recover in-process instead of pinning the session to |
| 2863 | /// an ad-hoc pool that has no files to re-read. |
| 2864 | pub(crate) fn empty_with_workspace_config_sources( |
| 2865 | path: &std::path::Path, |
| 2866 | workspace: &Path, |
| 2867 | plugins: Arc<crate::plugins::PluginRegistry>, |
| 2868 | ) -> Result<Self> { |
| 2869 | validate_mcp_config_path(path)?; |
| 2870 | if plugins.workspace() != workspace { |
| 2871 | anyhow::bail!("plugin registry workspace does not match MCP pool workspace"); |
| 2872 | } |
| 2873 | let workspace = checked_workspace_path(workspace)?; |
| 2874 | let mut pool = Self::new(McpConfig::default()); |
| 2875 | pool.config_sources = vec![ |
| 2876 | path.to_path_buf(), |
| 2877 | checked_workspace_mcp_config_path(&workspace)?, |
| 2878 | ]; |
| 2879 | pool.config_sources |
| 2880 | .extend(crate::config::workspace_trust_config_candidate_paths()); |
| 2881 | pool.last_mtimes = pool |
| 2882 | .config_sources |
| 2883 | .iter() |
| 2884 | .map(|source| mcp_config_mtime(source)) |
| 2885 | .collect(); |
| 2886 | pool.workspace = Some(workspace); |
| 2887 | pool.plugin_registry = Some(plugins); |
| 2888 | Ok(pool) |
| 2889 | } |
| 2890 | |
| 2891 | /// Install the session ceiling before any connection or model catalog is exposed. |
| 2892 | pub(crate) fn with_disallowed_tools(mut self, rules: Vec<String>) -> Self { |
| 2893 | self.disallowed_tools.extend(rules); |
| 2894 | self |
| 2895 | } |
| 2896 | |
| 2897 | /// Only a prefix covering the entire namespace suppresses a server. An |
| 2898 | /// individual tool denial must preserve its siblings and resource access. |
| 2899 | pub(crate) fn server_denied_by(rules: &[String], server: &str) -> bool { |
| 2900 | let namespace = format!("mcp_{server}_").to_ascii_lowercase(); |
| 2901 | rules.iter().any(|rule| { |
| 2902 | rule.to_ascii_lowercase() |
| 2903 | .strip_suffix('*') |
| 2904 | .is_some_and(|prefix| namespace.starts_with(prefix)) |
| 2905 | }) |
| 2906 | } |
| 2907 | |
| 2908 | fn server_allowed(&self, server: &str) -> bool { |
| 2909 | !Self::server_denied_by(&self.disallowed_tools, server) |
| 2910 | } |
| 2911 | |
| 2912 | pub(crate) fn tool_allowed(&self, name: &str) -> bool { |
| 2913 | !crate::core::engine::tool_catalog::tool_matches_any_rule(&self.disallowed_tools, name) |
| 2914 | } |
| 2915 | |
| 2916 | fn require_server(&self, server: &str) -> Result<()> { |
| 2917 | anyhow::ensure!( |
| 2918 | self.server_allowed(server), |
| 2919 | "Failed to find MCP server: {server}" |
| 2920 | ); |
| 2921 | Ok(()) |
| 2922 | } |
| 2923 | |
| 2924 | pub(crate) fn authorize_call( |
| 2925 | rules: &[String], |
| 2926 | name: &str, |
| 2927 | input: &serde_json::Value, |
| 2928 | ) -> Result<()> { |
| 2929 | anyhow::ensure!( |
| 2930 | !crate::core::engine::tool_catalog::tool_matches_any_rule(rules, name), |
| 2931 | "Unknown MCP tool name: {name}" |
| 2932 | ); |
| 2933 | if matches!( |
| 2934 | name, |
| 2935 | "list_mcp_resources" |
| 2936 | | "list_mcp_resource_templates" |
| 2937 | | "mcp_read_resource" |
| 2938 | | "read_mcp_resource" |
| 2939 | | "mcp_get_prompt" |
| 2940 | ) && let Some(server) = input.get("server").and_then(serde_json::Value::as_str) |
| 2941 | { |
| 2942 | anyhow::ensure!( |
| 2943 | !Self::server_denied_by(rules, server), |
| 2944 | "Failed to find MCP server: {server}" |
| 2945 | ); |
| 2946 | } |
| 2947 | Ok(()) |
| 2948 | } |
| 2949 | |
| 2950 | /// Attach a per-domain network policy (#135). When set, HTTP/SSE |
| 2951 | /// transports are gated through it; STDIO transports are unaffected. |
| 2952 | pub fn with_network_policy(mut self, policy: NetworkPolicyDecider) -> Self { |
| 2953 | self.network_policy = Some(policy); |
| 2954 | self |
| 2955 | } |
| 2956 | |
| 2957 | /// Configure the OAuth callback overrides (`mcp_oauth_callback_port` / |
| 2958 | /// `mcp_oauth_callback_url`) for installations whose OAuth client has a |
| 2959 | /// pre-registered redirect URI. The synthetic self-serve login tool must |
| 2960 | /// use the same overrides as `/mcp login`, or the provider rejects its |
| 2961 | /// ephemeral loopback redirect. |
| 2962 | pub fn with_oauth_callback(mut self, port: Option<u16>, url: Option<String>) -> Self { |
| 2963 | self.oauth_callback_port = port; |
| 2964 | self.oauth_callback_url = url; |
| 2965 | self |
| 2966 | } |
| 2967 | |
| 2968 | pub(crate) fn connect_timeouts(&self) -> McpTimeouts { |
| 2969 | self.config.timeouts |
| 2970 | } |
| 2971 | |
| 2972 | pub(crate) fn cloned_network_policy(&self) -> Option<NetworkPolicyDecider> { |
| 2973 | self.network_policy.clone() |
| 2974 | } |
| 2975 | |
| 2976 | pub(crate) fn current_catalog_generation(&self) -> u64 { |
| 2977 | self.catalog_generation.load(Ordering::SeqCst) |
| 2978 | } |
| 2979 | |
| 2980 | fn drop_connection(&mut self, server_name: &str, reason: &str) { |
| 2981 | if self.connections.remove(server_name).is_some() { |
| 2982 | tracing::debug!( |
| 2983 | target: "mcp", |
| 2984 | server = %server_name, |
| 2985 | reason = %reason, |
| 2986 | "dropped MCP connection" |
| 2987 | ); |
| 2988 | } |
| 2989 | } |
| 2990 | |
| 2991 | fn drop_all_connections(&mut self, reason: &str) { |
| 2992 | // Auth state is only known from a live connect attempt; once every |
| 2993 | // connection is dropped (config reload, source switch, shutdown) the |
| 2994 | // next attempt re-derives it. A reload is explicit intent, so every |
| 2995 | // cooldown lifts with it. |
| 2996 | self.connect_backoff.clear(); |
| 2997 | self.needs_auth_servers.clear(); |
| 2998 | self.needs_auth_generation = self.needs_auth_generation.wrapping_add(1); |
| 2999 | if self.connections.is_empty() { |
| 3000 | return; |
| 3001 | } |
| 3002 | let count = self.connections.len(); |
| 3003 | tracing::debug!( |
| 3004 | target: "mcp", |
| 3005 | count, |
| 3006 | reason = %reason, |
| 3007 | "dropping MCP connections" |
| 3008 | ); |
| 3009 | self.connections.clear(); |
| 3010 | } |
| 3011 | |
| 3012 | /// If the source config file's mtime has changed since the last check, |
| 3013 | /// re-read it and (only when the content hash also changed) drop all |
| 3014 | /// existing connections so the next `get_or_connect` reattaches under |
| 3015 | /// the new config. No-op when the pool was constructed via [`McpPool::new`] |
| 3016 | /// (no source path), when stat fails, or when the file content is |
| 3017 | /// byte-identical to what we last loaded. Returns `Ok(true)` if any |
| 3018 | /// connections were dropped, `Ok(false)` otherwise. |
| 3019 | /// |
| 3020 | /// This is the lazy half of the auto-reload story for #1267: instead of a |
| 3021 | /// long-lived file watcher, the next tool invocation pays a single `stat` |
| 3022 | /// call (and only re-reads the file when the mtime moved). On networked |
| 3023 | /// or remote filesystems where mtime granularity is poor, the hash |
| 3024 | /// compare keeps us from churning connections on every check. |
| 3025 | fn reload_from_config_sources(&mut self, force: bool) -> Result<bool> { |
| 3026 | if self.config_sources.is_empty() { |
| 3027 | if force { |
| 3028 | anyhow::bail!("MCP pool has no configuration source to reload"); |
| 3029 | } |
| 3030 | return Ok(false); |
| 3031 | } |
| 3032 | let current_mtimes: Vec<_> = self |
| 3033 | .config_sources |
| 3034 | .iter() |
| 3035 | .map(|path| mcp_config_mtime(path)) |
| 3036 | .collect(); |
| 3037 | if !force && current_mtimes == self.last_mtimes { |
| 3038 | return Ok(false); |
| 3039 | } |
| 3040 | // An mtime moved, or the user explicitly requested a reload: re-read |
| 3041 | // the complete global + workspace + plugin-backed config. |
| 3042 | let primary = self |
| 3043 | .config_sources |
| 3044 | .first() |
| 3045 | .context("MCP config source list unexpectedly empty")?; |
| 3046 | let new_config = if let Some(workspace) = self.workspace.as_deref() { |
| 3047 | match self.plugin_registry.as_deref() { |
| 3048 | Some(plugins) => { |
| 3049 | load_config_with_workspace_and_plugins(primary, workspace, plugins)? |
| 3050 | } |
| 3051 | None => load_config_with_workspace(primary, workspace)?, |
| 3052 | } |
| 3053 | } else { |
| 3054 | load_config(primary)? |
| 3055 | }; |
| 3056 | let new_hash = hash_mcp_config(&new_config); |
| 3057 | // Always advance mtimes so a touched-but-unchanged file doesn't |
| 3058 | // make us re-read on every subsequent call. |
| 3059 | self.last_mtimes = current_mtimes; |
| 3060 | if !force && new_hash == self.config_hash { |
| 3061 | return Ok(false); |
| 3062 | } |
| 3063 | // A real content change, or an explicit reload, invalidates every |
| 3064 | // advertised route and live transport. The latter matters when OAuth |
| 3065 | // credentials changed without changing the config bytes. |
| 3066 | self.drop_all_connections(if force { |
| 3067 | "explicit config reload" |
| 3068 | } else { |
| 3069 | "config reload" |
| 3070 | }); |
| 3071 | self.config = new_config; |
| 3072 | self.config_hash = new_hash; |
| 3073 | self.catalog_generation.fetch_add(1, Ordering::SeqCst); |
| 3074 | Ok(true) |
| 3075 | } |
| 3076 | |
| 3077 | pub async fn reload_if_config_changed(&mut self) -> Result<bool> { |
| 3078 | self.reload_from_config_sources(false) |
| 3079 | } |
| 3080 | |
| 3081 | /// Force a source re-read and drop every live connection so the next |
| 3082 | /// connect pass reattaches under the current configuration and |
| 3083 | /// credentials — without waiting for any handshake. An explicit reload is |
| 3084 | /// intent, so cooldowns lift and even a byte-identical config re-dials. |
| 3085 | /// |
| 3086 | /// An unreadable or malformed source returns `Err` **before** anything is |
| 3087 | /// dropped: a failed reload leaves the live tool pool intact. Dynamic |
| 3088 | /// in-memory servers remain registered because this mutates the existing |
| 3089 | /// pool rather than replacing it. |
| 3090 | pub(crate) fn force_reload_config_sources(&mut self) -> Result<()> { |
| 3091 | self.reload_from_config_sources(true).map(|_| ()) |
| 3092 | } |
| 3093 | |
| 3094 | /// Install a replacement global config source transactionally, preserving |
| 3095 | /// this shared pool (and its dynamic runtime servers) for parent and |
| 3096 | /// sub-agent holders. A malformed replacement leaves the current config, |
| 3097 | /// connections, and source paths unchanged. On success every live |
| 3098 | /// connection is dropped; the caller reattaches through its own connect |
| 3099 | /// pass so no pool lock is held across a handshake. |
| 3100 | pub(crate) fn switch_workspace_config_source( |
| 3101 | &mut self, |
| 3102 | path: &Path, |
| 3103 | workspace: &Path, |
| 3104 | plugins: Arc<crate::plugins::PluginRegistry>, |
| 3105 | ) -> Result<()> { |
| 3106 | validate_mcp_config_path(path)?; |
| 3107 | if plugins.workspace() != workspace { |
| 3108 | anyhow::bail!("plugin registry workspace does not match MCP pool workspace"); |
| 3109 | } |
| 3110 | let workspace = checked_workspace_path(workspace)?; |
| 3111 | let new_config = |
| 3112 | load_config_with_workspace_and_plugins(path, &workspace, plugins.as_ref())?; |
| 3113 | let mut new_sources = vec![ |
| 3114 | path.to_path_buf(), |
| 3115 | checked_workspace_mcp_config_path(&workspace)?, |
| 3116 | ]; |
| 3117 | new_sources.extend(crate::config::workspace_trust_config_candidate_paths()); |
| 3118 | let new_mtimes = new_sources |
| 3119 | .iter() |
| 3120 | .map(|source| mcp_config_mtime(source)) |
| 3121 | .collect(); |
| 3122 | |
| 3123 | self.drop_all_connections("config source switch"); |
| 3124 | self.config_hash = hash_mcp_config(&new_config); |
| 3125 | self.config = new_config; |
| 3126 | self.config_sources = new_sources; |
| 3127 | self.last_mtimes = new_mtimes; |
| 3128 | self.workspace = Some(workspace); |
| 3129 | self.plugin_registry = Some(plugins); |
| 3130 | self.catalog_generation.fetch_add(1, Ordering::SeqCst); |
| 3131 | Ok(()) |
| 3132 | } |
| 3133 | |
| 3134 | /// Get or create a connection to a server |
| 3135 | pub async fn get_or_connect(&mut self, server_name: &str) -> Result<&mut McpConnection> { |
| 3136 | self.require_server(server_name)?; |
| 3137 | // Lazy auto-reload (#1267 part 2): cheap mtime-then-hash check before |
| 3138 | // each connection lookup. Transient FS errors are logged but not |
| 3139 | // propagated so a brief hiccup can't take down the whole tool dispatch. |
| 3140 | if let Err(e) = self.reload_if_config_changed().await { |
| 3141 | tracing::warn!("MCP config reload check failed: {e:#}"); |
| 3142 | } |
| 3143 | |
| 3144 | let plugin_source = self |
| 3145 | .connections |
| 3146 | .get(server_name) |
| 3147 | .and_then(|connection| connection.config().reviewed_plugin.clone()) |
| 3148 | .or_else(|| { |
| 3149 | self.config |
| 3150 | .servers |
| 3151 | .get(server_name) |
| 3152 | .and_then(|config| config.reviewed_plugin.clone()) |
| 3153 | }); |
| 3154 | if let Some(source) = plugin_source |
| 3155 | && let Err(error) = source.validate_before_use(server_name, "use") |
| 3156 | { |
| 3157 | self.drop_connection(server_name, "plugin authority revoked or changed"); |
| 3158 | return Err(error); |
| 3159 | } |
| 3160 | |
| 3161 | // Authority was just validated above for this same source; checking |
| 3162 | // it again here would re-hash the bundle within one dispatch (#6209). |
| 3163 | let is_ready = self |
| 3164 | .connections |
| 3165 | .get(server_name) |
| 3166 | .map(McpConnection::is_transport_ready) |
| 3167 | .unwrap_or(false); |
| 3168 | if is_ready { |
| 3169 | return self |
| 3170 | .connections |
| 3171 | .get_mut(server_name) |
| 3172 | .ok_or_else(|| anyhow::anyhow!("MCP connection disappeared for {server_name}")); |
| 3173 | } |
| 3174 | |
| 3175 | // Take (don't drop) the stale connection: if the reconnect attempt |
| 3176 | // below fails, the previous connection is restored so its last-good |
| 3177 | // tool catalog stays model-visible during the outage instead of |
| 3178 | // disappearing with a dropped transport (#6187). |
| 3179 | let previous_connection = self.connections.remove(server_name); |
| 3180 | if previous_connection.is_some() { |
| 3181 | tracing::debug!( |
| 3182 | target: "mcp", |
| 3183 | server = %server_name, |
| 3184 | reason = "reconnect", |
| 3185 | "detached MCP connection for reconnect" |
| 3186 | ); |
| 3187 | } |
| 3188 | |
| 3189 | // Check static config first, then dynamic servers |
| 3190 | let server_config = self |
| 3191 | .config |
| 3192 | .servers |
| 3193 | .get(server_name) |
| 3194 | .cloned() |
| 3195 | .or_else(|| self.dynamic_servers.read().get(server_name).cloned()) |
| 3196 | .ok_or_else(|| anyhow::anyhow!("Failed to find MCP server: {server_name}"))?; |
| 3197 | |
| 3198 | if !server_config.is_enabled() { |
| 3199 | anyhow::bail!("Failed to connect MCP server '{server_name}': server is disabled"); |
| 3200 | } |
| 3201 | |
| 3202 | let mut connection = match McpConnection::connect_with_policy( |
| 3203 | server_name.to_string(), |
| 3204 | server_config, |
| 3205 | &self.config.timeouts, |
| 3206 | self.network_policy.as_ref(), |
| 3207 | ) |
| 3208 | .await |
| 3209 | { |
| 3210 | Ok(connection) => connection, |
| 3211 | Err(error) => { |
| 3212 | self.note_connect_failure(server_name, &error); |
| 3213 | if let Some(previous) = previous_connection { |
| 3214 | tracing::debug!( |
| 3215 | target: "mcp", |
| 3216 | server = %server_name, |
| 3217 | "reconnect failed; restored the previous MCP connection and its last-good catalog" |
| 3218 | ); |
| 3219 | self.connections.insert(server_name.to_string(), previous); |
| 3220 | } |
| 3221 | return Err(error); |
| 3222 | } |
| 3223 | }; |
| 3224 | connection.catalog_generation = self.catalog_generation.load(Ordering::SeqCst); |
| 3225 | |
| 3226 | self.store_ready_connection(server_name.to_string(), connection)?; |
| 3227 | self.connections |
| 3228 | .get_mut(server_name) |
| 3229 | .ok_or_else(|| anyhow::anyhow!("Failed to store MCP connection for {server_name}")) |
| 3230 | } |
| 3231 | |
| 3232 | /// Retry exactly one server against the configuration already owned by |
| 3233 | /// this pool. |
| 3234 | /// |
| 3235 | /// Unlike normal lazy tool dispatch, an explicit row retry must not notice |
| 3236 | /// a concurrent config mtime and invalidate healthy siblings. Config edits |
| 3237 | /// remain owned by the explicit reload path; this operation only replaces |
| 3238 | /// the named transport. |
| 3239 | pub async fn retry_connection(&mut self, server_name: &str) -> Result<&mut McpConnection> { |
| 3240 | self.require_server(server_name)?; |
| 3241 | // A person asked for this one by name. Clear the cooldown so the |
| 3242 | // attempt happens now and, if it fails again, the ladder restarts |
| 3243 | // from the short end rather than from wherever it had climbed to. |
| 3244 | // Explicit intent restarts supervision: the cooldown, the dead mark, |
| 3245 | // and any park all clear, so the supervisor resumes watching whatever |
| 3246 | // this retry stores — or stays quiet while the server is connectionless. |
| 3247 | self.connect_backoff.remove(server_name); |
| 3248 | self.supervised_dead.remove(server_name); |
| 3249 | self.supervised_parked.remove(server_name); |
| 3250 | let plugin_source = self |
| 3251 | .connections |
| 3252 | .get(server_name) |
| 3253 | .and_then(|connection| connection.config().reviewed_plugin.clone()) |
| 3254 | .or_else(|| { |
| 3255 | self.config |
| 3256 | .servers |
| 3257 | .get(server_name) |
| 3258 | .and_then(|config| config.reviewed_plugin.clone()) |
| 3259 | }); |
| 3260 | if let Some(source) = plugin_source |
| 3261 | && let Err(error) = source.validate_before_use(server_name, "use") |
| 3262 | { |
| 3263 | self.drop_connection(server_name, "plugin authority revoked or changed"); |
| 3264 | return Err(error); |
| 3265 | } |
| 3266 | |
| 3267 | self.drop_connection(server_name, "retry"); |
| 3268 | |
| 3269 | let server_config = self |
| 3270 | .config |
| 3271 | .servers |
| 3272 | .get(server_name) |
| 3273 | .cloned() |
| 3274 | .or_else(|| self.dynamic_servers.read().get(server_name).cloned()) |
| 3275 | .ok_or_else(|| anyhow::anyhow!("Failed to find MCP server: {server_name}"))?; |
| 3276 | |
| 3277 | if !server_config.is_enabled() { |
| 3278 | anyhow::bail!("Failed to connect MCP server '{server_name}': server is disabled"); |
| 3279 | } |
| 3280 | |
| 3281 | let mut connection = match McpConnection::connect_with_policy( |
| 3282 | server_name.to_string(), |
| 3283 | server_config, |
| 3284 | &self.config.timeouts, |
| 3285 | self.network_policy.as_ref(), |
| 3286 | ) |
| 3287 | .await |
| 3288 | { |
| 3289 | Ok(connection) => connection, |
| 3290 | Err(error) => { |
| 3291 | self.note_connect_failure(server_name, &error); |
| 3292 | return Err(error); |
| 3293 | } |
| 3294 | }; |
| 3295 | connection.catalog_generation = self.current_catalog_generation(); |
| 3296 | self.store_ready_connection(server_name.to_string(), connection)?; |
| 3297 | self.connections |
| 3298 | .get_mut(server_name) |
| 3299 | .ok_or_else(|| anyhow::anyhow!("Failed to store MCP connection for {server_name}")) |
| 3300 | } |
| 3301 | |
| 3302 | pub(crate) fn store_ready_connection( |
| 3303 | &mut self, |
| 3304 | name: String, |
| 3305 | connection: McpConnection, |
| 3306 | ) -> Result<()> { |
| 3307 | self.require_server(&name)?; |
| 3308 | anyhow::ensure!( |
| 3309 | connection.catalog_generation == self.current_catalog_generation(), |
| 3310 | "MCP configuration changed while connecting {name}; retry against the current config" |
| 3311 | ); |
| 3312 | if let Some(source) = connection.config().reviewed_plugin.as_ref() { |
| 3313 | source.validate_before_use(&name, "use")?; |
| 3314 | } |
| 3315 | // A successful connect settles the auth question for this server, |
| 3316 | // and the cooldown with it — plus any supervisor dead mark or park, |
| 3317 | // since a stored-ready connection is alive by construction. |
| 3318 | self.connecting.remove(&name); |
| 3319 | self.connect_backoff.remove(&name); |
| 3320 | self.supervised_dead.remove(&name); |
| 3321 | self.supervised_parked.remove(&name); |
| 3322 | if self.needs_auth_servers.remove(&name) { |
| 3323 | self.needs_auth_generation = self.needs_auth_generation.wrapping_add(1); |
| 3324 | } |
| 3325 | self.connections.insert(name, connection); |
| 3326 | Ok(()) |
| 3327 | } |
| 3328 | |
| 3329 | /// Record a connect failure's auth classification. When the failure looks |
| 3330 | /// like a missing/expired OAuth login, the next model catalog offers the |
| 3331 | /// synthetic `mcp_<server>_authenticate` tool so the model can self-serve |
| 3332 | /// the login instead of dead-ending on the error item. A non-auth |
| 3333 | /// failure replaces the verdict — the state is "the most recent connect |
| 3334 | /// failed auth-required", not "some connect once did". |
| 3335 | pub(crate) fn note_connect_failure(&mut self, name: &str, error: &anyhow::Error) { |
| 3336 | self.connecting.remove(name); |
| 3337 | if !self.server_allowed(name) { |
| 3338 | return; |
| 3339 | } |
| 3340 | let entry = self |
| 3341 | .connect_backoff |
| 3342 | .entry(name.to_string()) |
| 3343 | .or_insert(ConnectBackoff { |
| 3344 | consecutive_failures: 0, |
| 3345 | retry_after: std::time::Instant::now(), |
| 3346 | last_error: String::new(), |
| 3347 | }); |
| 3348 | entry.consecutive_failures = entry.consecutive_failures.saturating_add(1); |
| 3349 | entry.retry_after = |
| 3350 | std::time::Instant::now() + connect_backoff_delay(entry.consecutive_failures); |
| 3351 | entry.last_error = format_mcp_error_for_display(error); |
| 3352 | let changed = if oauth::error_looks_auth_required(error) { |
| 3353 | self.needs_auth_servers.insert(name.to_string()) |
| 3354 | } else { |
| 3355 | self.needs_auth_servers.remove(name) |
| 3356 | }; |
| 3357 | if changed { |
| 3358 | self.needs_auth_generation = self.needs_auth_generation.wrapping_add(1); |
| 3359 | } |
| 3360 | } |
| 3361 | |
| 3362 | /// Current needs-auth surface generation. Compare across a tool call to |
| 3363 | /// learn whether the call flipped a server into or out of the |
| 3364 | /// `◆ auth required` state (a live 401, or a login that landed). |
| 3365 | #[must_use] |
| 3366 | pub fn needs_auth_generation(&self) -> u64 { |
| 3367 | self.needs_auth_generation |
| 3368 | } |
| 3369 | |
| 3370 | /// Whether the server's most recent connect attempt failed auth-required |
| 3371 | /// (the typed `◆ auth required` state). Cleared by a successful connect |
| 3372 | /// and by any full connection drop (reload, source switch, shutdown). |
| 3373 | #[must_use] |
| 3374 | pub fn server_needs_auth(&self, name: &str) -> bool { |
| 3375 | self.server_allowed(name) && self.needs_auth_servers.contains(name) |
| 3376 | } |
| 3377 | |
| 3378 | /// The needs-auth server that owns a model tool name (`mcp_<server>_…`), |
| 3379 | /// if any: the server the model is trying to reach with a real tool name |
| 3380 | /// from a catalog built before its login lapsed. Longest configured name |
| 3381 | /// wins so `mcp_a_b_tool` routes to server `a_b` over `a`. |
| 3382 | fn needs_auth_server_for_tool_name(&self, prefixed_name: &str) -> Option<String> { |
| 3383 | let rest = prefixed_name.strip_prefix("mcp_")?; |
| 3384 | self.needs_auth_servers |
| 3385 | .iter() |
| 3386 | .filter(|server| { |
| 3387 | self.server_allowed(server) |
| 3388 | && self.tool_allowed(prefixed_name) |
| 3389 | && rest |
| 3390 | .strip_prefix(server.as_str()) |
| 3391 | .is_some_and(|suffix| suffix.starts_with('_')) |
| 3392 | }) |
| 3393 | .max_by_key(|server| server.len()) |
| 3394 | .cloned() |
| 3395 | } |
| 3396 | |
| 3397 | /// Peak concurrent spawn+handshake attempts. Uncapped, a config full of |
| 3398 | /// `npx` servers would start one node runtime per server at the same |
| 3399 | /// instant — a memory spike on low-end machines the sequential loop never |
| 3400 | /// produced. Eight keeps wall-clock wins (the connect timeout dominates) |
| 3401 | /// while bounding peak memory. |
| 3402 | const CONNECT_CONCURRENCY: usize = 8; |
| 3403 | |
| 3404 | /// Consecutive failed reconnects after which the supervisor parks a |
| 3405 | /// server instead of redialing it. The cooldown ladder already spaces |
| 3406 | /// attempts, but a server that never answers (wrong binary, dead port) |
| 3407 | /// should not burn a spawn+handshake every sweep forever. A stored-ready |
| 3408 | /// connection or an explicit `/mcp retry` clears the park — and every |
| 3409 | /// success resets the count, so an occasionally-crashing server keeps |
| 3410 | /// recovering instead of parking. |
| 3411 | const SUPERVISOR_PARK_AFTER_CONSECUTIVE_FAILURES: u32 = 5; |
| 3412 | |
| 3413 | /// Supervisor sweep cadence. Death is noticed within one tick; an idle |
| 3414 | /// tick costs one pool lock plus a `try_wait` per stdio child. |
| 3415 | const SUPERVISOR_TICK: std::time::Duration = std::time::Duration::from_secs(5); |
| 3416 | |
| 3417 | /// One supervisor sweep's reconnect candidates, computed under a brief |
| 3418 | /// pool lock. Handshakes run outside the lock via |
| 3419 | /// [`Self::spawn_pending_connects`], so a wedged server never blocks a |
| 3420 | /// live turn's pool access while it burns its connect timeout. |
| 3421 | pub(crate) fn plan_supervision(&mut self) -> SupervisionPlan { |
| 3422 | let dynamic = self.dynamic_servers.read(); |
| 3423 | let candidates: Vec<(String, McpServerConfig)> = self |
| 3424 | .config |
| 3425 | .servers |
| 3426 | .iter() |
| 3427 | .filter(|(name, server)| server.is_enabled() && self.server_allowed(name)) |
| 3428 | .map(|(name, server)| (name.clone(), server.clone())) |
| 3429 | .chain( |
| 3430 | dynamic |
| 3431 | .iter() |
| 3432 | .filter(|(_, server)| server.is_enabled()) |
| 3433 | .map(|(name, server)| (name.clone(), server.clone())), |
| 3434 | ) |
| 3435 | .collect(); |
| 3436 | drop(dynamic); |
| 3437 | let watched: HashSet<String> = candidates.iter().map(|(name, _)| name.clone()).collect(); |
| 3438 | // Silent prune: manual retries drop connections the supervisor never |
| 3439 | // re-spawns (on-demand reconnect owns connectionless servers), and |
| 3440 | // removed/disabled servers leave supervision without an event. |
| 3441 | self.supervised_dead |
| 3442 | .retain(|name| watched.contains(name) && self.connections.contains_key(name)); |
| 3443 | self.supervised_parked.retain(|name| watched.contains(name)); |
| 3444 | let mut due = Vec::new(); |
| 3445 | let mut recovered = Vec::new(); |
| 3446 | let mut parked = Vec::new(); |
| 3447 | let now = std::time::Instant::now(); |
| 3448 | for (name, config) in candidates { |
| 3449 | let Some(connection) = self.connections.get(&name) else { |
| 3450 | continue; |
| 3451 | }; |
| 3452 | if connection.is_transport_ready() { |
| 3453 | if self.supervised_dead.remove(&name) { |
| 3454 | self.supervised_parked.remove(&name); |
| 3455 | recovered.push(name); |
| 3456 | } |
| 3457 | continue; |
| 3458 | } |
| 3459 | // A login-pending server cannot be fixed by redialing; the auth |
| 3460 | // surface owns it. It stays out of the dead set so recovery via |
| 3461 | // login reports nothing stale. |
| 3462 | if self.needs_auth_servers.contains(&name) { |
| 3463 | continue; |
| 3464 | } |
| 3465 | let fresh_death = self.supervised_dead.insert(name.clone()); |
| 3466 | if self.connecting.contains(&name) { |
| 3467 | continue; |
| 3468 | } |
| 3469 | if let Some(backoff) = self.connect_backoff.get(&name) { |
| 3470 | if backoff.consecutive_failures >= Self::SUPERVISOR_PARK_AFTER_CONSECUTIVE_FAILURES |
| 3471 | { |
| 3472 | if self.supervised_parked.insert(name.clone()) { |
| 3473 | parked.push(name); |
| 3474 | } |
| 3475 | continue; |
| 3476 | } |
| 3477 | if now < backoff.retry_after { |
| 3478 | continue; |
| 3479 | } |
| 3480 | } |
| 3481 | due.push(SupervisionDue { |
| 3482 | name, |
| 3483 | config, |
| 3484 | fresh_death, |
| 3485 | }); |
| 3486 | } |
| 3487 | SupervisionPlan { |
| 3488 | due, |
| 3489 | recovered, |
| 3490 | parked, |
| 3491 | timeouts: self.config.timeouts, |
| 3492 | network_policy: self.network_policy.clone(), |
| 3493 | catalog_generation: self.catalog_generation.load(Ordering::SeqCst), |
| 3494 | } |
| 3495 | } |
| 3496 | |
| 3497 | /// Resolve one supervised reconnect attempt. Success stores the live |
| 3498 | /// connection (which clears the backoff, the dead mark, and any park); |
| 3499 | /// failure records the backoff and reports the death or the repeated |
| 3500 | /// failure with the diagnosis, parking on the threshold crossing. |
| 3501 | pub(crate) fn resolve_supervision_attempt( |
| 3502 | &mut self, |
| 3503 | name: &str, |
| 3504 | fresh_death: bool, |
| 3505 | result: Result<McpConnection, anyhow::Error>, |
| 3506 | ) -> McpSupervisorUpdate { |
| 3507 | let mut update = McpSupervisorUpdate::default(); |
| 3508 | let stored = |
| 3509 | result.and_then(|connection| self.store_ready_connection(name.to_string(), connection)); |
| 3510 | match stored { |
| 3511 | Ok(()) => { |
| 3512 | if !fresh_death { |
| 3513 | update.recovered.push(name.to_string()); |
| 3514 | } |
| 3515 | } |
| 3516 | Err(error) => { |
| 3517 | self.note_connect_failure(name, &error); |
| 3518 | let last_error = self |
| 3519 | .connect_backoff |
| 3520 | .get(name) |
| 3521 | .map(|backoff| backoff.last_error.clone()) |
| 3522 | .unwrap_or_else(|| format!("{error:#}")); |
| 3523 | if fresh_death { |
| 3524 | update.died.push((name.to_string(), last_error)); |
| 3525 | } else { |
| 3526 | update.failed.push((name.to_string(), last_error)); |
| 3527 | } |
| 3528 | if self.connect_backoff.get(name).is_some_and(|backoff| { |
| 3529 | backoff.consecutive_failures >= Self::SUPERVISOR_PARK_AFTER_CONSECUTIVE_FAILURES |
| 3530 | }) && self.supervised_parked.insert(name.to_string()) |
| 3531 | { |
| 3532 | update.parked.push(name.to_string()); |
| 3533 | } |
| 3534 | } |
| 3535 | } |
| 3536 | update |
| 3537 | } |
| 3538 | |
| 3539 | /// Watch every live connection and reconnect the dead ones. Exits when |
| 3540 | /// the pool is dropped (the engine holds the only strong reference) or |
| 3541 | /// the engine stops listening. Reports transitions only, so the engine |
| 3542 | /// emits a snapshot update exactly when something changed (#6187). |
| 3543 | pub(crate) async fn supervise_pool( |
| 3544 | pool: std::sync::Weak<tokio::sync::Mutex<McpPool>>, |
| 3545 | tx: tokio::sync::mpsc::Sender<McpSupervisorUpdate>, |
| 3546 | ) { |
| 3547 | loop { |
| 3548 | tokio::time::sleep(Self::SUPERVISOR_TICK).await; |
| 3549 | let Some(pool) = pool.upgrade() else { break }; |
| 3550 | let plan = pool.lock().await.plan_supervision(); |
| 3551 | if plan.due.is_empty() && plan.recovered.is_empty() && plan.parked.is_empty() { |
| 3552 | continue; |
| 3553 | } |
| 3554 | let mut connects = Self::spawn_pending_connects( |
| 3555 | plan.due |
| 3556 | .iter() |
| 3557 | .map(|due| (due.name.clone(), due.config.clone())) |
| 3558 | .collect(), |
| 3559 | plan.timeouts, |
| 3560 | plan.network_policy.clone(), |
| 3561 | plan.catalog_generation, |
| 3562 | ); |
| 3563 | let mut update = McpSupervisorUpdate { |
| 3564 | recovered: plan.recovered, |
| 3565 | parked: plan.parked, |
| 3566 | ..Default::default() |
| 3567 | }; |
| 3568 | let fresh_by_name: HashMap<String, bool> = plan |
| 3569 | .due |
| 3570 | .into_iter() |
| 3571 | .map(|due| (due.name, due.fresh_death)) |
| 3572 | .collect(); |
| 3573 | while let Some(joined) = connects.join_next().await { |
| 3574 | let (name, result) = joined |
| 3575 | .unwrap_or_else(|error| ("connection task".to_string(), Err(error.into()))); |
| 3576 | let fresh_death = fresh_by_name.get(&name).copied().unwrap_or(false); |
| 3577 | let resolution = |
| 3578 | pool.lock() |
| 3579 | .await |
| 3580 | .resolve_supervision_attempt(&name, fresh_death, result); |
| 3581 | update.merge(resolution); |
| 3582 | } |
| 3583 | if !update.is_empty() && tx.send(update).await.is_err() { |
| 3584 | break; |
| 3585 | } |
| 3586 | } |
| 3587 | } |
| 3588 | |
| 3589 | /// Collect the configured servers a connect pass should start. `only` |
| 3590 | /// scopes the pass to the given names; `None` connects every enabled, |
| 3591 | /// allowed server (`connect_all`). Dynamic runtime servers stay |
| 3592 | /// registered and connect via [`Self::get_or_connect`]; connect passes |
| 3593 | /// have never spawned them. Every emitted name is marked |
| 3594 | /// [`Self::connecting`] until its spawn resolves. |
| 3595 | pub(crate) fn collect_pending_connects( |
| 3596 | &mut self, |
| 3597 | only: Option<&HashSet<String>>, |
| 3598 | ) -> (Vec<McpPendingConnect>, Vec<McpConnectError>) { |
| 3599 | let names: Vec<String> = self |
| 3600 | .config |
| 3601 | .servers |
| 3602 | .iter() |
| 3603 | .filter(|(name, server)| server.is_enabled() && self.server_allowed(name)) |
| 3604 | .filter(|(name, _)| only.is_none_or(|set| set.contains(*name))) |
| 3605 | .map(|(name, _)| name.clone()) |
| 3606 | .collect(); |
| 3607 | let mut pending = Vec::new(); |
| 3608 | let mut errors = Vec::new(); |
| 3609 | for name in names { |
| 3610 | let Some(server_config) = self.config.servers.get(&name).cloned() else { |
| 3611 | continue; |
| 3612 | }; |
| 3613 | |
| 3614 | let plugin_source = self |
| 3615 | .connections |
| 3616 | .get(&name) |
| 3617 | .and_then(|connection| connection.config().reviewed_plugin.clone()) |
| 3618 | .or_else(|| server_config.reviewed_plugin.clone()); |
| 3619 | if let Some(source) = plugin_source |
| 3620 | && let Err(error) = source.validate_before_use(&name, "use") |
| 3621 | { |
| 3622 | self.drop_connection(&name, "plugin authority revoked or changed"); |
| 3623 | errors.push((name, error)); |
| 3624 | continue; |
| 3625 | } |
| 3626 | |
| 3627 | // Authority validated immediately above for this same source. |
| 3628 | if self |
| 3629 | .connections |
| 3630 | .get(&name) |
| 3631 | .is_some_and(McpConnection::is_transport_ready) |
| 3632 | { |
| 3633 | continue; |
| 3634 | } |
| 3635 | // Inside its cooldown a failed server costs nothing and still |
| 3636 | // tells the truth: the recorded diagnosis is replayed so the row |
| 3637 | // keeps reading `error`, rather than going quiet and looking |
| 3638 | // healthy because nobody asked. |
| 3639 | if let Some(backoff) = self.connect_backoff.get(&name) |
| 3640 | && std::time::Instant::now() < backoff.retry_after |
| 3641 | { |
| 3642 | errors.push((name, anyhow::anyhow!(backoff.last_error.clone()))); |
| 3643 | continue; |
| 3644 | } |
| 3645 | if self.connecting.contains(&name) { |
| 3646 | // An earlier pass spawned this connect and it has not |
| 3647 | // resolved; a second pass must not spawn a duplicate. |
| 3648 | continue; |
| 3649 | } |
| 3650 | self.drop_connection(&name, "reconnect"); |
| 3651 | self.connecting.insert(name.clone()); |
| 3652 | pending.push((name, server_config)); |
| 3653 | } |
| 3654 | (pending, errors) |
| 3655 | } |
| 3656 | |
| 3657 | /// Start connects for servers an explicit tool selection named. Unlike a |
| 3658 | /// boot pass the selection is the intent — cooldowns do not apply — but |
| 3659 | /// servers already ready or already in flight are left alone, and plugin |
| 3660 | /// authority is re-validated exactly as in |
| 3661 | /// [`Self::collect_pending_connects`]. Covers dynamic servers too: a |
| 3662 | /// selection can name one. |
| 3663 | pub(crate) fn take_pending_connects_for( |
| 3664 | &mut self, |
| 3665 | names: &[String], |
| 3666 | ) -> (Vec<McpPendingConnect>, Vec<McpConnectError>) { |
| 3667 | let mut pending = Vec::new(); |
| 3668 | let mut errors = Vec::new(); |
| 3669 | for name in names { |
| 3670 | let Some(server_config) = self.server_config(name) else { |
| 3671 | continue; |
| 3672 | }; |
| 3673 | if !server_config.is_enabled() || !self.server_allowed(name) { |
| 3674 | continue; |
| 3675 | } |
| 3676 | let plugin_source = self |
| 3677 | .connections |
| 3678 | .get(name) |
| 3679 | .and_then(|connection| connection.config().reviewed_plugin.clone()) |
| 3680 | .or_else(|| server_config.reviewed_plugin.clone()); |
| 3681 | if let Some(source) = plugin_source |
| 3682 | && let Err(error) = source.validate_before_use(name, "use") |
| 3683 | { |
| 3684 | self.drop_connection(name, "plugin authority revoked or changed"); |
| 3685 | errors.push((name.clone(), error)); |
| 3686 | continue; |
| 3687 | } |
| 3688 | // Authority validated immediately above for this same source. |
| 3689 | if self |
| 3690 | .connections |
| 3691 | .get(name) |
| 3692 | .is_some_and(McpConnection::is_transport_ready) |
| 3693 | || !self.connecting.insert(name.clone()) |
| 3694 | { |
| 3695 | continue; |
| 3696 | } |
| 3697 | self.drop_connection(name, "reconnect"); |
| 3698 | pending.push((name.clone(), server_config)); |
| 3699 | } |
| 3700 | (pending, errors) |
| 3701 | } |
| 3702 | |
| 3703 | /// Forget in-flight marks for connects whose spawns were aborted before |
| 3704 | /// resolution (boot-pass abort on config change, deadline expiry). |
| 3705 | pub(crate) fn cancel_connecting(&mut self, names: &HashSet<String>) { |
| 3706 | self.connecting.retain(|name| !names.contains(name)); |
| 3707 | } |
| 3708 | |
| 3709 | /// Servers with a connect in flight right now — the one honest answer to |
| 3710 | /// "which servers are connecting" (#6033). |
| 3711 | pub(crate) fn connecting_servers(&self) -> Vec<String> { |
| 3712 | self.connecting.iter().cloned().collect() |
| 3713 | } |
| 3714 | |
| 3715 | /// Enabled, allowed configured servers the boot pass must still start |
| 3716 | /// eagerly under lazy boot (#6033): servers marked `required`, plus any |
| 3717 | /// server the session's explicit tool selections cover. |
| 3718 | pub(crate) fn eager_boot_server_names(&self, requested: &[String]) -> HashSet<String> { |
| 3719 | self.config |
| 3720 | .servers |
| 3721 | .iter() |
| 3722 | .filter(|(name, server)| server.is_enabled() && self.server_allowed(name)) |
| 3723 | .filter(|(name, server)| { |
| 3724 | server.required || tool_selection_covers_server(requested, name) |
| 3725 | }) |
| 3726 | .map(|(name, _)| name.clone()) |
| 3727 | .collect() |
| 3728 | } |
| 3729 | |
| 3730 | /// Enabled, allowed servers — configured or dynamic — covered by an |
| 3731 | /// explicit tool selection (`mcp_<server>_*` names or `mcp_<prefix>*` |
| 3732 | /// globs). These are the names a turn is allowed to start on demand. |
| 3733 | pub(crate) fn explicitly_selected_server_names(&self, requested: &[String]) -> Vec<String> { |
| 3734 | let dynamic = self.dynamic_servers.read(); |
| 3735 | self.config |
| 3736 | .servers |
| 3737 | .iter() |
| 3738 | .chain(dynamic.iter()) |
| 3739 | .filter(|(name, server)| server.is_enabled() && self.server_allowed(name)) |
| 3740 | .filter(|(name, _)| tool_selection_covers_server(requested, name)) |
| 3741 | .map(|(name, _)| name.clone()) |
| 3742 | .collect() |
| 3743 | } |
| 3744 | |
| 3745 | pub(crate) fn push_required_server_errors(&self, errors: &mut Vec<McpConnectError>) { |
| 3746 | for (name, server_cfg) in &self.config.servers { |
| 3747 | // Only stand in for a missing diagnosis. When the connect attempt |
| 3748 | // above already reported why this server failed, appending a |
| 3749 | // second, contentless entry for the same name buries it: callers |
| 3750 | // fold these pairs into a `HashMap<name, message>`, so the later |
| 3751 | // generic string silently replaced the real cause. |
| 3752 | if self.server_allowed(name) |
| 3753 | && server_cfg.required |
| 3754 | && server_cfg.is_enabled() |
| 3755 | && !self |
| 3756 | .connections |
| 3757 | .get(name) |
| 3758 | .is_some_and(McpConnection::is_ready) |
| 3759 | && !errors.iter().any(|(failed, _)| failed == name) |
| 3760 | { |
| 3761 | errors.push(( |
| 3762 | name.clone(), |
| 3763 | anyhow::anyhow!("required MCP server failed to initialize"), |
| 3764 | )); |
| 3765 | } |
| 3766 | } |
| 3767 | } |
| 3768 | |
| 3769 | /// Handshake the pending servers concurrently without holding the pool |
| 3770 | /// lock. Callers insert results under a short lock so a live turn can |
| 3771 | /// snapshot ready tools while optional servers are still connecting. |
| 3772 | pub(crate) fn spawn_pending_connects( |
| 3773 | pending: Vec<McpPendingConnect>, |
| 3774 | timeouts: McpTimeouts, |
| 3775 | network_policy: Option<NetworkPolicyDecider>, |
| 3776 | catalog_generation: u64, |
| 3777 | ) -> tokio::task::JoinSet<(String, Result<McpConnection, anyhow::Error>)> { |
| 3778 | let semaphore = std::sync::Arc::new(tokio::sync::Semaphore::new(Self::CONNECT_CONCURRENCY)); |
| 3779 | let mut joins: tokio::task::JoinSet<(String, Result<McpConnection, anyhow::Error>)> = |
| 3780 | tokio::task::JoinSet::new(); |
| 3781 | for (name, config) in pending { |
| 3782 | let permit = semaphore.clone(); |
| 3783 | let network_policy = network_policy.clone(); |
| 3784 | joins.spawn(async move { |
| 3785 | let connection = std::panic::AssertUnwindSafe(async { |
| 3786 | let _permit = permit.acquire_owned().await; |
| 3787 | McpConnection::connect_with_policy( |
| 3788 | name.clone(), |
| 3789 | config, |
| 3790 | &timeouts, |
| 3791 | network_policy.as_ref(), |
| 3792 | ) |
| 3793 | .await |
| 3794 | .map(|mut connection| { |
| 3795 | connection.catalog_generation = catalog_generation; |
| 3796 | connection |
| 3797 | }) |
| 3798 | }) |
| 3799 | .catch_unwind() |
| 3800 | .await |
| 3801 | .unwrap_or_else(|_| Err(anyhow::anyhow!("MCP connection task panicked"))); |
| 3802 | (name, connection) |
| 3803 | }); |
| 3804 | } |
| 3805 | |
| 3806 | joins |
| 3807 | } |
| 3808 | |
| 3809 | /// Connect to all enabled servers, returning errors for failed connections. |
| 3810 | /// |
| 3811 | /// Servers connect **concurrently** (bounded by [`Self::CONNECT_CONCURRENCY`]). |
| 3812 | /// This used to be a sequential loop over `get_or_connect`, so every |
| 3813 | /// server paid the slowest server's spawn+handshake from its own budget: |
| 3814 | /// with the default 10s connect timeout, N servers meant a worst case of |
| 3815 | /// N×10s before the pool was usable. Each connection still gets its own |
| 3816 | /// configured connect timeout; one wedged server can no longer serialize |
| 3817 | /// the rest. |
| 3818 | /// |
| 3819 | /// Semantics preserved from the sequential loop: only configured servers |
| 3820 | /// are connected (dynamic runtime entries stay registered), the config is |
| 3821 | /// reloaded before the name snapshot (so a server added mid-session |
| 3822 | /// connects on this call, not the next), plugin-authority revocation |
| 3823 | /// drops the connection instead of silently reconnecting, and the |
| 3824 | /// required-server sweep reports at most one error per name. Config edits |
| 3825 | /// that land while the batch is in flight are reconciled by one retry |
| 3826 | /// pass: a content change drops every connection the previous pass |
| 3827 | /// inserted. |
| 3828 | pub async fn connect_all(&mut self) -> Vec<(String, anyhow::Error)> { |
| 3829 | let mut errors = Vec::new(); |
| 3830 | // Reload before taking the configured-name snapshot. Previously the |
| 3831 | // first call after adding a server captured the old names, then only |
| 3832 | // noticed the config change inside `get_or_connect`, delaying the new |
| 3833 | // server until a second turn. |
| 3834 | if let Err(err) = self.reload_if_config_changed().await { |
| 3835 | errors.push(("configuration".to_string(), err)); |
| 3836 | return errors; |
| 3837 | } |
| 3838 | |
| 3839 | // Drop needs-auth markers for servers that are no longer connectable |
| 3840 | // (removed from config or disabled); the loop below re-marks any |
| 3841 | // server whose fresh connect attempt still fails auth-required. |
| 3842 | { |
| 3843 | let dynamic = self.dynamic_servers.read(); |
| 3844 | let before = self.needs_auth_servers.len(); |
| 3845 | self.needs_auth_servers.retain(|name| { |
| 3846 | self.config |
| 3847 | .servers |
| 3848 | .get(name) |
| 3849 | .is_some_and(|server| server.is_enabled()) |
| 3850 | || dynamic.get(name).is_some_and(|server| server.is_enabled()) |
| 3851 | }); |
| 3852 | if self.needs_auth_servers.len() != before { |
| 3853 | self.needs_auth_generation = self.needs_auth_generation.wrapping_add(1); |
| 3854 | } |
| 3855 | } |
| 3856 | |
| 3857 | for _pass in 0..2 { |
| 3858 | let (pending, auth_errors) = self.collect_pending_connects(None); |
| 3859 | errors.extend(auth_errors); |
| 3860 | if pending.is_empty() { |
| 3861 | break; |
| 3862 | } |
| 3863 | |
| 3864 | let mut connects = Self::spawn_pending_connects( |
| 3865 | pending, |
| 3866 | self.config.timeouts, |
| 3867 | self.network_policy.clone(), |
| 3868 | self.catalog_generation.load(Ordering::SeqCst), |
| 3869 | ); |
| 3870 | while let Some(joined) = connects.join_next().await { |
| 3871 | let (name, result) = joined |
| 3872 | .unwrap_or_else(|error| ("connection task".to_string(), Err(error.into()))); |
| 3873 | let result = result |
| 3874 | .and_then(|connection| self.store_ready_connection(name.clone(), connection)); |
| 3875 | if let Err(error) = result { |
| 3876 | self.note_connect_failure(&name, &error); |
| 3877 | errors.push((name, error)); |
| 3878 | } |
| 3879 | } |
| 3880 | |
| 3881 | // Reconcile a config edit that landed mid-batch: a content |
| 3882 | // change dropped every connection this pass inserted, so run one |
| 3883 | // more pass against the new config and drop the stale pass's |
| 3884 | // errors with it. |
| 3885 | match self.reload_if_config_changed().await { |
| 3886 | Ok(true) => { |
| 3887 | errors.clear(); |
| 3888 | continue; |
| 3889 | } |
| 3890 | Ok(false) => break, |
| 3891 | Err(error) => { |
| 3892 | errors.push(("configuration".to_string(), error)); |
| 3893 | break; |
| 3894 | } |
| 3895 | } |
| 3896 | } |
| 3897 | |
| 3898 | self.push_required_server_errors(&mut errors); |
| 3899 | errors |
| 3900 | } |
| 3901 | |
| 3902 | /// The single definition of an MCP tool's model-facing name. |
| 3903 | /// |
| 3904 | /// [`Self::all_tools`] (which builds the model catalog) and |
| 3905 | /// [`Self::resolved_tool_servers`] (which tells tool inspection which |
| 3906 | /// server owns a name) both call this, so a human-facing server |
| 3907 | /// attribution can never drift from the name the model actually received. |
| 3908 | #[must_use] |
| 3909 | pub fn mcp_model_tool_name(server: &str, tool: &str) -> String { |
| 3910 | format!("mcp_{server}_{tool}") |
| 3911 | } |
| 3912 | |
| 3913 | /// Map an exact `mcp_<server>_authenticate` model tool name to its server |
| 3914 | /// when the synthetic self-serve OAuth login tool should answer for it: |
| 3915 | /// the server's last connect failed auth-required, or the name matches a |
| 3916 | /// configured OAuth-capable server whose catalog entry is stale (a login |
| 3917 | /// completed since the catalog was built). Never fires when a ready |
| 3918 | /// connection advertises a real tool under the same model name — the |
| 3919 | /// server's own `authenticate` tool always wins. |
| 3920 | pub(crate) fn authenticate_tool_target(&self, prefixed_name: &str) -> Option<String> { |
| 3921 | if !self.tool_allowed(prefixed_name) { |
| 3922 | return None; |
| 3923 | } |
| 3924 | let target = self |
| 3925 | .needs_auth_servers |
| 3926 | .iter() |
| 3927 | .find(|server| { |
| 3928 | Self::mcp_model_tool_name(server, AUTHENTICATE_TOOL_NAME) == prefixed_name |
| 3929 | }) |
| 3930 | .cloned() |
| 3931 | .or_else(|| { |
| 3932 | let dynamic = self.dynamic_servers.read(); |
| 3933 | self.config |
| 3934 | .servers |
| 3935 | .iter() |
| 3936 | .chain(dynamic.iter()) |
| 3937 | .find(|(name, _)| { |
| 3938 | Self::mcp_model_tool_name(name, AUTHENTICATE_TOOL_NAME) == prefixed_name |
| 3939 | }) |
| 3940 | .map(|(name, _)| name.clone()) |
| 3941 | })?; |
| 3942 | let capable = { |
| 3943 | let dynamic = self.dynamic_servers.read(); |
| 3944 | self.config |
| 3945 | .servers |
| 3946 | .get(&target) |
| 3947 | .or_else(|| dynamic.get(&target)) |
| 3948 | .is_some_and(oauth::server_supports_oauth_login) |
| 3949 | }; |
| 3950 | if !capable || !self.server_allowed(&target) { |
| 3951 | return None; |
| 3952 | } |
| 3953 | if self.parse_prefixed_name(prefixed_name).is_ok() { |
| 3954 | return None; |
| 3955 | } |
| 3956 | Some(target) |
| 3957 | } |
| 3958 | |
| 3959 | /// The configured server by name, static config first, then the |
| 3960 | /// session's dynamically added servers. |
| 3961 | fn server_config(&self, server_name: &str) -> Option<McpServerConfig> { |
| 3962 | self.config |
| 3963 | .servers |
| 3964 | .get(server_name) |
| 3965 | .cloned() |
| 3966 | .or_else(|| self.dynamic_servers.read().get(server_name).cloned()) |
| 3967 | } |
| 3968 | |
| 3969 | /// Phase one of the synthetic `mcp_<server>_authenticate` tool: decide |
| 3970 | /// whether a browser login is needed and, if so, start it. Only touches |
| 3971 | /// config, the connection map, and the token store — it returns as soon |
| 3972 | /// as the authorization URL exists, so a caller holding the pool lock can |
| 3973 | /// release it before the (up to five minute) browser wait in |
| 3974 | /// [`oauth::McpOAuthToolLogin::finish`]. Holding the lock across that wait |
| 3975 | /// would freeze every other MCP call, the `/mcp` manager, and the |
| 3976 | /// Extensions view for the whole sign-in. |
| 3977 | pub(crate) async fn begin_authenticate_tool( |
| 3978 | &self, |
| 3979 | server_name: &str, |
| 3980 | ) -> Result<AuthenticateToolStart> { |
| 3981 | self.require_server(server_name)?; |
| 3982 | Self::authorize_call( |
| 3983 | &self.disallowed_tools, |
| 3984 | &Self::mcp_model_tool_name(server_name, AUTHENTICATE_TOOL_NAME), |
| 3985 | &serde_json::json!({}), |
| 3986 | )?; |
| 3987 | let server = self |
| 3988 | .server_config(server_name) |
| 3989 | .ok_or_else(|| anyhow::anyhow!("MCP server '{server_name}' is no longer configured"))?; |
| 3990 | if !server.is_enabled() { |
| 3991 | anyhow::bail!("MCP server '{server_name}' is disabled"); |
| 3992 | } |
| 3993 | |
| 3994 | // Already-authorized branch: a login that completed since the catalog |
| 3995 | // was built (e.g. `codewhale mcp login` in another window) must not |
| 3996 | // restart the browser flow — adopt the stored tokens by reconnecting. |
| 3997 | let ready = self |
| 3998 | .connections |
| 3999 | .get(server_name) |
| 4000 | .is_some_and(McpConnection::is_ready); |
| 4001 | if ready || oauth::has_usable_stored_tokens(server_name, &server) { |
| 4002 | return Ok(AuthenticateToolStart::AlreadyAuthorized); |
| 4003 | } |
| 4004 | let login = oauth::begin_oauth_login_for_server_tool( |
| 4005 | server_name, |
| 4006 | &server, |
| 4007 | None, |
| 4008 | self.oauth_callback_port, |
| 4009 | self.oauth_callback_url.as_deref(), |
| 4010 | self.network_policy.as_ref(), |
| 4011 | ) |
| 4012 | .await?; |
| 4013 | Ok(AuthenticateToolStart::Login(Box::new(login))) |
| 4014 | } |
| 4015 | |
| 4016 | /// Phase two of the synthetic authenticate tool: reconnect the server so |
| 4017 | /// its real tools resolve in this session, and describe the outcome for |
| 4018 | /// the model. `get_or_connect` drops any non-ready connection itself, so |
| 4019 | /// both the fresh-login and already-authorized outcomes fall through to |
| 4020 | /// it. Errors are never swallowed — a reconnect that still fails |
| 4021 | /// auth-required re-marks the server (via `get_or_connect`) and surfaces |
| 4022 | /// truthfully for the model to relay. |
| 4023 | pub(crate) async fn finish_authenticate_tool( |
| 4024 | &mut self, |
| 4025 | server_name: &str, |
| 4026 | outcome: AuthenticateToolOutcome, |
| 4027 | ) -> Result<serde_json::Value> { |
| 4028 | self.require_server(server_name)?; |
| 4029 | let rules = self.disallowed_tools.clone(); |
| 4030 | match self.get_or_connect(server_name).await { |
| 4031 | Ok(conn) => { |
| 4032 | let tools: Vec<String> = conn |
| 4033 | .tools() |
| 4034 | .iter() |
| 4035 | .filter(|tool| conn.config().is_tool_enabled(&tool.name)) |
| 4036 | .map(|tool| Self::mcp_model_tool_name(server_name, &tool.name)) |
| 4037 | .filter(|name| { |
| 4038 | !crate::core::engine::tool_catalog::tool_matches_any_rule(&rules, name) |
| 4039 | }) |
| 4040 | .collect(); |
| 4041 | let (status, detail) = match &outcome { |
| 4042 | AuthenticateToolOutcome::Authenticated { .. } => ( |
| 4043 | "authenticated", |
| 4044 | "authenticated successfully and is now connected", |
| 4045 | ), |
| 4046 | AuthenticateToolOutcome::AlreadyAuthorized => ( |
| 4047 | "already_authorized", |
| 4048 | "already had valid OAuth credentials and is now connected", |
| 4049 | ), |
| 4050 | }; |
| 4051 | let mut result = serde_json::json!({ |
| 4052 | "status": status, |
| 4053 | "server": server_name, |
| 4054 | "tools": tools, |
| 4055 | "message": format!( |
| 4056 | "MCP server '{server_name}' {detail}. Its real MCP tools (listed in 'tools') replaced the synthetic authenticate tool and are callable from the next model request in this session." |
| 4057 | ), |
| 4058 | }); |
| 4059 | if let AuthenticateToolOutcome::Authenticated { authorization_url } = outcome { |
| 4060 | result["authorization_url"] = serde_json::Value::String(authorization_url); |
| 4061 | } |
| 4062 | Ok(result) |
| 4063 | } |
| 4064 | Err(error) => { |
| 4065 | // A stored token the server just rejected must not feed the |
| 4066 | // loop again: `begin_authenticate_tool` short-circuits to |
| 4067 | // AlreadyAuthorized whenever usable-looking tokens exist, so |
| 4068 | // keeping them means the model re-calls the synthetic tool |
| 4069 | // forever while every reconnect fails auth-required. Drop |
| 4070 | // the durable copy so the next begin starts a real login. |
| 4071 | if matches!(outcome, AuthenticateToolOutcome::AlreadyAuthorized) |
| 4072 | && oauth::error_looks_auth_required(&error) |
| 4073 | && let Some(server) = self.server_config(server_name) |
| 4074 | { |
| 4075 | match oauth::delete_oauth_tokens_for_server(server_name, &server) { |
| 4076 | Ok(true) => tracing::info!( |
| 4077 | target: "mcp", |
| 4078 | server = %server_name, |
| 4079 | "rejected stored OAuth token removed after failed AlreadyAuthorized reconnect" |
| 4080 | ), |
| 4081 | Ok(false) => {} |
| 4082 | Err(delete_err) => tracing::warn!( |
| 4083 | target: "mcp", |
| 4084 | server = %server_name, |
| 4085 | error = %delete_err, |
| 4086 | "could not remove rejected stored OAuth token" |
| 4087 | ), |
| 4088 | } |
| 4089 | } |
| 4090 | Err(error).with_context(|| { |
| 4091 | format!( |
| 4092 | "MCP server '{server_name}' completed OAuth login but the reconnect failed" |
| 4093 | ) |
| 4094 | }) |
| 4095 | } |
| 4096 | } |
| 4097 | } |
| 4098 | |
| 4099 | /// Execute the synthetic `mcp_<server>_authenticate` tool in place, |
| 4100 | /// holding `&mut self` (and therefore any enclosing pool lock) for the |
| 4101 | /// whole flow. Engine tool execution uses |
| 4102 | /// [`authenticate_tool_via_pool`] instead, which releases the shared |
| 4103 | /// pool between the two phases. |
| 4104 | async fn run_authenticate_tool(&mut self, server_name: &str) -> Result<serde_json::Value> { |
| 4105 | let outcome = match self.begin_authenticate_tool(server_name).await? { |
| 4106 | AuthenticateToolStart::AlreadyAuthorized => AuthenticateToolOutcome::AlreadyAuthorized, |
| 4107 | AuthenticateToolStart::Login(login) => { |
| 4108 | let authorization_url = login.authorization_url().to_string(); |
| 4109 | login.finish().await?; |
| 4110 | AuthenticateToolOutcome::Authenticated { authorization_url } |
| 4111 | } |
| 4112 | }; |
| 4113 | self.finish_authenticate_tool(server_name, outcome).await |
| 4114 | } |
| 4115 | |
| 4116 | /// The model-facing recovery for a server in the `◆ auth required` |
| 4117 | /// state: the one call that recovers it. Names the synthetic |
| 4118 | /// `mcp_<server>_authenticate` tool when the server is OAuth-servable, and |
| 4119 | /// otherwise the credential source (plugin environment header, manual |
| 4120 | /// bearer) the server is allowed to authenticate with. |
| 4121 | fn auth_required_hint(&self, server_name: &str) -> String { |
| 4122 | let tool_name = Self::mcp_model_tool_name(server_name, AUTHENTICATE_TOOL_NAME); |
| 4123 | if self.authenticate_tool_target(&tool_name).is_some() { |
| 4124 | return format!( |
| 4125 | "MCP server '{server_name}' requires OAuth login (◆ auth required); call the `{tool_name}` tool to authenticate, or run `/mcp login {server_name}`" |
| 4126 | ); |
| 4127 | } |
| 4128 | let recovery = match self.server_config(server_name) { |
| 4129 | Some(server) => oauth::auth_required_recovery_hint(server_name, &server), |
| 4130 | None => oauth::auth_required_login_hint(server_name), |
| 4131 | }; |
| 4132 | format!("MCP server '{server_name}' requires authentication (◆ auth required); {recovery}") |
| 4133 | } |
| 4134 | |
| 4135 | /// Route an auth-required failure from a live tool call into the same |
| 4136 | /// typed state a failed connect produces: drop the connection (its |
| 4137 | /// credential is no longer accepted), mark the server needs-auth so the |
| 4138 | /// next catalog offers the synthetic login tool, and name the recovery on |
| 4139 | /// the error. Any other error passes through untouched. |
| 4140 | fn note_live_call_failure(&mut self, server_name: &str, error: anyhow::Error) -> anyhow::Error { |
| 4141 | if !oauth::error_looks_auth_required(&error) { |
| 4142 | return error; |
| 4143 | } |
| 4144 | self.drop_connection(server_name, "auth required on live call"); |
| 4145 | self.note_connect_failure(server_name, &error); |
| 4146 | error.context(self.auth_required_hint(server_name)) |
| 4147 | } |
| 4148 | |
| 4149 | /// Fold `(server, tool)` pairs into `model name -> owning server`. |
| 4150 | /// |
| 4151 | /// Mirrors [`Self::all_tools`]' ambiguity rule: when two servers produce |
| 4152 | /// the same model name, the name is dropped entirely rather than |
| 4153 | /// attributed to an arbitrary winner. Callers then report it as unknown. |
| 4154 | #[must_use] |
| 4155 | pub fn resolve_tool_server_map<'a>( |
| 4156 | pairs: impl Iterator<Item = (&'a str, &'a str)>, |
| 4157 | ) -> std::collections::BTreeMap<String, String> { |
| 4158 | let mut resolved: std::collections::BTreeMap<String, Option<String>> = |
| 4159 | std::collections::BTreeMap::new(); |
| 4160 | for (server, tool) in pairs { |
| 4161 | match resolved.entry(Self::mcp_model_tool_name(server, tool)) { |
| 4162 | std::collections::btree_map::Entry::Vacant(entry) => { |
| 4163 | entry.insert(Some(server.to_string())); |
| 4164 | } |
| 4165 | std::collections::btree_map::Entry::Occupied(mut entry) => { |
| 4166 | entry.insert(None); |
| 4167 | } |
| 4168 | } |
| 4169 | } |
| 4170 | resolved |
| 4171 | .into_iter() |
| 4172 | .filter_map(|(name, server)| server.map(|server| (name, server))) |
| 4173 | .collect() |
| 4174 | } |
| 4175 | |
| 4176 | /// Model tool name -> owning server name, for the tools this pool actually |
| 4177 | /// resolved. Names the pool did not resolve are simply absent, so callers |
| 4178 | /// report them as unknown instead of parsing `mcp_{server}_{tool}` (a |
| 4179 | /// server name may itself contain `_`, so that split is a guess). |
| 4180 | /// |
| 4181 | /// Read-only projection used by tool inspection; it never connects or |
| 4182 | /// executes. |
| 4183 | #[must_use] |
| 4184 | pub fn resolved_tool_servers(&self) -> std::collections::BTreeMap<String, String> { |
| 4185 | Self::resolve_tool_server_map(self.connections.iter().flat_map(|(server, conn)| { |
| 4186 | let authorized = self.server_allowed(server) && conn.catalog_authorized(); |
| 4187 | conn.tools().iter().filter_map(move |tool| { |
| 4188 | (authorized |
| 4189 | && conn.config().is_tool_enabled(&tool.name) |
| 4190 | && self.tool_allowed(&Self::mcp_model_tool_name(server, &tool.name))) |
| 4191 | .then_some((server.as_str(), tool.name.as_str())) |
| 4192 | }) |
| 4193 | })) |
| 4194 | } |
| 4195 | |
| 4196 | /// Get all discovered tools with server-prefixed names |
| 4197 | pub fn all_tools(&self) -> Vec<(String, &McpTool)> { |
| 4198 | let mut by_name: std::collections::BTreeMap<String, Option<&McpTool>> = |
| 4199 | std::collections::BTreeMap::new(); |
| 4200 | for (server, conn) in &self.connections { |
| 4201 | if !self.server_allowed(server) || !conn.catalog_authorized() { |
| 4202 | continue; |
| 4203 | } |
| 4204 | for tool in conn.tools() { |
| 4205 | if !conn.config().is_tool_enabled(&tool.name) { |
| 4206 | continue; |
| 4207 | } |
| 4208 | let name = Self::mcp_model_tool_name(server, &tool.name); |
| 4209 | if !self.tool_allowed(&name) { |
| 4210 | continue; |
| 4211 | } |
| 4212 | match by_name.entry(name.clone()) { |
| 4213 | std::collections::btree_map::Entry::Vacant(entry) => { |
| 4214 | entry.insert(Some(tool)); |
| 4215 | } |
| 4216 | std::collections::btree_map::Entry::Occupied(mut entry) => { |
| 4217 | tracing::warn!( |
| 4218 | target: "mcp", |
| 4219 | model_tool = %name, |
| 4220 | "hiding ambiguous MCP model tool name" |
| 4221 | ); |
| 4222 | entry.insert(None); |
| 4223 | } |
| 4224 | } |
| 4225 | } |
| 4226 | } |
| 4227 | by_name |
| 4228 | .into_iter() |
| 4229 | .filter_map(|(name, tool)| tool.map(|tool| (name, tool))) |
| 4230 | .collect() |
| 4231 | } |
| 4232 | |
| 4233 | /// Get all discovered resources with server-prefixed names |
| 4234 | pub fn all_resources(&self) -> Vec<(String, &McpResource)> { |
| 4235 | let mut resources = Vec::new(); |
| 4236 | for (server, conn) in &self.connections { |
| 4237 | if !self.server_allowed(server) || !conn.catalog_authorized() { |
| 4238 | continue; |
| 4239 | } |
| 4240 | for resource in conn.resources() { |
| 4241 | // Format: mcp_{server}_{resource_name} |
| 4242 | // Note: resource names might contain spaces, we should probably slugify them |
| 4243 | let safe_name = resource.name.replace(' ', "_").to_lowercase(); |
| 4244 | resources.push((format!("mcp_{server}_{safe_name}"), resource)); |
| 4245 | } |
| 4246 | } |
| 4247 | resources |
| 4248 | } |
| 4249 | |
| 4250 | /// Get all discovered resource templates with server-prefixed names |
| 4251 | #[allow(dead_code)] // Public API for MCP resource discovery |
| 4252 | pub fn all_resource_templates(&self) -> Vec<(String, &McpResourceTemplate)> { |
| 4253 | let mut templates = Vec::new(); |
| 4254 | for (server, conn) in &self.connections { |
| 4255 | if !self.server_allowed(server) || !conn.catalog_authorized() { |
| 4256 | continue; |
| 4257 | } |
| 4258 | for template in conn.resource_templates() { |
| 4259 | let safe_name = template.name.replace(' ', "_").to_lowercase(); |
| 4260 | templates.push((format!("mcp_{server}_{safe_name}"), template)); |
| 4261 | } |
| 4262 | } |
| 4263 | templates |
| 4264 | } |
| 4265 | |
| 4266 | async fn list_resources(&mut self, server: Option<String>) -> Result<Vec<serde_json::Value>> { |
| 4267 | if let Some(server_name) = server { |
| 4268 | let conn = self.get_or_connect(&server_name).await?; |
| 4269 | let resources = conn |
| 4270 | .resources() |
| 4271 | .iter() |
| 4272 | .map(|resource| { |
| 4273 | serde_json::json!({ |
| 4274 | "server": server_name.clone(), |
| 4275 | "uri": resource.uri, |
| 4276 | "name": resource.name, |
| 4277 | "description": resource.description, |
| 4278 | "mime_type": resource.mime_type, |
| 4279 | }) |
| 4280 | }) |
| 4281 | .collect(); |
| 4282 | return Ok(resources); |
| 4283 | } |
| 4284 | |
| 4285 | let mut items = Vec::new(); |
| 4286 | let errors = self.connect_all().await; |
| 4287 | for (server, err) in errors { |
| 4288 | tracing::warn!("Failed to connect MCP server '{server}' for resources: {err:#}"); |
| 4289 | if oauth::error_looks_auth_required(&err) { |
| 4290 | items.push(self.mcp_auth_required_error_item(&server)); |
| 4291 | } |
| 4292 | } |
| 4293 | for (server, conn) in &self.connections { |
| 4294 | if !self.server_allowed(server) || !conn.catalog_authorized() { |
| 4295 | continue; |
| 4296 | } |
| 4297 | for resource in conn.resources() { |
| 4298 | items.push(serde_json::json!({ |
| 4299 | "server": server, |
| 4300 | "uri": resource.uri, |
| 4301 | "name": resource.name, |
| 4302 | "description": resource.description, |
| 4303 | "mime_type": resource.mime_type, |
| 4304 | })); |
| 4305 | } |
| 4306 | } |
| 4307 | Ok(items) |
| 4308 | } |
| 4309 | |
| 4310 | async fn list_resource_templates( |
| 4311 | &mut self, |
| 4312 | server: Option<String>, |
| 4313 | ) -> Result<Vec<serde_json::Value>> { |
| 4314 | if let Some(server_name) = server { |
| 4315 | let conn = self.get_or_connect(&server_name).await?; |
| 4316 | let templates = conn |
| 4317 | .resource_templates() |
| 4318 | .iter() |
| 4319 | .map(|template| { |
| 4320 | serde_json::json!({ |
| 4321 | "server": server_name.clone(), |
| 4322 | "uri_template": template.uri_template, |
| 4323 | "name": template.name, |
| 4324 | "description": template.description, |
| 4325 | "mime_type": template.mime_type, |
| 4326 | }) |
| 4327 | }) |
| 4328 | .collect(); |
| 4329 | return Ok(templates); |
| 4330 | } |
| 4331 | |
| 4332 | let mut items = Vec::new(); |
| 4333 | let errors = self.connect_all().await; |
| 4334 | for (server, err) in errors { |
| 4335 | tracing::warn!( |
| 4336 | "Failed to connect MCP server '{server}' for resource templates: {err:#}" |
| 4337 | ); |
| 4338 | if oauth::error_looks_auth_required(&err) { |
| 4339 | items.push(self.mcp_auth_required_error_item(&server)); |
| 4340 | } |
| 4341 | } |
| 4342 | for (server, conn) in &self.connections { |
| 4343 | if !self.server_allowed(server) || !conn.catalog_authorized() { |
| 4344 | continue; |
| 4345 | } |
| 4346 | for template in conn.resource_templates() { |
| 4347 | items.push(serde_json::json!({ |
| 4348 | "server": server, |
| 4349 | "uri_template": template.uri_template, |
| 4350 | "name": template.name, |
| 4351 | "description": template.description, |
| 4352 | "mime_type": template.mime_type, |
| 4353 | })); |
| 4354 | } |
| 4355 | } |
| 4356 | Ok(items) |
| 4357 | } |
| 4358 | |
| 4359 | /// Listing-time error item for a needs-auth server. Carries the same |
| 4360 | /// recovery the tool-call path names (the synthetic authenticate tool |
| 4361 | /// when OAuth-servable) so a resource or prompt listing never dead-ends |
| 4362 | /// on a login the model could have self-served. |
| 4363 | fn mcp_auth_required_error_item(&self, server: &str) -> serde_json::Value { |
| 4364 | let mut item = serde_json::json!({ |
| 4365 | "error": "authentication_required", |
| 4366 | "server": server, |
| 4367 | "message": self.auth_required_hint(server), |
| 4368 | }); |
| 4369 | let tool_name = Self::mcp_model_tool_name(server, AUTHENTICATE_TOOL_NAME); |
| 4370 | if self.authenticate_tool_target(&tool_name).is_some() { |
| 4371 | item["authenticate_tool"] = serde_json::Value::String(tool_name); |
| 4372 | } |
| 4373 | item |
| 4374 | } |
| 4375 | |
| 4376 | /// Get all discovered prompts with server-prefixed names |
| 4377 | pub fn all_prompts(&self) -> Vec<(String, &McpPrompt)> { |
| 4378 | let mut prompts = Vec::new(); |
| 4379 | for (server, conn) in &self.connections { |
| 4380 | if !self.server_allowed(server) || !conn.catalog_authorized() { |
| 4381 | continue; |
| 4382 | } |
| 4383 | for prompt in conn.prompts() { |
| 4384 | // Format: mcp_{server}_{prompt} |
| 4385 | prompts.push((format!("mcp_{}_{}", server, prompt.name), prompt)); |
| 4386 | } |
| 4387 | } |
| 4388 | prompts |
| 4389 | } |
| 4390 | |
| 4391 | /// Read a resource from a specific server |
| 4392 | pub async fn read_resource( |
| 4393 | &mut self, |
| 4394 | server_name: &str, |
| 4395 | uri: &str, |
| 4396 | ) -> Result<serde_json::Value> { |
| 4397 | let global_timeouts = self.config.timeouts; |
| 4398 | let conn = self.get_or_connect(server_name).await?; |
| 4399 | let advertised_literal = conn.resources().iter().any(|resource| resource.uri == uri); |
| 4400 | let advertised_template = conn |
| 4401 | .resource_templates() |
| 4402 | .iter() |
| 4403 | .any(|template| resource_uri_matches_template(uri, &template.uri_template)); |
| 4404 | if !advertised_literal && !advertised_template { |
| 4405 | anyhow::bail!("MCP resource URI '{uri}' was not advertised by server '{server_name}'"); |
| 4406 | } |
| 4407 | let timeout = conn.config().effective_read_timeout(&global_timeouts); |
| 4408 | conn.read_resource(uri, timeout).await |
| 4409 | } |
| 4410 | |
| 4411 | /// Get a prompt from a specific server |
| 4412 | pub async fn get_prompt( |
| 4413 | &mut self, |
| 4414 | server_name: &str, |
| 4415 | prompt_name: &str, |
| 4416 | arguments: serde_json::Value, |
| 4417 | ) -> Result<serde_json::Value> { |
| 4418 | let global_timeouts = self.config.timeouts; |
| 4419 | let conn = self.get_or_connect(server_name).await?; |
| 4420 | if !conn |
| 4421 | .prompts() |
| 4422 | .iter() |
| 4423 | .any(|prompt| prompt.name == prompt_name) |
| 4424 | { |
| 4425 | anyhow::bail!( |
| 4426 | "MCP prompt '{prompt_name}' was not advertised by server '{server_name}'" |
| 4427 | ); |
| 4428 | } |
| 4429 | let timeout = conn.config().effective_execute_timeout(&global_timeouts); |
| 4430 | conn.get_prompt(prompt_name, arguments, timeout).await |
| 4431 | } |
| 4432 | |
| 4433 | /// Parse a prefixed name into (server_name, tool_name) |
| 4434 | pub(crate) fn parse_prefixed_name(&self, prefixed_name: &str) -> Result<(String, String)> { |
| 4435 | Self::authorize_call( |
| 4436 | &self.disallowed_tools, |
| 4437 | prefixed_name, |
| 4438 | &serde_json::json!({}), |
| 4439 | )?; |
| 4440 | let Some(rest) = prefixed_name.strip_prefix("mcp_") else { |
| 4441 | anyhow::bail!("Invalid MCP tool name: {prefixed_name}"); |
| 4442 | }; |
| 4443 | |
| 4444 | let mut matched: Option<(String, String)> = None; |
| 4445 | for (server, connection) in &self.connections { |
| 4446 | if !self.server_allowed(server) || !connection.catalog_authorized() { |
| 4447 | continue; |
| 4448 | } |
| 4449 | for tool in connection.tools() { |
| 4450 | if !connection.config().is_tool_enabled(&tool.name) |
| 4451 | || format!("{server}_{}", tool.name) != rest |
| 4452 | { |
| 4453 | continue; |
| 4454 | } |
| 4455 | if matched.is_some() { |
| 4456 | anyhow::bail!( |
| 4457 | "Ambiguous MCP tool name '{prefixed_name}' matches more than one server/tool authority" |
| 4458 | ); |
| 4459 | } |
| 4460 | matched = Some((server.clone(), tool.name.clone())); |
| 4461 | } |
| 4462 | } |
| 4463 | if let Some(matched) = matched { |
| 4464 | return Ok(matched); |
| 4465 | } |
| 4466 | |
| 4467 | Err(anyhow::anyhow!("Unknown MCP tool name: {prefixed_name}")) |
| 4468 | } |
| 4469 | |
| 4470 | /// Resolve an MCP tool through an exact advertised catalog. A configured |
| 4471 | /// but lazy server may be connected and asked for `tools/list`; the |
| 4472 | /// requested suffix is never treated as authority on its own. |
| 4473 | async fn resolve_advertised_tool(&mut self, prefixed_name: &str) -> Result<McpToolRoute> { |
| 4474 | Self::authorize_call( |
| 4475 | &self.disallowed_tools, |
| 4476 | prefixed_name, |
| 4477 | &serde_json::json!({}), |
| 4478 | )?; |
| 4479 | if let Ok((server_name, tool_name)) = self.parse_prefixed_name(prefixed_name) { |
| 4480 | return self.capture_tool_route(server_name, tool_name); |
| 4481 | } |
| 4482 | let Some(rest) = prefixed_name.strip_prefix("mcp_") else { |
| 4483 | anyhow::bail!("Invalid MCP tool name: {prefixed_name}"); |
| 4484 | }; |
| 4485 | let mut candidates = { |
| 4486 | let dynamic = self.dynamic_servers.read(); |
| 4487 | self.config |
| 4488 | .servers |
| 4489 | .iter() |
| 4490 | .filter_map(|(name, config)| { |
| 4491 | (config.is_enabled() |
| 4492 | && self.server_allowed(name) |
| 4493 | && rest |
| 4494 | .strip_prefix(name) |
| 4495 | .is_some_and(|suffix| suffix.starts_with('_'))) |
| 4496 | .then_some(name.clone()) |
| 4497 | }) |
| 4498 | .chain(dynamic.iter().filter_map(|(name, config)| { |
| 4499 | (config.is_enabled() |
| 4500 | && self.server_allowed(name) |
| 4501 | && rest |
| 4502 | .strip_prefix(name) |
| 4503 | .is_some_and(|suffix| suffix.starts_with('_'))) |
| 4504 | .then_some(name.clone()) |
| 4505 | })) |
| 4506 | .collect::<Vec<_>>() |
| 4507 | }; |
| 4508 | candidates.sort(); |
| 4509 | candidates.dedup(); |
| 4510 | for server in candidates { |
| 4511 | // Connecting and catalog discovery are the only lazy side effects. |
| 4512 | // A guessed method is never sent to the transport. |
| 4513 | let _ = self.get_or_connect(&server).await?; |
| 4514 | } |
| 4515 | let (server_name, tool_name) = self.parse_prefixed_name(prefixed_name)?; |
| 4516 | self.capture_tool_route(server_name, tool_name) |
| 4517 | } |
| 4518 | |
| 4519 | fn capture_tool_route(&self, server_name: String, tool_name: String) -> Result<McpToolRoute> { |
| 4520 | let connection = self |
| 4521 | .connections |
| 4522 | .get(&server_name) |
| 4523 | .context("advertised MCP connection disappeared during resolution")?; |
| 4524 | let plugin_authority = connection |
| 4525 | .config() |
| 4526 | .reviewed_plugin |
| 4527 | .as_ref() |
| 4528 | .map(|source| source.authority.clone()); |
| 4529 | Ok(McpToolRoute { |
| 4530 | server_name, |
| 4531 | tool_name, |
| 4532 | catalog_generation: connection.catalog_generation, |
| 4533 | plugin_authority, |
| 4534 | }) |
| 4535 | } |
| 4536 | |
| 4537 | /// Every model-facing tool name the runtime MCP pool can own right now: |
| 4538 | /// the current `to_api_tools` output plus the synthetic |
| 4539 | /// `mcp_<server>_authenticate` name for every enabled OAuth-servable |
| 4540 | /// server, whether or not it is currently needs-auth. The turn loop |
| 4541 | /// uses this universe to REPLACE the pool's slice of the tool catalog |
| 4542 | /// instead of additively merging it — the synthetic entry must leave |
| 4543 | /// after a login, and dead real tools must leave after a live 401. |
| 4544 | /// The model-visible tool-name universe for an already-built catalog. |
| 4545 | /// |
| 4546 | /// Takes the catalog rather than rebuilding it: `to_api_tools` re-verifies |
| 4547 | /// every reviewed plugin bundle, so calling both meant hashing each bundle |
| 4548 | /// twice per turn to produce two views of one thing — and the two could |
| 4549 | /// disagree if authority drifted between them (#6209). |
| 4550 | pub fn model_tool_names( |
| 4551 | &self, |
| 4552 | api_tools: &[codewhale_models::Tool], |
| 4553 | ) -> std::collections::HashSet<String> { |
| 4554 | let mut names: std::collections::HashSet<String> = |
| 4555 | api_tools.iter().map(|tool| tool.name.clone()).collect(); |
| 4556 | let dynamic = self.dynamic_servers.read(); |
| 4557 | for (server, config) in self.config.servers.iter().chain(dynamic.iter()) { |
| 4558 | if self.server_allowed(server) |
| 4559 | && config.is_enabled() |
| 4560 | && oauth::server_supports_oauth_login(config) |
| 4561 | && self.tool_allowed(&Self::mcp_model_tool_name(server, AUTHENTICATE_TOOL_NAME)) |
| 4562 | { |
| 4563 | names.insert(Self::mcp_model_tool_name(server, AUTHENTICATE_TOOL_NAME)); |
| 4564 | } |
| 4565 | } |
| 4566 | names |
| 4567 | } |
| 4568 | |
| 4569 | /// Convert discovered tools to API Tool format |
| 4570 | pub fn to_api_tools(&self) -> Vec<codewhale_models::Tool> { |
| 4571 | let mut api_tools = Vec::new(); |
| 4572 | // Add regular tools |
| 4573 | for (name, tool) in self.all_tools() { |
| 4574 | api_tools.push(codewhale_models::Tool { |
| 4575 | tool_type: None, |
| 4576 | name, |
| 4577 | description: tool.description.clone().unwrap_or_default(), |
| 4578 | input_schema: tool.input_schema.clone(), |
| 4579 | allowed_callers: Some(vec!["direct".to_string()]), |
| 4580 | defer_loading: Some(false), |
| 4581 | input_examples: None, |
| 4582 | strict: None, |
| 4583 | cache_control: None, |
| 4584 | }); |
| 4585 | } |
| 4586 | |
| 4587 | // A server whose last connect failed auth-required has no real tools |
| 4588 | // to advertise. In their place offer exactly one synthetic |
| 4589 | // `mcp_<server>_authenticate` tool so the model can self-serve the |
| 4590 | // OAuth login instead of dead-ending on the listing's error item. |
| 4591 | // Never shadow a real tool that owns the same model name. |
| 4592 | { |
| 4593 | let dynamic = self.dynamic_servers.read(); |
| 4594 | for server in &self.needs_auth_servers { |
| 4595 | let Some(config) = self |
| 4596 | .config |
| 4597 | .servers |
| 4598 | .get(server) |
| 4599 | .or_else(|| dynamic.get(server)) |
| 4600 | else { |
| 4601 | continue; |
| 4602 | }; |
| 4603 | if !self.server_allowed(server) |
| 4604 | || !config.is_enabled() |
| 4605 | || !oauth::server_supports_oauth_login(config) |
| 4606 | { |
| 4607 | continue; |
| 4608 | } |
| 4609 | let name = Self::mcp_model_tool_name(server, AUTHENTICATE_TOOL_NAME); |
| 4610 | if !self.tool_allowed(&name) { |
| 4611 | continue; |
| 4612 | } |
| 4613 | if api_tools.iter().any(|tool| tool.name == name) { |
| 4614 | continue; |
| 4615 | } |
| 4616 | api_tools.push(codewhale_models::Tool { |
| 4617 | tool_type: None, |
| 4618 | name, |
| 4619 | description: oauth::authenticate_tool_description(server), |
| 4620 | input_schema: serde_json::json!({ |
| 4621 | "type": "object", |
| 4622 | "properties": {} |
| 4623 | }), |
| 4624 | allowed_callers: Some(vec!["direct".to_string()]), |
| 4625 | defer_loading: Some(false), |
| 4626 | input_examples: None, |
| 4627 | strict: None, |
| 4628 | cache_control: None, |
| 4629 | }); |
| 4630 | } |
| 4631 | } |
| 4632 | |
| 4633 | // Only advertise each resource-listing meta-tool when the servers actually |
| 4634 | // expose the corresponding kind. Previously both were injected whenever any |
| 4635 | // MCP server was configured, so tools-only servers left the model with |
| 4636 | // meta-tools that can only ever return empty results — a wasted tool slot |
| 4637 | // and prompt tokens. Gate each on its own non-empty collection, mirroring |
| 4638 | // the `mcp_read_resource` guard below (`!resources.is_empty()`). |
| 4639 | if !self.all_resources().is_empty() { |
| 4640 | api_tools.push(codewhale_models::Tool { |
| 4641 | tool_type: None, |
| 4642 | name: "list_mcp_resources".to_string(), |
| 4643 | description: "List available MCP resources across servers (optionally filtered by server).".to_string(), |
| 4644 | input_schema: serde_json::json!({ |
| 4645 | "type": "object", |
| 4646 | "properties": { |
| 4647 | "server": { "type": "string", "description": "Optional MCP server name to filter by" } |
| 4648 | } |
| 4649 | }), |
| 4650 | allowed_callers: Some(vec!["direct".to_string()]), |
| 4651 | defer_loading: Some(false), |
| 4652 | input_examples: None, |
| 4653 | strict: None, |
| 4654 | cache_control: None, |
| 4655 | }); |
| 4656 | } |
| 4657 | if !self.all_resource_templates().is_empty() { |
| 4658 | api_tools.push(codewhale_models::Tool { |
| 4659 | tool_type: None, |
| 4660 | name: "list_mcp_resource_templates".to_string(), |
| 4661 | description: "List available MCP resource templates across servers (optionally filtered by server).".to_string(), |
| 4662 | input_schema: serde_json::json!({ |
| 4663 | "type": "object", |
| 4664 | "properties": { |
| 4665 | "server": { "type": "string", "description": "Optional MCP server name to filter by" } |
| 4666 | } |
| 4667 | }), |
| 4668 | allowed_callers: Some(vec!["direct".to_string()]), |
| 4669 | defer_loading: Some(false), |
| 4670 | input_examples: None, |
| 4671 | strict: None, |
| 4672 | cache_control: None, |
| 4673 | }); |
| 4674 | } |
| 4675 | |
| 4676 | // Add resource reading tools if resources exist |
| 4677 | let resources = self.all_resources(); |
| 4678 | if !resources.is_empty() { |
| 4679 | api_tools.push(codewhale_models::Tool { |
| 4680 | tool_type: None, |
| 4681 | name: "mcp_read_resource".to_string(), |
| 4682 | description: "Read a resource from an MCP server using its URI".to_string(), |
| 4683 | input_schema: serde_json::json!({ |
| 4684 | "type": "object", |
| 4685 | "properties": { |
| 4686 | "server": { "type": "string", "description": "The name of the MCP server" }, |
| 4687 | "uri": { "type": "string", "description": "The URI of the resource to read" } |
| 4688 | }, |
| 4689 | "required": ["server", "uri"] |
| 4690 | }), |
| 4691 | allowed_callers: Some(vec!["direct".to_string()]), |
| 4692 | defer_loading: Some(false), |
| 4693 | input_examples: None, |
| 4694 | strict: None, |
| 4695 | cache_control: None, |
| 4696 | }); |
| 4697 | api_tools.push(codewhale_models::Tool { |
| 4698 | tool_type: None, |
| 4699 | name: "read_mcp_resource".to_string(), |
| 4700 | description: "Alias for mcp_read_resource.".to_string(), |
| 4701 | input_schema: serde_json::json!({ |
| 4702 | "type": "object", |
| 4703 | "properties": { |
| 4704 | "server": { "type": "string", "description": "The name of the MCP server" }, |
| 4705 | "uri": { "type": "string", "description": "The URI of the resource to read" } |
| 4706 | }, |
| 4707 | "required": ["server", "uri"] |
| 4708 | }), |
| 4709 | allowed_callers: Some(vec!["direct".to_string()]), |
| 4710 | defer_loading: Some(false), |
| 4711 | input_examples: None, |
| 4712 | strict: None, |
| 4713 | cache_control: None, |
| 4714 | }); |
| 4715 | } |
| 4716 | |
| 4717 | // Add prompt getting tools if prompts exist |
| 4718 | let prompts = self.all_prompts(); |
| 4719 | if !prompts.is_empty() { |
| 4720 | api_tools.push(codewhale_models::Tool { |
| 4721 | tool_type: None, |
| 4722 | name: "mcp_get_prompt".to_string(), |
| 4723 | description: "Get a prompt from an MCP server".to_string(), |
| 4724 | input_schema: serde_json::json!({ |
| 4725 | "type": "object", |
| 4726 | "properties": { |
| 4727 | "server": { "type": "string", "description": "The name of the MCP server" }, |
| 4728 | "name": { "type": "string", "description": "The name of the prompt" }, |
| 4729 | "arguments": { |
| 4730 | "type": "object", |
| 4731 | "description": "Optional arguments for the prompt", |
| 4732 | "additionalProperties": { "type": "string" } |
| 4733 | } |
| 4734 | }, |
| 4735 | "required": ["server", "name"] |
| 4736 | }), |
| 4737 | allowed_callers: Some(vec!["direct".to_string()]), |
| 4738 | defer_loading: Some(false), |
| 4739 | input_examples: None, |
| 4740 | strict: None, |
| 4741 | cache_control: None, |
| 4742 | }); |
| 4743 | } |
| 4744 | |
| 4745 | // Sort by name for prefix-cache stability — the tool block sent to |
| 4746 | // the model needs to be deterministic across runs (#1319). |
| 4747 | api_tools.retain(|tool| self.tool_allowed(&tool.name)); |
| 4748 | api_tools.sort_by(|a, b| a.name.cmp(&b.name)); |
| 4749 | api_tools |
| 4750 | } |
| 4751 | |
| 4752 | /// Apply a child's narrower ceiling without changing the shared pool. |
| 4753 | pub(crate) async fn call_tool_with_disallowed( |
| 4754 | &mut self, |
| 4755 | name: &str, |
| 4756 | input: serde_json::Value, |
| 4757 | rules: &[String], |
| 4758 | ) -> Result<serde_json::Value> { |
| 4759 | Self::authorize_call(&self.disallowed_tools, name, &input)?; |
| 4760 | Self::authorize_call(rules, name, &input)?; |
| 4761 | if !rules.is_empty() |
| 4762 | && rules != self.disallowed_tools.as_slice() |
| 4763 | && matches!(name, "list_mcp_resources" | "list_mcp_resource_templates") |
| 4764 | && input |
| 4765 | .get("server") |
| 4766 | .and_then(serde_json::Value::as_str) |
| 4767 | .is_none() |
| 4768 | { |
| 4769 | self.reload_if_config_changed().await?; |
| 4770 | let servers = self.enabled_server_names(); |
| 4771 | let mut items = Vec::new(); |
| 4772 | for server in servers { |
| 4773 | if Self::server_denied_by(rules, &server) { |
| 4774 | continue; |
| 4775 | } |
| 4776 | let result = if name == "list_mcp_resources" { |
| 4777 | self.list_resources(Some(server.clone())).await |
| 4778 | } else { |
| 4779 | self.list_resource_templates(Some(server.clone())).await |
| 4780 | }; |
| 4781 | match result { |
| 4782 | Ok(mut resources) => items.append(&mut resources), |
| 4783 | Err(error) if oauth::error_looks_auth_required(&error) => { |
| 4784 | let mut item = self.mcp_auth_required_error_item(&server); |
| 4785 | let auth_name = Self::mcp_model_tool_name(&server, AUTHENTICATE_TOOL_NAME); |
| 4786 | if crate::core::engine::tool_catalog::tool_matches_any_rule( |
| 4787 | rules, &auth_name, |
| 4788 | ) { |
| 4789 | item.as_object_mut() |
| 4790 | .expect("error item object") |
| 4791 | .remove("authenticate_tool"); |
| 4792 | item["message"] = |
| 4793 | serde_json::json!("MCP server requires authentication"); |
| 4794 | } |
| 4795 | items.push(item); |
| 4796 | } |
| 4797 | Err(error) => tracing::warn!("MCP resource discovery failed: {error:#}"), |
| 4798 | } |
| 4799 | } |
| 4800 | let field = if name == "list_mcp_resources" { |
| 4801 | "resources" |
| 4802 | } else { |
| 4803 | "templates" |
| 4804 | }; |
| 4805 | return Ok(serde_json::json!({ field: items })); |
| 4806 | } |
| 4807 | let synthetic_auth = self.authenticate_tool_target(name).is_some(); |
| 4808 | let mut result = self.call_tool(name, input).await?; |
| 4809 | if synthetic_auth { |
| 4810 | Self::filter_authenticate_result(&mut result, rules); |
| 4811 | } |
| 4812 | Ok(result) |
| 4813 | } |
| 4814 | |
| 4815 | pub(crate) fn filter_authenticate_result(result: &mut serde_json::Value, rules: &[String]) { |
| 4816 | if let Some(tools) = result |
| 4817 | .get_mut("tools") |
| 4818 | .and_then(serde_json::Value::as_array_mut) |
| 4819 | { |
| 4820 | tools.retain(|name| { |
| 4821 | name.as_str().is_some_and(|name| { |
| 4822 | !crate::core::engine::tool_catalog::tool_matches_any_rule(rules, name) |
| 4823 | }) |
| 4824 | }); |
| 4825 | } |
| 4826 | } |
| 4827 | |
| 4828 | /// Call a tool by its prefixed name (mcp_{server}_{tool}) |
| 4829 | pub async fn call_tool( |
| 4830 | &mut self, |
| 4831 | prefixed_name: &str, |
| 4832 | arguments: serde_json::Value, |
| 4833 | ) -> Result<serde_json::Value> { |
| 4834 | Self::authorize_call(&self.disallowed_tools, prefixed_name, &arguments)?; |
| 4835 | if prefixed_name == "list_mcp_resources" { |
| 4836 | let server = arguments |
| 4837 | .get("server") |
| 4838 | .and_then(|v| v.as_str()) |
| 4839 | .map(str::to_string); |
| 4840 | let resources = self.list_resources(server).await?; |
| 4841 | return Ok(serde_json::json!({ "resources": resources })); |
| 4842 | } |
| 4843 | |
| 4844 | if prefixed_name == "list_mcp_resource_templates" { |
| 4845 | let server = arguments |
| 4846 | .get("server") |
| 4847 | .and_then(|v| v.as_str()) |
| 4848 | .map(str::to_string); |
| 4849 | let templates = self.list_resource_templates(server).await?; |
| 4850 | return Ok(serde_json::json!({ "templates": templates })); |
| 4851 | } |
| 4852 | |
| 4853 | if prefixed_name == "mcp_read_resource" { |
| 4854 | let server_name = arguments |
| 4855 | .get("server") |
| 4856 | .and_then(|v| v.as_str()) |
| 4857 | .context("Missing 'server' argument")?; |
| 4858 | let uri = arguments |
| 4859 | .get("uri") |
| 4860 | .and_then(|v| v.as_str()) |
| 4861 | .context("Missing 'uri' argument")?; |
| 4862 | return self.read_resource(server_name, uri).await; |
| 4863 | } |
| 4864 | |
| 4865 | if prefixed_name == "read_mcp_resource" { |
| 4866 | let server_name = arguments |
| 4867 | .get("server") |
| 4868 | .and_then(|v| v.as_str()) |
| 4869 | .context("Missing 'server' argument")?; |
| 4870 | let uri = arguments |
| 4871 | .get("uri") |
| 4872 | .and_then(|v| v.as_str()) |
| 4873 | .context("Missing 'uri' argument")?; |
| 4874 | return self.read_resource(server_name, uri).await; |
| 4875 | } |
| 4876 | |
| 4877 | if prefixed_name == "mcp_get_prompt" { |
| 4878 | let server_name = arguments |
| 4879 | .get("server") |
| 4880 | .and_then(|v| v.as_str()) |
| 4881 | .context("Missing 'server' argument")?; |
| 4882 | let name = arguments |
| 4883 | .get("name") |
| 4884 | .and_then(|v| v.as_str()) |
| 4885 | .context("Missing 'name' argument")?; |
| 4886 | let args = arguments |
| 4887 | .get("arguments") |
| 4888 | .cloned() |
| 4889 | .unwrap_or(serde_json::json!({})); |
| 4890 | return self.get_prompt(server_name, name, args).await; |
| 4891 | } |
| 4892 | |
| 4893 | // Synthetic self-serve OAuth login: `mcp_<server>_authenticate` runs |
| 4894 | // the same flow `/mcp login` uses and reconnects the server on |
| 4895 | // success, so the real tools resolve in this session. |
| 4896 | if let Some(server_name) = self.authenticate_tool_target(prefixed_name) { |
| 4897 | return self.run_authenticate_tool(&server_name).await; |
| 4898 | } |
| 4899 | |
| 4900 | let route = match self.resolve_advertised_tool(prefixed_name).await { |
| 4901 | Ok(route) => route, |
| 4902 | Err(error) => { |
| 4903 | // A real tool name reached a server whose login has lapsed |
| 4904 | // (stale catalog, or a mid-session 401). Name the one call |
| 4905 | // that recovers it instead of leaving a dead error. |
| 4906 | if let Some(server) = self.needs_auth_server_for_tool_name(prefixed_name) { |
| 4907 | return Err(error.context(self.auth_required_hint(&server))); |
| 4908 | } |
| 4909 | return Err(error); |
| 4910 | } |
| 4911 | }; |
| 4912 | let server_name = route.server_name.clone(); |
| 4913 | let tool_name = route.tool_name.clone(); |
| 4914 | // Copy the global timeouts to avoid borrow conflict |
| 4915 | let global_timeouts = self.config.timeouts; |
| 4916 | let conn = self.get_or_connect(&server_name).await?; |
| 4917 | if conn.catalog_generation != route.catalog_generation { |
| 4918 | anyhow::bail!("MCP catalog changed after tool resolution; retry the call"); |
| 4919 | } |
| 4920 | if conn |
| 4921 | .config() |
| 4922 | .reviewed_plugin |
| 4923 | .as_ref() |
| 4924 | .map(|source| &source.authority) |
| 4925 | != route.plugin_authority.as_ref() |
| 4926 | || !conn.config().is_tool_enabled(&tool_name) |
| 4927 | || !conn.tools().iter().any(|tool| tool.name == tool_name) |
| 4928 | { |
| 4929 | anyhow::bail!("MCP tool '{tool_name}' is disabled for server '{server_name}'"); |
| 4930 | } |
| 4931 | let timeout = conn.config().effective_execute_timeout(&global_timeouts); |
| 4932 | let result = match conn.call_tool(&tool_name, arguments.clone(), timeout).await { |
| 4933 | Ok(result) => Ok(result), |
| 4934 | // A rejected credential is not a stale session: reconnecting |
| 4935 | // replays the same rejection, so it takes the auth-required |
| 4936 | // path below instead of the transparent retry. |
| 4937 | Err(err) |
| 4938 | if is_retriable_mcp_call_error(&err) && !oauth::error_looks_auth_required(&err) => |
| 4939 | { |
| 4940 | tracing::debug!( |
| 4941 | target: "mcp", |
| 4942 | server = server_name, |
| 4943 | tool = tool_name, |
| 4944 | error = %err, |
| 4945 | "retrying MCP tool call after stale session" |
| 4946 | ); |
| 4947 | self.drop_connection(&server_name, "stale session retry"); |
| 4948 | // No `?` here: a reconnect that fails auth-required must |
| 4949 | // still reach the classification below. |
| 4950 | match self.get_or_connect(&server_name).await { |
| 4951 | Ok(conn) => { |
| 4952 | if conn.catalog_generation != route.catalog_generation |
| 4953 | || conn |
| 4954 | .config() |
| 4955 | .reviewed_plugin |
| 4956 | .as_ref() |
| 4957 | .map(|source| &source.authority) |
| 4958 | != route.plugin_authority.as_ref() |
| 4959 | || !conn.config().is_tool_enabled(&tool_name) |
| 4960 | || !conn.tools().iter().any(|tool| tool.name == tool_name) |
| 4961 | { |
| 4962 | Err(anyhow::anyhow!( |
| 4963 | "MCP tool '{tool_name}' is disabled for server '{server_name}'" |
| 4964 | )) |
| 4965 | } else { |
| 4966 | let timeout = conn.config().effective_execute_timeout(&global_timeouts); |
| 4967 | conn.call_tool(&tool_name, arguments, timeout).await |
| 4968 | } |
| 4969 | } |
| 4970 | // A reconnect that fails must not swallow the call error |
| 4971 | // that triggered it: report both, original first. |
| 4972 | Err(reconnect_err) => Err(anyhow::anyhow!( |
| 4973 | "{err:#}; reconnect failed: {reconnect_err:#}" |
| 4974 | )), |
| 4975 | } |
| 4976 | } |
| 4977 | Err(err) => Err(err), |
| 4978 | }; |
| 4979 | // A credential the server stopped accepting mid-session (revoked or |
| 4980 | // rotated elsewhere, refresh rejected) is the same `◆ auth required` |
| 4981 | // class as a failed connect: land it in the same typed state so the |
| 4982 | // next catalog offers the login tool instead of a dead error. |
| 4983 | result.map_err(|err| self.note_live_call_failure(&server_name, err)) |
| 4984 | } |
| 4985 | |
| 4986 | /// Get list of configured server names (static + dynamic) |
| 4987 | #[allow(dead_code)] // Public API for MCP consumers |
| 4988 | pub fn server_names(&self) -> Vec<String> { |
| 4989 | let mut names: Vec<String> = self |
| 4990 | .config |
| 4991 | .servers |
| 4992 | .keys() |
| 4993 | .filter(|name| self.server_allowed(name)) |
| 4994 | .cloned() |
| 4995 | .collect(); |
| 4996 | let dynamic = self.dynamic_servers.read(); |
| 4997 | for name in dynamic.keys() { |
| 4998 | if self.server_allowed(name) && !names.contains(name) { |
| 4999 | names.push(name.clone()); |
| 5000 | } |
| 5001 | } |
| 5002 | names |
| 5003 | } |
| 5004 | |
| 5005 | /// Add a runtime server configuration (in-memory only, not persisted). |
| 5006 | /// |
| 5007 | /// This is used for dynamically started MCP servers from chat context. |
| 5008 | /// Stored in `dynamic_servers` so it doesn't interfere with file-based config reload. |
| 5009 | /// |
| 5010 | /// Returns `Err` if a server with the same name already exists as a static config |
| 5011 | /// or a dynamic config. The caller should surface the error to the LLM/user. |
| 5012 | pub fn add_runtime_server_config( |
| 5013 | &self, |
| 5014 | name: String, |
| 5015 | config: McpServerConfig, |
| 5016 | ) -> Result<(), String> { |
| 5017 | self.require_server(&name) |
| 5018 | .map_err(|error| error.to_string())?; |
| 5019 | if self.config.servers.contains_key(&name) { |
| 5020 | return Err(format!( |
| 5021 | "MCP server '{}' already exists in the config file. \ |
| 5022 | Reconnect with start_mcp_server using only its exact name (omit server), \ |
| 5023 | or run /mcp retry with that name. This preserves its stored credentials.", |
| 5024 | name |
| 5025 | )); |
| 5026 | } |
| 5027 | let mut dynamic = self.dynamic_servers.write(); |
| 5028 | if dynamic.contains_key(&name) { |
| 5029 | return Err(format!( |
| 5030 | "MCP server '{}' was already started earlier in this session. \ |
| 5031 | Reconnect with start_mcp_server using only its exact name (omit server).", |
| 5032 | name |
| 5033 | )); |
| 5034 | } |
| 5035 | let mut config = config; |
| 5036 | config.runtime_added = true; |
| 5037 | dynamic.insert(name, config); |
| 5038 | self.catalog_generation.fetch_add(1, Ordering::SeqCst); |
| 5039 | Ok(()) |
| 5040 | } |
| 5041 | |
| 5042 | /// Remove an in-memory runtime server after a failed start attempt. |
| 5043 | /// This makes dynamic registration transactional: callers may retry the |
| 5044 | /// same deterministic name after correcting an argument or install issue. |
| 5045 | pub fn remove_runtime_server_config(&mut self, name: &str) { |
| 5046 | self.drop_connection(name, "runtime server start rolled back"); |
| 5047 | if self.dynamic_servers.write().remove(name).is_some() { |
| 5048 | self.catalog_generation.fetch_add(1, Ordering::SeqCst); |
| 5049 | } |
| 5050 | } |
| 5051 | |
| 5052 | /// Get list of connected server names |
| 5053 | pub fn connected_servers(&self) -> Vec<&str> { |
| 5054 | self.connections |
| 5055 | .iter() |
| 5056 | .filter(|(name, c)| self.server_allowed(name) && c.is_ready()) |
| 5057 | .map(|(n, _)| n.as_str()) |
| 5058 | .collect() |
| 5059 | } |
| 5060 | |
| 5061 | /// Names of every *enabled* server this pool would connect on the next |
| 5062 | /// turn (static config + dynamic runtime entries). |
| 5063 | /// |
| 5064 | /// Read-only: unlike [`Self::connect_all`], it neither reloads the config |
| 5065 | /// sources nor starts a process. `/preview-request` uses it, together |
| 5066 | /// with [`Self::connected_servers`] and |
| 5067 | /// [`Self::config_sources_unchanged`], to decide whether the currently |
| 5068 | /// connected tool set is *exactly* what the next turn would send — and to |
| 5069 | /// report the tool surface as unavailable when it is not (#1004). |
| 5070 | pub fn enabled_server_names(&self) -> Vec<String> { |
| 5071 | let mut names: Vec<String> = self |
| 5072 | .config |
| 5073 | .servers |
| 5074 | .iter() |
| 5075 | .filter(|(name, server)| server.is_enabled() && self.server_allowed(name)) |
| 5076 | .map(|(name, _)| name.clone()) |
| 5077 | .collect(); |
| 5078 | let dynamic = self.dynamic_servers.read(); |
| 5079 | for (name, server) in dynamic.iter() { |
| 5080 | if self.server_allowed(name) && server.is_enabled() && !names.contains(name) { |
| 5081 | names.push(name.clone()); |
| 5082 | } |
| 5083 | } |
| 5084 | names |
| 5085 | } |
| 5086 | |
| 5087 | /// Compare against the freshly authorized merged configuration without |
| 5088 | /// reloading or disconnecting any sibling transport. |
| 5089 | pub(crate) fn config_matches(&self, config: &McpConfig) -> bool { |
| 5090 | hash_mcp_config(config) == self.config_hash |
| 5091 | } |
| 5092 | |
| 5093 | /// Whether every configured MCP source still has the mtime this pool last |
| 5094 | /// read, i.e. whether `connect_all` would find anything new. |
| 5095 | /// |
| 5096 | /// Stats files; never reads, parses, reloads, or drops a connection. A |
| 5097 | /// pool with no configured source is trivially unchanged. |
| 5098 | pub fn config_sources_unchanged(&self) -> bool { |
| 5099 | if self.config_sources.is_empty() { |
| 5100 | return true; |
| 5101 | } |
| 5102 | let current: Vec<_> = self |
| 5103 | .config_sources |
| 5104 | .iter() |
| 5105 | .map(|path| mcp_config_mtime(path)) |
| 5106 | .collect(); |
| 5107 | current == self.last_mtimes |
| 5108 | } |
| 5109 | |
| 5110 | /// Graceful shutdown of every connection in the pool: send SIGTERM to |
| 5111 | /// each stdio child and give them a short grace period before drop |
| 5112 | /// fires SIGKILL. Whalescale#420. |
| 5113 | /// |
| 5114 | /// Call from the TUI exit path *before* dropping the pool to give |
| 5115 | /// MCP servers a chance to flush state. The fallback Drop on |
| 5116 | /// `StdioTransport` still sends SIGTERM if this never runs, so even |
| 5117 | /// abnormal exits avoid leaking PIDs without a signal. |
| 5118 | pub async fn shutdown_all(&mut self) { |
| 5119 | let names: Vec<String> = self.connections.keys().cloned().collect(); |
| 5120 | for name in names { |
| 5121 | if let Some(conn) = self.connections.get_mut(&name) { |
| 5122 | conn.transport.shutdown().await; |
| 5123 | } |
| 5124 | } |
| 5125 | self.connections.clear(); |
| 5126 | } |
| 5127 | |
| 5128 | /// Check if a tool name is an MCP tool |
| 5129 | pub fn is_mcp_tool(name: &str) -> bool { |
| 5130 | name.starts_with("mcp_") |
| 5131 | || matches!( |
| 5132 | name, |
| 5133 | "list_mcp_resources" | "list_mcp_resource_templates" | "read_mcp_resource" |
| 5134 | ) |
| 5135 | } |
| 5136 | } |
| 5137 | |
| 5138 | #[derive(Debug, Clone, Copy, PartialEq, Eq)] |
| 5139 | pub enum McpWriteStatus { |
| 5140 | Created, |
| 5141 | Overwritten, |
| 5142 | SkippedExists, |
| 5143 | } |
| 5144 | |
| 5145 | #[derive(Debug, Clone, PartialEq, Eq)] |
| 5146 | pub struct McpDiscoveredItem { |
| 5147 | pub name: String, |
| 5148 | pub model_name: String, |
| 5149 | pub description: Option<String>, |
| 5150 | } |
| 5151 | |
| 5152 | #[derive(Debug, Clone, PartialEq, Eq)] |
| 5153 | pub struct McpServerSnapshot { |
| 5154 | pub name: String, |
| 5155 | pub enabled: bool, |
| 5156 | pub required: bool, |
| 5157 | pub transport: String, |
| 5158 | pub command_or_url: String, |
| 5159 | pub connect_timeout: u64, |
| 5160 | pub execute_timeout: u64, |
| 5161 | pub read_timeout: u64, |
| 5162 | pub connected: bool, |
| 5163 | pub error: Option<String>, |
| 5164 | /// Typed `◆ auth required` state: the server's most recent connect |
| 5165 | /// attempt failed because a login is missing, expired, or revoked |
| 5166 | /// (401 / OAuth not logged in / `invalid_grant`). Derived from the pool's |
| 5167 | /// needs-auth set — the same source that decides whether the model gets |
| 5168 | /// the synthetic `mcp_<server>_authenticate` tool — so every surface |
| 5169 | /// (session boot row, `/mcp` manager, Extensions, model catalog) agrees. |
| 5170 | pub auth_required: bool, |
| 5171 | pub capability_metadata: McpServerCapabilityMetadata, |
| 5172 | pub tools: Vec<McpDiscoveredItem>, |
| 5173 | pub resources: Vec<McpDiscoveredItem>, |
| 5174 | pub prompts: Vec<McpDiscoveredItem>, |
| 5175 | } |
| 5176 | |
| 5177 | impl McpServerSnapshot { |
| 5178 | /// Recovery for this observed server. The typed auth-required state wins |
| 5179 | /// over error-text sniffing so a needs-auth server always routes to |
| 5180 | /// `/mcp login <name>`. |
| 5181 | #[must_use] |
| 5182 | pub fn recovery_kind(&self, oauth_capable: bool) -> Option<McpRecoveryKind> { |
| 5183 | if self.enabled && !self.connected && self.auth_required { |
| 5184 | return Some(McpRecoveryKind::Reauth); |
| 5185 | } |
| 5186 | mcp_recovery_kind( |
| 5187 | self.enabled, |
| 5188 | true, |
| 5189 | self.connected, |
| 5190 | self.error.as_deref(), |
| 5191 | oauth_capable, |
| 5192 | ) |
| 5193 | } |
| 5194 | } |
| 5195 | |
| 5196 | #[derive(Debug, Clone, PartialEq, Eq)] |
| 5197 | pub struct McpManagerSnapshot { |
| 5198 | pub config_path: std::path::PathBuf, |
| 5199 | pub config_exists: bool, |
| 5200 | pub reload_required: bool, |
| 5201 | pub servers: Vec<McpServerSnapshot>, |
| 5202 | } |
| 5203 | |
| 5204 | /// First-class recovery for a configured MCP server. Commands named here exist: |
| 5205 | /// `/mcp login`, `/mcp reload`, `/mcp validate`, `/mcp enable`. There is no |
| 5206 | /// `/mcp auth`. |
| 5207 | #[derive(Debug, Clone, Copy, PartialEq, Eq)] |
| 5208 | pub enum McpRecoveryKind { |
| 5209 | Enable, |
| 5210 | Connect, |
| 5211 | Reconnect, |
| 5212 | Reauth, |
| 5213 | Diagnose, |
| 5214 | } |
| 5215 | |
| 5216 | impl McpRecoveryKind { |
| 5217 | #[must_use] |
| 5218 | pub fn slash_command(self, name: &str) -> String { |
| 5219 | match self { |
| 5220 | Self::Enable => format!("/mcp enable {name}"), |
| 5221 | // Reconnect one server, not all of them. A row that reads |
| 5222 | // `[reconnect] aws` and then reloads all 23 configured servers is |
| 5223 | // not the action it advertised: it takes ~40 s, it disturbs every |
| 5224 | // healthy connection, and the row the user aimed at is still |
| 5225 | // pending when the list comes back. `/mcp retry <name>` reaches |
| 5226 | // `retry_mcp_server`, which reconnects exactly that server. |
| 5227 | Self::Connect | Self::Reconnect if mcp_name_is_command_safe(name) => { |
| 5228 | format!("/mcp retry {name}") |
| 5229 | } |
| 5230 | // A name the command line cannot carry safely still gets the |
| 5231 | // blunt instrument rather than a quoted-argument hazard. |
| 5232 | Self::Connect | Self::Reconnect => "/mcp reload".to_string(), |
| 5233 | Self::Reauth => format!("/mcp login {name}"), |
| 5234 | Self::Diagnose if mcp_name_is_command_safe(name) => format!("/mcp validate {name}"), |
| 5235 | Self::Diagnose => "/mcp validate".to_string(), |
| 5236 | } |
| 5237 | } |
| 5238 | |
| 5239 | #[must_use] |
| 5240 | pub fn label_key(self) -> codewhale_localization::MessageId { |
| 5241 | match self { |
| 5242 | Self::Enable => codewhale_localization::MessageId::ExtensionsActionEnable, |
| 5243 | Self::Connect => codewhale_localization::MessageId::ExtensionsActionConnect, |
| 5244 | Self::Reconnect => codewhale_localization::MessageId::ExtensionsActionReconnect, |
| 5245 | Self::Reauth => codewhale_localization::MessageId::ExtensionsActionReauth, |
| 5246 | Self::Diagnose => codewhale_localization::MessageId::ExtensionsActionDiagnose, |
| 5247 | } |
| 5248 | } |
| 5249 | } |
| 5250 | |
| 5251 | #[must_use] |
| 5252 | pub fn mcp_name_is_command_safe(name: &str) -> bool { |
| 5253 | !name.is_empty() |
| 5254 | && name |
| 5255 | .chars() |
| 5256 | .all(|ch| ch.is_ascii_alphanumeric() || matches!(ch, '-' | '_' | '.')) |
| 5257 | } |
| 5258 | |
| 5259 | /// Display target for an MCP server row: stdio servers show the command |
| 5260 | /// name only — no `./…` relative-path prefix, no directories, no args — |
| 5261 | /// while URL servers keep their full URL, which is the identity. The |
| 5262 | /// snapshot's `command_or_url` keeps full fidelity for the engine and the |
| 5263 | /// wire; this is presentation only. |
| 5264 | #[must_use] |
| 5265 | pub fn mcp_display_target(transport: &str, command_or_url: &str) -> String { |
| 5266 | if transport != "stdio" { |
| 5267 | return command_or_url.to_string(); |
| 5268 | } |
| 5269 | let command = command_or_url |
| 5270 | .split_whitespace() |
| 5271 | .next() |
| 5272 | .unwrap_or(command_or_url); |
| 5273 | let name = command.rsplit(['/', '\\']).next().unwrap_or(command); |
| 5274 | if name.is_empty() { |
| 5275 | command_or_url.to_string() |
| 5276 | } else { |
| 5277 | name.to_string() |
| 5278 | } |
| 5279 | } |
| 5280 | |
| 5281 | #[must_use] |
| 5282 | pub fn mcp_server_oauth_capable(config: &McpServerConfig) -> bool { |
| 5283 | config.url.is_some() |
| 5284 | && (config.oauth.is_some() || !config.scopes.is_empty() || config.oauth_resource.is_some()) |
| 5285 | } |
| 5286 | |
| 5287 | #[must_use] |
| 5288 | pub fn mcp_recovery_kind( |
| 5289 | enabled: bool, |
| 5290 | inspected: bool, |
| 5291 | connected: bool, |
| 5292 | error: Option<&str>, |
| 5293 | oauth_capable: bool, |
| 5294 | ) -> Option<McpRecoveryKind> { |
| 5295 | if !enabled { |
| 5296 | return Some(McpRecoveryKind::Enable); |
| 5297 | } |
| 5298 | if let Some(error) = error { |
| 5299 | if oauth::error_text_looks_auth_required(error) { |
| 5300 | return Some(McpRecoveryKind::Reauth); |
| 5301 | } |
| 5302 | return Some(McpRecoveryKind::Diagnose); |
| 5303 | } |
| 5304 | if !inspected { |
| 5305 | return Some(McpRecoveryKind::Connect); |
| 5306 | } |
| 5307 | if connected { |
| 5308 | // A server that is enabled, inspected, connected and erroring on |
| 5309 | // nothing needs no recovery. It used to be labelled `diagnose`, so |
| 5310 | // every healthy row advertised a repair it did not need — founder |
| 5311 | // live-test: "even the ones that are connected say diagnose lol". |
| 5312 | return None; |
| 5313 | } |
| 5314 | if oauth_capable { |
| 5315 | return Some(McpRecoveryKind::Reauth); |
| 5316 | } |
| 5317 | Some(McpRecoveryKind::Reconnect) |
| 5318 | } |
| 5319 | |
| 5320 | pub fn load_config(path: &Path) -> Result<McpConfig> { |
| 5321 | validate_mcp_config_path(path)?; |
| 5322 | let Some(contents) = read_mcp_config_file(path)? else { |
| 5323 | return Ok(McpConfig::default()); |
| 5324 | }; |
| 5325 | serde_json::from_str(&contents).map_err(|_| { |
| 5326 | anyhow::anyhow!( |
| 5327 | "Failed to parse MCP config {}; file contents were omitted", |
| 5328 | codewhale_config::quote_os_path(path) |
| 5329 | ) |
| 5330 | }) |
| 5331 | } |
| 5332 | |
| 5333 | /// Maximum bytes read from an MCP config file. Configs are kilobytes. |
| 5334 | const MAX_MCP_CONFIG_BYTES: u64 = 1024 * 1024; |
| 5335 | |
| 5336 | fn read_mcp_config_file(path: &Path) -> Result<Option<String>> { |
| 5337 | let metadata = match fs::symlink_metadata(path) { |
| 5338 | Ok(metadata) => metadata, |
| 5339 | Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(None), |
| 5340 | Err(err) => { |
| 5341 | return Err(err) |
| 5342 | .with_context(|| format!("Failed to inspect MCP config {}", path.display())); |
| 5343 | } |
| 5344 | }; |
| 5345 | let file_type = metadata.file_type(); |
| 5346 | if file_type.is_symlink() || !file_type.is_file() { |
| 5347 | anyhow::bail!("MCP config path must be a regular file: {}", path.display()); |
| 5348 | } |
| 5349 | |
| 5350 | let file = open_mcp_config_file(path) |
| 5351 | .with_context(|| format!("Failed to read MCP config {}", path.display()))?; |
| 5352 | let mut contents = String::new(); |
| 5353 | file.take(MAX_MCP_CONFIG_BYTES + 1) |
| 5354 | .read_to_string(&mut contents) |
| 5355 | .with_context(|| format!("Failed to read MCP config {}", path.display()))?; |
| 5356 | if contents.len() as u64 > MAX_MCP_CONFIG_BYTES { |
| 5357 | anyhow::bail!("MCP config {} exceeds the 1 MiB limit", path.display()); |
| 5358 | } |
| 5359 | Ok(Some(contents)) |
| 5360 | } |
| 5361 | |
| 5362 | #[cfg(unix)] |
| 5363 | fn open_mcp_config_file(path: &Path) -> std::io::Result<fs::File> { |
| 5364 | use std::os::unix::fs::OpenOptionsExt; |
| 5365 | |
| 5366 | fs::OpenOptions::new() |
| 5367 | .read(true) |
| 5368 | .custom_flags(libc::O_NOFOLLOW) |
| 5369 | .open(path) |
| 5370 | } |
| 5371 | |
| 5372 | #[cfg(not(unix))] |
| 5373 | fn open_mcp_config_file(path: &Path) -> std::io::Result<fs::File> { |
| 5374 | fs::File::open(path) |
| 5375 | } |
| 5376 | |
| 5377 | pub fn workspace_mcp_config_path(workspace: &Path) -> PathBuf { |
| 5378 | normalize_workspace_path(workspace) |
| 5379 | .join(".codewhale") |
| 5380 | .join("mcp.json") |
| 5381 | } |
| 5382 | |
| 5383 | pub fn load_config_with_workspace(global_path: &Path, workspace: &Path) -> Result<McpConfig> { |
| 5384 | let plugins = crate::plugins::PluginRegistry::empty(workspace); |
| 5385 | load_config_with_workspace_and_plugins(global_path, workspace, &plugins) |
| 5386 | } |
| 5387 | |
| 5388 | pub fn load_config_with_workspace_and_plugins( |
| 5389 | global_path: &Path, |
| 5390 | workspace: &Path, |
| 5391 | plugins: &crate::plugins::PluginRegistry, |
| 5392 | ) -> Result<McpConfig> { |
| 5393 | let mut merged = load_config(global_path)?; |
| 5394 | let workspace = checked_workspace_path(workspace)?; |
| 5395 | let project_path = checked_workspace_mcp_config_path(&workspace)?; |
| 5396 | if !project_path.exists() || paths_refer_to_same_config(global_path, &project_path) { |
| 5397 | return merge_plugin_mcp_servers(merged, plugins); |
| 5398 | } |
| 5399 | // Workspace-local MCP can spawn stdio servers, so it is only honored after |
| 5400 | // the user has trusted this workspace in user-owned config. Do not accept |
| 5401 | // project-local legacy trust markers here: a repository could carry those |
| 5402 | // files itself and silently reintroduce the project-scope `mcp_config_path` |
| 5403 | // risk denied in #417. |
| 5404 | if !workspace_allows_project_mcp_config(&workspace) { |
| 5405 | return merge_plugin_mcp_servers(merged, plugins); |
| 5406 | } |
| 5407 | |
| 5408 | let mut project = load_config(&project_path)?; |
| 5409 | for server in project.servers.values_mut() { |
| 5410 | if server.command.is_some() && server.url.is_none() { |
| 5411 | server.cwd = Some(resolve_project_mcp_cwd(&workspace, server.cwd.as_deref())?); |
| 5412 | } |
| 5413 | } |
| 5414 | merged.servers.extend(project.servers); |
| 5415 | |
| 5416 | merge_plugin_mcp_servers(merged, plugins) |
| 5417 | } |
| 5418 | |
| 5419 | fn merge_plugin_mcp_servers( |
| 5420 | config: McpConfig, |
| 5421 | registry: &crate::plugins::PluginRegistry, |
| 5422 | ) -> Result<McpConfig> { |
| 5423 | let Some(state_path) = registry.state_path().map(Path::to_path_buf) else { |
| 5424 | return Ok(config); |
| 5425 | }; |
| 5426 | let plugins = registry |
| 5427 | .active_plugins() |
| 5428 | .into_iter() |
| 5429 | .filter_map(|plugin| { |
| 5430 | plugin |
| 5431 | .authority(state_path.clone(), registry.workspace().to_path_buf()) |
| 5432 | .map(|authority| (plugin.name().to_string(), plugin.clone(), authority)) |
| 5433 | }) |
| 5434 | .collect::<Vec<_>>(); |
| 5435 | |
| 5436 | let host_environment = registry.host_environment().ok_or_else(|| { |
| 5437 | anyhow::anyhow!("active plugin registry is missing its pre-dotenv environment snapshot") |
| 5438 | })?; |
| 5439 | merge_plugin_mcp_servers_from_plugins_with_environment(config, plugins, host_environment) |
| 5440 | } |
| 5441 | |
| 5442 | fn merge_plugin_mcp_servers_from_plugins_with_environment( |
| 5443 | mut config: McpConfig, |
| 5444 | plugins: impl IntoIterator< |
| 5445 | Item = ( |
| 5446 | String, |
| 5447 | crate::plugins::types::LoadedPlugin, |
| 5448 | crate::plugins::types::PluginAuthority, |
| 5449 | ), |
| 5450 | >, |
| 5451 | host_environment: Arc<crate::plugins::HostEnvironment>, |
| 5452 | ) -> Result<McpConfig> { |
| 5453 | for (plugin_name, plugin, authority) in plugins { |
| 5454 | // Adapter-level denial keeps headless paths fail-closed even if a |
| 5455 | // future caller accidentally passes the full inventory instead of the |
| 5456 | // registry's active-only view. |
| 5457 | if !plugin.active() { |
| 5458 | continue; |
| 5459 | } |
| 5460 | if let Some(mcp_servers) = &plugin.manifest.mcp_servers { |
| 5461 | let mut mcp_servers = mcp_servers.iter().collect::<Vec<_>>(); |
| 5462 | mcp_servers.sort_by_key(|(name, _)| *name); |
| 5463 | for (server_name, server_config) in mcp_servers { |
| 5464 | let required_capability = if server_config.command.is_some() |
| 5465 | && server_config.url.is_none() |
| 5466 | { |
| 5467 | crate::plugins::activation::PluginActivationCapability::McpStdio |
| 5468 | } else if server_config.url.is_some() && server_config.command.is_none() { |
| 5469 | crate::plugins::activation::PluginActivationCapability::McpRemote |
| 5470 | } else { |
| 5471 | tracing::warn!( |
| 5472 | target: "mcp", |
| 5473 | plugin = %plugin_name, |
| 5474 | server = %server_name, |
| 5475 | "plugin MCP server is neither a reviewed stdio nor remote transport; denying it" |
| 5476 | ); |
| 5477 | continue; |
| 5478 | }; |
| 5479 | if !plugin.component_active(required_capability) |
| 5480 | || crate::plugins::registry::verify_plugin_component_authority( |
| 5481 | &authority, |
| 5482 | required_capability, |
| 5483 | ) |
| 5484 | .is_err() |
| 5485 | { |
| 5486 | tracing::warn!( |
| 5487 | target: "mcp", |
| 5488 | plugin = %plugin_name, |
| 5489 | server = %server_name, |
| 5490 | capability = required_capability.as_str(), |
| 5491 | "plugin bundle changed after review or this transport is inactive; denying the MCP server until reload and re-review" |
| 5492 | ); |
| 5493 | continue; |
| 5494 | } |
| 5495 | let qualified_name = qualified_plugin_server_name(&plugin_name, server_name); |
| 5496 | if config.servers.contains_key(&qualified_name) { |
| 5497 | tracing::warn!( |
| 5498 | target: "mcp", |
| 5499 | plugin = %plugin_name, |
| 5500 | server = %server_name, |
| 5501 | qualified_name = %qualified_name, |
| 5502 | "explicit MCP configuration keeps precedence over a colliding plugin server" |
| 5503 | ); |
| 5504 | continue; |
| 5505 | } |
| 5506 | let mut server_config = server_config.clone(); |
| 5507 | |
| 5508 | if server_config.command.is_some() && server_config.url.is_none() { |
| 5509 | let staged_root = plugin |
| 5510 | .staged_root |
| 5511 | .as_deref() |
| 5512 | .context("active plugin is missing its runtime snapshot")?; |
| 5513 | server_config.cwd = Some(resolve_plugin_mcp_cwd( |
| 5514 | staged_root, |
| 5515 | server_config.cwd.as_deref(), |
| 5516 | )?); |
| 5517 | freeze_plugin_stdio_paths(&mut server_config, staged_root)?; |
| 5518 | } |
| 5519 | server_config.reviewed_plugin = Some(ReviewedPluginMcpSource::from_authority( |
| 5520 | authority.clone(), |
| 5521 | server_config.url.as_deref(), |
| 5522 | Arc::clone(&host_environment), |
| 5523 | )?); |
| 5524 | |
| 5525 | config.servers.insert(qualified_name, server_config); |
| 5526 | } |
| 5527 | } |
| 5528 | } |
| 5529 | |
| 5530 | Ok(config) |
| 5531 | } |
| 5532 | |
| 5533 | #[cfg(test)] |
| 5534 | fn merge_plugin_mcp_servers_from_plugins( |
| 5535 | config: McpConfig, |
| 5536 | plugins: impl IntoIterator< |
| 5537 | Item = ( |
| 5538 | String, |
| 5539 | crate::plugins::types::LoadedPlugin, |
| 5540 | crate::plugins::types::PluginAuthority, |
| 5541 | ), |
| 5542 | >, |
| 5543 | ) -> Result<McpConfig> { |
| 5544 | merge_plugin_mcp_servers_from_plugins_with_environment( |
| 5545 | config, |
| 5546 | plugins, |
| 5547 | Arc::new(crate::plugins::HostEnvironment::capture()), |
| 5548 | ) |
| 5549 | } |
| 5550 | |
| 5551 | fn qualified_plugin_server_name(plugin_name: &str, server_name: &str) -> String { |
| 5552 | format!( |
| 5553 | "plugin-{}-{}-{}", |
| 5554 | plugin_name.len(), |
| 5555 | plugin_name, |
| 5556 | server_name |
| 5557 | ) |
| 5558 | } |
| 5559 | |
| 5560 | fn freeze_plugin_stdio_paths(config: &mut McpServerConfig, staged_root: &Path) -> Result<()> { |
| 5561 | if let Some(command) = config.command.as_mut() |
| 5562 | && (command.contains('/') || command.contains('\\')) |
| 5563 | { |
| 5564 | let frozen = resolve_plugin_mcp_cwd(staged_root, Some(Path::new(command)))?; |
| 5565 | *command = frozen.display().to_string(); |
| 5566 | } |
| 5567 | let runtime_cwd = config.cwd.as_deref().unwrap_or(staged_root).to_path_buf(); |
| 5568 | for argument in &mut config.args { |
| 5569 | if argument.starts_with('-') || Path::new(argument).is_absolute() { |
| 5570 | continue; |
| 5571 | } |
| 5572 | let candidate = normalize_path_components(&runtime_cwd.join(argument.as_str())); |
| 5573 | if candidate.exists() { |
| 5574 | let frozen = candidate |
| 5575 | .canonicalize() |
| 5576 | .context("failed to freeze reviewed plugin MCP argument path")?; |
| 5577 | if !frozen.starts_with(staged_root) { |
| 5578 | anyhow::bail!("reviewed plugin MCP argument path escaped its staged root"); |
| 5579 | } |
| 5580 | *argument = frozen.display().to_string(); |
| 5581 | } |
| 5582 | } |
| 5583 | Ok(()) |
| 5584 | } |
| 5585 | |
| 5586 | fn resolve_plugin_mcp_cwd(plugin_path: &Path, cwd: Option<&Path>) -> Result<PathBuf> { |
| 5587 | let cwd = match cwd { |
| 5588 | Some(cwd) if cwd.is_relative() => normalize_path_components(&plugin_path.join(cwd)), |
| 5589 | Some(cwd) => normalize_path_components(cwd), |
| 5590 | None => plugin_path.to_path_buf(), |
| 5591 | }; |
| 5592 | let resolved = cwd |
| 5593 | .canonicalize() |
| 5594 | .unwrap_or_else(|_| normalize_path_components(&cwd)); |
| 5595 | if !resolved.starts_with(plugin_path) { |
| 5596 | anyhow::bail!("reviewed plugin MCP path escaped its staged root"); |
| 5597 | } |
| 5598 | Ok(resolved) |
| 5599 | } |
| 5600 | |
| 5601 | fn workspace_allows_project_mcp_config(workspace: &Path) -> bool { |
| 5602 | crate::config::is_workspace_trusted(workspace) |
| 5603 | } |
| 5604 | |
| 5605 | fn checked_workspace_mcp_config_path(workspace: &Path) -> Result<PathBuf> { |
| 5606 | Ok(checked_workspace_path(workspace)? |
| 5607 | .join(".codewhale") |
| 5608 | .join("mcp.json")) |
| 5609 | } |
| 5610 | |
| 5611 | fn checked_workspace_path(workspace: &Path) -> Result<PathBuf> { |
| 5612 | if workspace.as_os_str().is_empty() { |
| 5613 | anyhow::bail!("workspace path cannot be empty"); |
| 5614 | } |
| 5615 | if workspace |
| 5616 | .components() |
| 5617 | .any(|component| matches!(component, Component::ParentDir)) |
| 5618 | { |
| 5619 | anyhow::bail!("workspace path cannot contain '..' components"); |
| 5620 | } |
| 5621 | let absolute = if workspace.is_absolute() { |
| 5622 | workspace.to_path_buf() |
| 5623 | } else { |
| 5624 | std::env::current_dir() |
| 5625 | .context("failed to resolve current directory for workspace")? |
| 5626 | .join(workspace) |
| 5627 | }; |
| 5628 | match absolute.canonicalize() { |
| 5629 | Ok(path) => Ok(path), |
| 5630 | Err(err) if err.kind() == std::io::ErrorKind::NotFound => { |
| 5631 | Ok(normalize_path_components(&absolute)) |
| 5632 | } |
| 5633 | Err(err) => { |
| 5634 | Err(err).with_context(|| format!("failed to resolve workspace {}", workspace.display())) |
| 5635 | } |
| 5636 | } |
| 5637 | } |
| 5638 | |
| 5639 | fn normalize_workspace_path(workspace: &Path) -> PathBuf { |
| 5640 | if let Ok(canonical) = workspace.canonicalize() { |
| 5641 | return canonical; |
| 5642 | } |
| 5643 | let absolute = if workspace.is_absolute() { |
| 5644 | workspace.to_path_buf() |
| 5645 | } else { |
| 5646 | std::env::current_dir() |
| 5647 | .unwrap_or_else(|_| PathBuf::from(".")) |
| 5648 | .join(workspace) |
| 5649 | }; |
| 5650 | normalize_path_components(&absolute) |
| 5651 | } |
| 5652 | |
| 5653 | fn resolve_project_mcp_cwd(workspace: &Path, cwd: Option<&Path>) -> Result<PathBuf> { |
| 5654 | let cwd = match cwd { |
| 5655 | Some(cwd) if cwd.is_relative() => normalize_path_components(&workspace.join(cwd)), |
| 5656 | Some(cwd) => normalize_path_components(cwd), |
| 5657 | None => workspace.to_path_buf(), |
| 5658 | }; |
| 5659 | let resolved = cwd |
| 5660 | .canonicalize() |
| 5661 | .unwrap_or_else(|_| normalize_path_components(&cwd)); |
| 5662 | if !resolved.starts_with(workspace) { |
| 5663 | anyhow::bail!( |
| 5664 | "Project MCP server cwd must stay within workspace: {}", |
| 5665 | resolved.display() |
| 5666 | ); |
| 5667 | } |
| 5668 | Ok(resolved) |
| 5669 | } |
| 5670 | |
| 5671 | fn normalize_path_components(path: &Path) -> PathBuf { |
| 5672 | let mut normalized = PathBuf::new(); |
| 5673 | for component in path.components() { |
| 5674 | match component { |
| 5675 | Component::Prefix(_) | Component::RootDir => { |
| 5676 | normalized.push(component.as_os_str()); |
| 5677 | } |
| 5678 | Component::CurDir => {} |
| 5679 | Component::ParentDir => { |
| 5680 | normalized.pop(); |
| 5681 | } |
| 5682 | Component::Normal(part) => normalized.push(part), |
| 5683 | } |
| 5684 | } |
| 5685 | if normalized.as_os_str().is_empty() { |
| 5686 | PathBuf::from(".") |
| 5687 | } else { |
| 5688 | normalized |
| 5689 | } |
| 5690 | } |
| 5691 | |
| 5692 | fn paths_refer_to_same_config(left: &Path, right: &Path) -> bool { |
| 5693 | match (left.canonicalize(), right.canonicalize()) { |
| 5694 | (Ok(left), Ok(right)) => left == right, |
| 5695 | _ => normalize_workspace_path(left) == normalize_workspace_path(right), |
| 5696 | } |
| 5697 | } |
| 5698 | |
| 5699 | /// Rebuild a JSON value with every object's keys in sorted order. |
| 5700 | /// |
| 5701 | /// [`McpConfig`] is full of `HashMap`s (`servers`, and per-server `env`, |
| 5702 | /// `headers`, `env_headers`), and `serde_json` is built here with |
| 5703 | /// `preserve_order`, so a serialization inherits whatever order the source |
| 5704 | /// `HashMap` happened to iterate in. Two `HashMap`s built separately in one |
| 5705 | /// process do *not* share an iteration order — `RandomState` re-seeds per map |
| 5706 | /// — so the raw serialization of two structurally identical configs differs. |
| 5707 | /// Sorting first makes the byte form depend only on content. The value tree |
| 5708 | /// here is the fixed `McpConfig` shape, so the recursion depth is bounded by |
| 5709 | /// that struct, not by untrusted input. |
| 5710 | fn canonicalize_json_keys(value: serde_json::Value) -> serde_json::Value { |
| 5711 | match value { |
| 5712 | serde_json::Value::Object(map) => { |
| 5713 | let mut entries: Vec<(String, serde_json::Value)> = map.into_iter().collect(); |
| 5714 | entries.sort_by(|(left, _), (right, _)| left.cmp(right)); |
| 5715 | serde_json::Value::Object( |
| 5716 | entries |
| 5717 | .into_iter() |
| 5718 | .map(|(key, value)| (key, canonicalize_json_keys(value))) |
| 5719 | .collect(), |
| 5720 | ) |
| 5721 | } |
| 5722 | serde_json::Value::Array(items) => { |
| 5723 | serde_json::Value::Array(items.into_iter().map(canonicalize_json_keys).collect()) |
| 5724 | } |
| 5725 | other => other, |
| 5726 | } |
| 5727 | } |
| 5728 | |
| 5729 | /// 64-bit content hash of an [`McpConfig`]. Used by [`McpPool`] to decide |
| 5730 | /// whether a freshly-read config differs from the one currently driving the |
| 5731 | /// live connections. Hashing the JSON serialization avoids forcing every |
| 5732 | /// nested config type to derive `Hash` (the timeouts struct, network policy |
| 5733 | /// stubs, etc.); the serialization is key-sorted first so the hash depends on |
| 5734 | /// content alone and not on per-`HashMap` iteration order. The hash is stable |
| 5735 | /// within a process for structurally identical configs, and across runs of the |
| 5736 | /// same Rust toolchain for byte-identical input. |
| 5737 | fn hash_mcp_config(config: &McpConfig) -> u64 { |
| 5738 | use std::hash::{Hash, Hasher}; |
| 5739 | let canonical = serde_json::to_value(config) |
| 5740 | .map(canonicalize_json_keys) |
| 5741 | .unwrap_or(serde_json::Value::Null); |
| 5742 | let bytes = serde_json::to_vec(&canonical).unwrap_or_default(); |
| 5743 | let mut hasher = std::collections::hash_map::DefaultHasher::new(); |
| 5744 | bytes.hash(&mut hasher); |
| 5745 | hasher.finish() |
| 5746 | } |
| 5747 | |
| 5748 | /// Best-effort fetch of the MCP config file's last-modified time. Returns |
| 5749 | /// `None` when the file is missing, when stat fails, when the platform |
| 5750 | /// doesn't expose mtime, or when the path fails the same allow-list check |
| 5751 | /// that MCP configuration reads and mutations apply. The lazy-reload check in |
| 5752 | /// `McpPool::get_or_connect` treats `None` as "skip the check this turn", |
| 5753 | /// so a rejected path simply degrades to "no auto-reload" rather than an |
| 5754 | /// error path. Callers already validate via `validate_mcp_config_path` at |
| 5755 | /// construction time; the redundant validation here keeps this helper |
| 5756 | /// safe-by-construction for any future caller and ties the validation to |
| 5757 | /// the call site rather than relying on cross-function reasoning. |
| 5758 | fn mcp_config_mtime(path: &Path) -> Option<std::time::SystemTime> { |
| 5759 | validate_mcp_config_path(path).ok()?; |
| 5760 | fs::metadata(path).ok()?.modified().ok() |
| 5761 | } |
| 5762 | |
| 5763 | /// A stale caller must reload instead of overwriting another process's edit. |
| 5764 | #[derive(Debug)] |
| 5765 | pub struct McpRevisionConflict; |
| 5766 | impl std::fmt::Display for McpRevisionConflict { |
| 5767 | fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { |
| 5768 | f.write_str("MCP configuration changed; reload it before saving") |
| 5769 | } |
| 5770 | } |
| 5771 | impl std::error::Error for McpRevisionConflict {} |
| 5772 | |
| 5773 | fn config_revision(raw: Option<&str>) -> String { |
| 5774 | raw.map_or_else( |
| 5775 | || "mcp-v1-absent".to_owned(), |
| 5776 | |raw| format!("mcp-v1-{}", crate::hashing::sha256_hex(raw.as_bytes())), |
| 5777 | ) |
| 5778 | } |
| 5779 | |
| 5780 | pub fn read_config_revision(path: &Path) -> Result<String> { |
| 5781 | validate_mcp_config_path(path)?; |
| 5782 | let raw = read_mcp_config_file(path)?; |
| 5783 | Ok(config_revision(raw.as_deref())) |
| 5784 | } |
| 5785 | |
| 5786 | /// Apply only changed known fields to the original JSON. Unknown fields in |
| 5787 | /// unrelated objects and in edited server entries remain operator-owned. |
| 5788 | fn apply_json_delta( |
| 5789 | raw: &mut serde_json::Value, |
| 5790 | before: &serde_json::Value, |
| 5791 | after: &serde_json::Value, |
| 5792 | ) { |
| 5793 | if before == after { |
| 5794 | return; |
| 5795 | } |
| 5796 | if let (Some(raw), Some(before), Some(after)) = |
| 5797 | (raw.as_object_mut(), before.as_object(), after.as_object()) |
| 5798 | { |
| 5799 | for key in before.keys().chain(after.keys()).collect::<BTreeSet<_>>() { |
| 5800 | match (before.get(key), after.get(key)) { |
| 5801 | (Some(old), Some(new)) if old != new => { |
| 5802 | apply_json_delta(raw.entry(key.clone()).or_insert(old.clone()), old, new); |
| 5803 | } |
| 5804 | (None, Some(new)) => { |
| 5805 | raw.insert(key.clone(), new.clone()); |
| 5806 | } |
| 5807 | (Some(_), None) => { |
| 5808 | raw.remove(key); |
| 5809 | } |
| 5810 | _ => {} |
| 5811 | } |
| 5812 | } |
| 5813 | } else { |
| 5814 | *raw = after.clone(); |
| 5815 | } |
| 5816 | } |
| 5817 | |
| 5818 | /// Every managed MCP writer rereads under the same OS-process lock. This is |
| 5819 | /// a delta operation, not a save of a previously loaded typed snapshot. |
| 5820 | pub fn mutate_config<T>( |
| 5821 | path: &Path, |
| 5822 | expected_revision: Option<&str>, |
| 5823 | mutate: impl FnOnce(&mut McpConfig) -> Result<T>, |
| 5824 | ) -> Result<(T, String)> { |
| 5825 | validate_mcp_config_path(path)?; |
| 5826 | codewhale_config::with_config_write_lock(path, |path| { |
| 5827 | let original = read_mcp_config_file(path)?; |
| 5828 | let revision = config_revision(original.as_deref()); |
| 5829 | if expected_revision.is_some_and(|expected| expected != revision) { |
| 5830 | return Err(McpRevisionConflict.into()); |
| 5831 | } |
| 5832 | let mut raw: serde_json::Value = match original.as_deref() { |
| 5833 | Some(raw) => serde_json::from_str(raw).map_err(|_| { |
| 5834 | anyhow::anyhow!("Failed to parse MCP config; file contents were omitted") |
| 5835 | })?, |
| 5836 | None => serde_json::json!({}), |
| 5837 | }; |
| 5838 | anyhow::ensure!(raw.is_object(), "MCP config must be an object"); |
| 5839 | let mut config: McpConfig = serde_json::from_value(raw.clone()) |
| 5840 | .map_err(|_| anyhow::anyhow!("Invalid MCP config; file contents were omitted"))?; |
| 5841 | let before = serde_json::to_value(&config)?; |
| 5842 | let result = mutate(&mut config)?; |
| 5843 | let after = serde_json::to_value(&config)?; |
| 5844 | if before == after { |
| 5845 | return Ok((result, revision)); |
| 5846 | } |
| 5847 | // Preserve legacy spelling while applying the canonical typed delta. |
| 5848 | let legacy = raw.get("mcpServers").is_some(); |
| 5849 | if legacy { |
| 5850 | let object = raw |
| 5851 | .as_object_mut() |
| 5852 | .context("MCP config must be an object")?; |
| 5853 | let servers = object.remove("mcpServers").expect("checked above"); |
| 5854 | object.insert("servers".into(), servers); |
| 5855 | } |
| 5856 | apply_json_delta(&mut raw, &before, &after); |
| 5857 | if legacy { |
| 5858 | let object = raw |
| 5859 | .as_object_mut() |
| 5860 | .context("MCP config must be an object")?; |
| 5861 | if let Some(servers) = object.remove("servers") { |
| 5862 | object.insert("mcpServers".into(), servers); |
| 5863 | } |
| 5864 | } |
| 5865 | let rendered = serde_json::to_string_pretty(&raw)?; |
| 5866 | if rendered.len() as u64 > MAX_MCP_CONFIG_BYTES { |
| 5867 | anyhow::bail!("MCP config exceeds the 1 MiB limit"); |
| 5868 | } |
| 5869 | write_atomic(path, rendered.as_bytes())?; |
| 5870 | Ok((result, config_revision(Some(&rendered)))) |
| 5871 | }) |
| 5872 | } |
| 5873 | |
| 5874 | fn mcp_template_json() -> Result<String> { |
| 5875 | let mut cfg = McpConfig::default(); |
| 5876 | cfg.servers.insert( |
| 5877 | "example".to_string(), |
| 5878 | McpServerConfig { |
| 5879 | command: Some("node".to_string()), |
| 5880 | args: vec!["./path/to/your-mcp-server.js".to_string()], |
| 5881 | env: HashMap::new(), |
| 5882 | cwd: None, |
| 5883 | url: None, |
| 5884 | transport: None, |
| 5885 | connect_timeout: None, |
| 5886 | execute_timeout: None, |
| 5887 | read_timeout: None, |
| 5888 | disabled: true, |
| 5889 | enabled: true, |
| 5890 | required: false, |
| 5891 | enabled_tools: Vec::new(), |
| 5892 | disabled_tools: Vec::new(), |
| 5893 | headers: HashMap::new(), |
| 5894 | env_headers: HashMap::new(), |
| 5895 | bearer_token_env_var: None, |
| 5896 | scopes: Vec::new(), |
| 5897 | oauth: None, |
| 5898 | oauth_resource: None, |
| 5899 | reviewed_plugin: None, |
| 5900 | runtime_added: false, |
| 5901 | allow_private_network: false, |
| 5902 | }, |
| 5903 | ); |
| 5904 | serde_json::to_string_pretty(&cfg).context("Failed to render MCP template JSON") |
| 5905 | } |
| 5906 | |
| 5907 | pub fn init_config(path: &Path, force: bool) -> Result<McpWriteStatus> { |
| 5908 | validate_mcp_config_path(path)?; |
| 5909 | codewhale_config::with_config_write_lock(path, |path| { |
| 5910 | let original = read_mcp_config_file(path)?; |
| 5911 | if let Some(raw) = original.as_deref() { |
| 5912 | let _: McpConfig = serde_json::from_str(raw) |
| 5913 | .map_err(|_| anyhow::anyhow!("Invalid MCP config; file contents were omitted"))?; |
| 5914 | if !force { |
| 5915 | return Ok(McpWriteStatus::SkippedExists); |
| 5916 | } |
| 5917 | } |
| 5918 | let template = mcp_template_json()?; |
| 5919 | write_atomic(path, template.as_bytes())?; |
| 5920 | Ok(if original.is_some() { |
| 5921 | McpWriteStatus::Overwritten |
| 5922 | } else { |
| 5923 | McpWriteStatus::Created |
| 5924 | }) |
| 5925 | }) |
| 5926 | } |
| 5927 | |
| 5928 | pub fn add_server_config( |
| 5929 | path: &Path, |
| 5930 | name: String, |
| 5931 | command: Option<String>, |
| 5932 | url: Option<String>, |
| 5933 | args: Vec<String>, |
| 5934 | transport: Option<String>, |
| 5935 | ) -> Result<()> { |
| 5936 | if command.is_none() && url.is_none() { |
| 5937 | anyhow::bail!("Provide either a command or URL for MCP server '{name}'."); |
| 5938 | } |
| 5939 | validate_mcp_transport(transport.as_deref())?; |
| 5940 | mutate_config(path, None, |cfg| { |
| 5941 | cfg.servers.insert( |
| 5942 | name, |
| 5943 | McpServerConfig { |
| 5944 | command, |
| 5945 | args, |
| 5946 | env: HashMap::new(), |
| 5947 | cwd: None, |
| 5948 | url, |
| 5949 | transport, |
| 5950 | connect_timeout: None, |
| 5951 | execute_timeout: None, |
| 5952 | read_timeout: None, |
| 5953 | disabled: false, |
| 5954 | enabled: true, |
| 5955 | required: false, |
| 5956 | enabled_tools: Vec::new(), |
| 5957 | disabled_tools: Vec::new(), |
| 5958 | headers: HashMap::new(), |
| 5959 | env_headers: HashMap::new(), |
| 5960 | bearer_token_env_var: None, |
| 5961 | scopes: Vec::new(), |
| 5962 | oauth: None, |
| 5963 | oauth_resource: None, |
| 5964 | reviewed_plugin: None, |
| 5965 | runtime_added: false, |
| 5966 | allow_private_network: false, |
| 5967 | }, |
| 5968 | ); |
| 5969 | Ok(()) |
| 5970 | }) |
| 5971 | .map(|_| ()) |
| 5972 | } |
| 5973 | |
| 5974 | pub fn remove_server_config(path: &Path, name: &str) -> Result<()> { |
| 5975 | mutate_config(path, None, |cfg| { |
| 5976 | if cfg.servers.remove(name).is_none() { |
| 5977 | anyhow::bail!("MCP server '{name}' not found"); |
| 5978 | } |
| 5979 | Ok(()) |
| 5980 | }) |
| 5981 | .map(|_| ()) |
| 5982 | } |
| 5983 | |
| 5984 | pub fn set_server_enabled(path: &Path, name: &str, enabled: bool) -> Result<()> { |
| 5985 | mutate_config(path, None, |cfg| { |
| 5986 | let server = cfg |
| 5987 | .servers |
| 5988 | .get_mut(name) |
| 5989 | .ok_or_else(|| anyhow::anyhow!("MCP server '{name}' not found"))?; |
| 5990 | server.enabled = enabled; |
| 5991 | server.disabled = !enabled; |
| 5992 | Ok(()) |
| 5993 | }) |
| 5994 | .map(|_| ()) |
| 5995 | } |
| 5996 | |
| 5997 | #[cfg(test)] |
| 5998 | pub fn manager_snapshot_from_config( |
| 5999 | path: &Path, |
| 6000 | reload_required: bool, |
| 6001 | ) -> Result<McpManagerSnapshot> { |
| 6002 | let cfg = load_config(path)?; |
| 6003 | Ok(snapshot_from_config( |
| 6004 | path, |
| 6005 | path.exists(), |
| 6006 | reload_required, |
| 6007 | &cfg, |
| 6008 | None, |
| 6009 | )) |
| 6010 | } |
| 6011 | |
| 6012 | #[cfg(test)] |
| 6013 | pub fn manager_snapshot_from_config_with_workspace( |
| 6014 | path: &Path, |
| 6015 | workspace: &Path, |
| 6016 | reload_required: bool, |
| 6017 | ) -> Result<McpManagerSnapshot> { |
| 6018 | let plugins = crate::plugins::PluginRegistry::empty(workspace); |
| 6019 | manager_snapshot_from_config_with_workspace_and_plugins( |
| 6020 | path, |
| 6021 | workspace, |
| 6022 | reload_required, |
| 6023 | &plugins, |
| 6024 | ) |
| 6025 | } |
| 6026 | |
| 6027 | pub fn manager_snapshot_from_config_with_workspace_and_plugins( |
| 6028 | path: &Path, |
| 6029 | workspace: &Path, |
| 6030 | reload_required: bool, |
| 6031 | plugins: &crate::plugins::PluginRegistry, |
| 6032 | ) -> Result<McpManagerSnapshot> { |
| 6033 | let cfg = load_config_with_workspace_and_plugins(path, workspace, plugins)?; |
| 6034 | Ok(snapshot_from_config( |
| 6035 | path, |
| 6036 | path.exists(), |
| 6037 | reload_required, |
| 6038 | &cfg, |
| 6039 | None, |
| 6040 | )) |
| 6041 | } |
| 6042 | |
| 6043 | #[cfg(test)] |
| 6044 | pub async fn discover_manager_snapshot( |
| 6045 | path: &Path, |
| 6046 | network_policy: Option<NetworkPolicyDecider>, |
| 6047 | reload_required: bool, |
| 6048 | ) -> Result<McpManagerSnapshot> { |
| 6049 | let cfg = load_config(path)?; |
| 6050 | let mut pool = McpPool::new(cfg.clone()); |
| 6051 | if let Some(policy) = network_policy { |
| 6052 | pool = pool.with_network_policy(policy); |
| 6053 | } |
| 6054 | let errors = pool |
| 6055 | .connect_all() |
| 6056 | .await |
| 6057 | .into_iter() |
| 6058 | .map(|(name, err)| (name, format_mcp_error_for_display(&err))) |
| 6059 | .collect::<HashMap<_, _>>(); |
| 6060 | Ok(snapshot_from_config( |
| 6061 | path, |
| 6062 | path.exists(), |
| 6063 | reload_required, |
| 6064 | &cfg, |
| 6065 | Some((&pool, &errors)), |
| 6066 | )) |
| 6067 | } |
| 6068 | |
| 6069 | pub async fn discover_manager_snapshot_with_workspace_and_plugins( |
| 6070 | path: &Path, |
| 6071 | workspace: &Path, |
| 6072 | network_policy: Option<NetworkPolicyDecider>, |
| 6073 | reload_required: bool, |
| 6074 | plugins: Arc<crate::plugins::PluginRegistry>, |
| 6075 | ) -> Result<McpManagerSnapshot> { |
| 6076 | let cfg = load_config_with_workspace_and_plugins(path, workspace, plugins.as_ref())?; |
| 6077 | let mut pool = McpPool::new(cfg.clone()); |
| 6078 | pool.workspace = Some(checked_workspace_path(workspace)?); |
| 6079 | pool.plugin_registry = Some(plugins); |
| 6080 | if let Some(policy) = network_policy { |
| 6081 | pool = pool.with_network_policy(policy); |
| 6082 | } |
| 6083 | let errors = pool |
| 6084 | .connect_all() |
| 6085 | .await |
| 6086 | .into_iter() |
| 6087 | .map(|(name, err)| (name, format_mcp_error_for_display(&err))) |
| 6088 | .collect::<HashMap<_, _>>(); |
| 6089 | Ok(snapshot_from_config( |
| 6090 | path, |
| 6091 | path.exists(), |
| 6092 | reload_required, |
| 6093 | &cfg, |
| 6094 | Some((&pool, &errors)), |
| 6095 | )) |
| 6096 | } |
| 6097 | |
| 6098 | pub(crate) fn format_mcp_error_for_display(error: &anyhow::Error) -> String { |
| 6099 | codewhale_config::persistence::redact_secrets(&format!("{error:#}")) |
| 6100 | } |
| 6101 | |
| 6102 | impl McpPool { |
| 6103 | /// Snapshot the live pool rather than starting a second discovery pool. |
| 6104 | /// This keeps the manager, hotbar, and next model turn aligned on one |
| 6105 | /// exact config/catalog generation. |
| 6106 | pub(crate) fn manager_snapshot( |
| 6107 | &self, |
| 6108 | path: &Path, |
| 6109 | reload_required: bool, |
| 6110 | errors: &HashMap<String, String>, |
| 6111 | ) -> McpManagerSnapshot { |
| 6112 | snapshot_from_config( |
| 6113 | path, |
| 6114 | path.exists(), |
| 6115 | reload_required, |
| 6116 | &self.config, |
| 6117 | Some((self, errors)), |
| 6118 | ) |
| 6119 | } |
| 6120 | } |
| 6121 | |
| 6122 | fn snapshot_from_config( |
| 6123 | path: &Path, |
| 6124 | config_exists: bool, |
| 6125 | reload_required: bool, |
| 6126 | cfg: &McpConfig, |
| 6127 | discovery: Option<(&McpPool, &HashMap<String, String>)>, |
| 6128 | ) -> McpManagerSnapshot { |
| 6129 | let mut servers = cfg |
| 6130 | .servers |
| 6131 | .iter() |
| 6132 | .filter(|(name, _)| discovery.is_none_or(|(pool, _)| pool.server_allowed(name))) |
| 6133 | .map(|(name, server)| { |
| 6134 | let transport = if server.url.is_some() { |
| 6135 | if is_legacy_sse_transport(server) { |
| 6136 | "sse" |
| 6137 | } else { |
| 6138 | "http/sse" |
| 6139 | } |
| 6140 | } else { |
| 6141 | "stdio" |
| 6142 | }; |
| 6143 | let command_or_url = server.url.clone().unwrap_or_else(|| { |
| 6144 | let mut command = server |
| 6145 | .command |
| 6146 | .clone() |
| 6147 | .unwrap_or_else(|| "(missing)".to_string()); |
| 6148 | if !server.args.is_empty() { |
| 6149 | command.push(' '); |
| 6150 | command.push_str(&server.args.join(" ")); |
| 6151 | } |
| 6152 | command |
| 6153 | }); |
| 6154 | let mut snapshot = McpServerSnapshot { |
| 6155 | name: name.clone(), |
| 6156 | enabled: server.is_enabled(), |
| 6157 | required: server.required, |
| 6158 | transport: transport.to_string(), |
| 6159 | command_or_url, |
| 6160 | connect_timeout: server.effective_connect_timeout(&cfg.timeouts), |
| 6161 | execute_timeout: server.effective_execute_timeout(&cfg.timeouts), |
| 6162 | read_timeout: server.effective_read_timeout(&cfg.timeouts), |
| 6163 | connected: false, |
| 6164 | error: if server.is_enabled() { |
| 6165 | None |
| 6166 | } else { |
| 6167 | Some("disabled".to_string()) |
| 6168 | }, |
| 6169 | auth_required: false, |
| 6170 | capability_metadata: McpServerCapabilityMetadata::NotObserved, |
| 6171 | tools: Vec::new(), |
| 6172 | resources: Vec::new(), |
| 6173 | prompts: Vec::new(), |
| 6174 | }; |
| 6175 | |
| 6176 | if let Some((pool, errors)) = discovery { |
| 6177 | if let Some(error) = errors.get(name) { |
| 6178 | snapshot.error = Some(error.clone()); |
| 6179 | } |
| 6180 | // The pool's needs-auth set is the authority; the error text |
| 6181 | // fallback keeps a boot-time error map (held by the engine |
| 6182 | // after the pool's live state was rebuilt) on the same |
| 6183 | // classification instead of downgrading to a plain failure. |
| 6184 | snapshot.auth_required = server.is_enabled() |
| 6185 | && (pool.server_needs_auth(name) |
| 6186 | || snapshot |
| 6187 | .error |
| 6188 | .as_deref() |
| 6189 | .is_some_and(oauth::error_text_looks_auth_required)); |
| 6190 | if let Some(conn) = pool.connections.get(name) { |
| 6191 | snapshot.connected = conn.is_ready(); |
| 6192 | snapshot.capability_metadata = conn.server_capabilities.map_or( |
| 6193 | McpServerCapabilityMetadata::LegacyFallback, |
| 6194 | McpServerCapabilityMetadata::Advertised, |
| 6195 | ); |
| 6196 | if snapshot.connected { |
| 6197 | // A count of connected servers and nothing else. The |
| 6198 | // name, the command or URL, and the error string are |
| 6199 | // user-chosen and routinely name internal infra. |
| 6200 | codewhale_telemetry::session_counters() |
| 6201 | .bump(codewhale_telemetry::Counter::McpServerConnected); |
| 6202 | } |
| 6203 | snapshot.tools = conn |
| 6204 | .tools() |
| 6205 | .iter() |
| 6206 | .filter(|tool| { |
| 6207 | conn.config().is_tool_enabled(&tool.name) |
| 6208 | && pool |
| 6209 | .tool_allowed(&McpPool::mcp_model_tool_name(name, &tool.name)) |
| 6210 | }) |
| 6211 | .map(|tool| McpDiscoveredItem { |
| 6212 | name: tool.name.clone(), |
| 6213 | model_name: format!("mcp_{}_{}", name, tool.name), |
| 6214 | description: tool.description.clone(), |
| 6215 | }) |
| 6216 | .collect(); |
| 6217 | snapshot.resources = |
| 6218 | conn.resources() |
| 6219 | .iter() |
| 6220 | .map(|resource| McpDiscoveredItem { |
| 6221 | name: resource.name.clone(), |
| 6222 | model_name: format!( |
| 6223 | "mcp_{}_{}", |
| 6224 | name, |
| 6225 | resource.name.replace(' ', "_").to_lowercase() |
| 6226 | ), |
| 6227 | description: resource.description.clone(), |
| 6228 | }) |
| 6229 | .chain(conn.resource_templates().iter().map(|template| { |
| 6230 | McpDiscoveredItem { |
| 6231 | name: template.name.clone(), |
| 6232 | model_name: format!( |
| 6233 | "mcp_{}_{}", |
| 6234 | name, |
| 6235 | template.name.replace(' ', "_").to_lowercase() |
| 6236 | ), |
| 6237 | description: template.description.clone(), |
| 6238 | } |
| 6239 | })) |
| 6240 | .collect(); |
| 6241 | snapshot.prompts = conn |
| 6242 | .prompts() |
| 6243 | .iter() |
| 6244 | .map(|prompt| McpDiscoveredItem { |
| 6245 | name: prompt.name.clone(), |
| 6246 | model_name: format!("mcp_{}_{}", name, prompt.name), |
| 6247 | description: prompt.description.clone(), |
| 6248 | }) |
| 6249 | .collect(); |
| 6250 | } |
| 6251 | } |
| 6252 | |
| 6253 | snapshot |
| 6254 | }) |
| 6255 | .collect::<Vec<_>>(); |
| 6256 | servers.sort_by(|a, b| a.name.cmp(&b.name)); |
| 6257 | McpManagerSnapshot { |
| 6258 | config_path: path.to_path_buf(), |
| 6259 | config_exists, |
| 6260 | reload_required, |
| 6261 | servers, |
| 6262 | } |
| 6263 | } |
| 6264 | |
| 6265 | // === Unit Tests === |
| 6266 | |
| 6267 | #[cfg(test)] |
| 6268 | mod tests; |
| 6269 |