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